Policy Gradients and Misgeneralization

Prerequisites

• D.1.2 - Reinforcement Learning
• Be familar with the PyTorch prerequisites.
  • In particular, if you're unfamilar with more advanced indexing like use of torch.gather, feel free to solve the exercise with a sequential python for-loop, and then use the solution code when you want the training loop to run quickly.

What you’ll learn

• Understand the derivation of (Vanilla) Policy Gradient (VPG).
• Implement Policy Gradient in PyTorch, and use to train an agent to play a toy RL environment.
• Understand reward functions, what behaviour they incentivise and whow they can be used as probes for behaviour.
• Understand what potential shaping is, and how it provides hints without allowing reward hacking.
• The difference between Goal Misgeneralisation and Goal Misspecification.

Roadmap for today

• 10:00–10:30
  • Lecture on intro to Policy Gradient
• 10:30–13:00
  • Work through VPG material
• 13:00-14:00
  • (Lunch + Break Time)
• 14:00–14:30
  • Lecture on intro to Goal Misgeneralization
• 14:30–17:30
  • Work through Goal Misgen. material
• 17:30–18:00
  • Buffer + Feedback

Vanilla Policy Gradient

Work through Section [2.2.2] of the ARENA program on Vanilla Policy Gradient. Follow the provided link, and then through to the exercises on colab.

• The network is small enough, and the training loop fast enough, that you can run this material locally on your laptop without problems if you prefer.

Goal Misgeneralization

Work through the provided colab notebook on Goal Misgeneralisation.

• The network is again small, but some of the later runs might be slow without a GPU, so we would recommend against running this locally unless you have a GPU in your laptop.

Potential shaping

The notebook shapes the reward with a potential function. These exercises show this cannot change the optimal policy.

Definition (Potential shaping). Take an MDP with states $\mathcal{S}$, reward $R(s,a,s^{\prime })$ and discount $\gamma \in [0,1)$. A potential is a bounded $\Phi :\mathcal{S}\to \mathbb{R}$. The shaped reward is

$$R^{\Phi }(s,a,s^{\prime })=R(s,a,s^{\prime })+\gamma \Phi (s^{\prime })-\Phi (s).$$

For a trajectory $s_0,a_0,r_1,s_1,a_1,r_2,\dots$, taking action $a_t$ in state $s_t$ returns the next state $s_{t+1}$ and reward $r_{t+1}=R(s_t,a_t,s_{t+1})$. Write $r_{t+1}^{\Phi }=R^{\Phi }(s_t,a_t,s_{t+1})$, $G={\sum }_{t=0}^{\infty }{\gamma }^t{r}_{t+1}$ and $G^{\Phi }={\sum }_{t=0}^{\infty }{\gamma }^t{r}_{t+1}^{\Phi }$.

Exercise 1. Show that

$$\sum _{t=0}^{T-1}{\gamma }^t(\gamma \Phi (s_{t+1})-\Phi (s_t))={\gamma }^T\Phi (s_T)-\Phi (s_0).$$

Solution

With $x_t={\gamma }^t\Phi (s_t)$,

$${\gamma }^t(\gamma \Phi (s_{t+1})-\Phi (s_t))={\gamma }^{t+1}\Phi (s_{t+1})-{\gamma }^t\Phi (s_t)=x_{t+1}-x_t,$$

so

$$\sum _{t=0}^{T-1}(x_{t+1}-x_t)=x_T-x_0={\gamma }^T\Phi (s_T)-\Phi (s_0).$$

Exercise 2. Show that $G^{\Phi }=G-\Phi (s_0)$. Where are boundedness of $\Phi$ and $\gamma <1$ used?

Solution

Truncate at $T$ and apply Exercise 1:

$$\sum _{t=0}^{T-1}{\gamma }^t{r}_{t+1}^{\Phi }=\sum _{t=0}^{T-1}{\gamma }^t{r}_{t+1}+\sum _{t=0}^{T-1}{\gamma }^t(\gamma \Phi (s_{t+1})-\Phi (s_t))=\sum _{t=0}^{T-1}{\gamma }^t{r}_{t+1}+{\gamma }^T\Phi (s_T)-\Phi (s_0).$$

If $|\Phi |\le M$ then $|{\gamma }^T\Phi (s_T)|\le {\gamma }^{T}M\to 0$ as $T\to \infty$, since $\gamma <1$. Letting $T\to \infty$ gives $G^{\Phi }=G-\Phi (s_0)$. (An unbounded $\Phi$ could grow like ${\gamma }^{-T}$, and the boundary term would not vanish.)

Exercise 3. Show that $Q_{R^{\Phi }}^{\pi }(s,a)=Q_R^{\pi }(s,a)-\Phi (s)$ for every policy $\pi$.

Solution

Using Exercise 2, and that $s_0=s$ under the conditioning:

$$Q_{R^{\Phi }}^{\pi }(s,a)={\mathbb{E}}_{\pi }[G^{\Phi }\mid s_0=s,a_0=a]={\mathbb{E}}_{\pi }[G-\Phi (s_0)\mid s_0=s,a_0=a]=Q_R^{\pi }(s,a)-\Phi (s).$$

Exercise 4. Assume $\mathcal{S}$ and the action set $\mathcal{A}$ are finite, so an optimal policy exists (D.1.2, Exercise 3.7). Show that $R$ and $R^{\Phi }$ have the same optimal policies, and that $Q_{R^{\Phi }}^{\ast }(s,a)=Q_R^{\ast }(s,a)-\Phi (s)$.

Solution

Averaging Exercise 3 over $a\sim \pi (\cdot \mid s)$ gives $V_{R^{\Phi }}^{\pi }(s)=V_R^{\pi }(s)-\Phi (s)$. The shift $\Phi (s)$ is the same for every policy, so for all policies $\pi ,{\pi }^{\prime }$ and states $s$

$$V_{R^{\Phi }}^{\pi }(s)\ge V_{R^{\Phi }}^{{\pi }^{\prime }}(s)⟺V_R^{\pi }(s)\ge V_R^{{\pi }^{\prime }}(s).$$

A policy is optimal if its value is at least that of every other policy in every state, so $\pi$ is optimal for $R$ iff it is optimal for $R^{\Phi }$.

Now take an optimal policy ${\pi }^{\ast }$ for $R$, which exists by D.1.2. It is also optimal for $R^{\Phi }$, so by Exercise 3

$$Q_{R^{\Phi }}^{\ast }(s,a)=Q_{R^{\Phi }}^{{\pi }^{\ast }}(s,a)=Q_R^{{\pi }^{\ast }}(s,a)-\Phi (s)=Q_R^{\ast }(s,a)-\Phi (s).$$

(For episodic tasks the same holds if $\Phi =0$ at terminal states, so the boundary term in Exercise 2 vanishes.)

Exercise 5. In the notebook, $\Phi =0.5$ while holding a shard and $0$ otherwise. What shaping bonus does picking up a shard and dropping it again earn? What if we instead paid $+0.5$ per PICKUP?

Solution

The robot goes $s\to s^{\prime }\to s^{\prime \prime }$ with $\Phi (s)=\Phi (s^{\prime \prime })=0$ and $\Phi (s^{\prime })=0.5$. The discounted shaping bonus is

$$(\gamma \Phi (s^{\prime })-\Phi (s))+\gamma (\gamma \Phi (s^{\prime \prime })-\Phi (s^{\prime }))=0.5\gamma -0.5\gamma =0.$$

Paying $+0.5$ per pickup instead, repeating the cycle forever earns $0.5+0.5{\gamma }^2+0.5{\gamma }^4+\cdots =\frac{0.5}{1-{\gamma }^2}>0$, so an agent can farm it rather than bin shards.