1️⃣ Policy Gradient Theorem

Learning Objectives
  • Understand the Policy Gradient Theorem
  • Understand the VPG algorithm: how to perform on-policy policy gradient

Policy Gradient Theorem

Instead of learning action-values and deriving a policy (as in Q-learning or DQN), policy gradient methods learn the policy directly. - Policy is parameterized: $\pi_\theta(a|s)$ with parameters $\theta$ (often a neural network).
- Objective: Choose $\theta$ to maximize expected return $J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}[G(\tau)]$ (joy), where $\tau$ is a trajectory and $G(\tau)$ its return.

We would desire to update the policy directly via gradient ascent against $J(\theta)$:

$$ \theta \leftarrow \theta + \alpha \nabla_\theta J(\theta) $$

The problem is that the return is a sum of rewards from the trajectory, and the trajectory itself is a result of sampling from the policy, over and over, as well as being dependent on the environmental distribution, which we do not have access to. There is no clear way to directly compute the gradient of the return with respect to the policy parameters. The solution here is the policy gradient theorem, which states that we can instead use the return weighted by the gradient of the log-probability as an unbiased estimator of the gradient of the return.

$$ \nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \sum_t G_t \nabla_\theta \log \pi_\theta(a_t|s_t) \right] $$

Derivation

The probability of sampling a trajectory $\tau = (s_0, a_0, s_1, a_1, \dots, s_T)$ is given by

$$ \Pr(\tau|\theta) = \prod_{t=0}^{T-1} \pi_\theta(a_t|s_t)\,\mu(s_{t+1}|s_t, a_t) $$
where $\mu$ is the environment transition probability.

$$ \begin{align*} \nabla_\theta J(\theta) &= \nabla_\theta \mathbb{E}_{\tau \sim \pi_\theta}[G(\tau)] \\ &= \nabla_\theta \sum_\tau \Pr(\tau|\theta) \, G(\tau) \\ &= \sum_\tau \nabla_\theta \Pr(\tau|\theta) \, G(\tau) \\ &= \sum_\tau \Pr(\tau|\theta)\,\nabla_\theta \log \Pr(\tau|\theta) \, G(\tau) \\ &= \mathbb{E}_{\tau \sim \pi_\theta}\left[ \nabla_\theta \log \Pr(\tau|\theta) \, G(\tau) \right] \end{align*} $$
where we made use of the log-derivative trick: $\nabla_\theta p(x) = p(x) \nabla_\theta \log p(x)$.

The dynamics $\mu$ do not depend on $\theta$, so:

$$ \begin{align*} \log \Pr(\tau|\theta) &= \log \left( \prod_{t=0}^{T-1} \pi_\theta(a_t|s_t)\,\mu(s_{t+1}|s_t, a_t) \right) \\ &= \sum_{t=0}^{T-1} \log \pi_\theta(a_t|s_t) + \sum_{t=0}^{T-1} \log \mu(s_{t+1}|s_t, a_t) \\ &= \sum_{t=0}^{T-1} \log \pi_\theta(a_t|s_t) + \text{const.} \end{align*} $$
where the const. term is independent of $\theta$, so when we take the gradient, it vanishes.

Thus:

$$ \nabla_\theta \log \Pr(\tau|\theta) = \sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t|s_t) $$

Plugging back into the gradient:

$$ \nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t|s_t)\, G(\tau) \right] $$

This is the Vanilla Policy Gradient estimator, also called REINFORCE. Each $\log \pi_\theta(a_t|s_t)$ is multiplied by the full return $G(\tau)$. However, the action $a_t$ cannot influence rewards before time $t$, only those afterwards. This means that all the rewards before timestep $t$ merely add noise, as no changes to the policy can affect them. To reduce variance, replace $G(\tau)$ with the return $G_t$ at timestep $t$, also called the reward-to-go:

$$ G_t = \sum_{i=t}^{T} \gamma^{i-t} r_{i} $$

Thus, the lower-variance unbiased estimator is:

$$ \nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t|s_t)\, G_t \right] $$

There are many other variants of the policy gradient estimator, as described in Schulman, 2018.

Implementation

We make use of exactly the same vectorised GPU CartPole environment as we did for DQN (see chapter2_rl/exercises/gpu_env.py): it's defined entirely in terms of tensor operations, so

  • we don't need to constantly convert between numpy and torch tensors
  • we can run large numbers of environments in parallel (thousands of environments, for millions of environment steps per second)
  • we avoid copying data back and forth between the CPU and GPU, which can be a significant bottleneck

For DQN we only ran a handful of environments at once, because DQN's sample efficiency comes from the replay buffer rather than from how much fresh data we collect per step. (Vanilla) Policy gradient methods can't reuse old data in the same way, so here we'll lean on the vectorisation much harder and collect a rollout from thousands of environments at once every update.

The Policy Network

The Policy Network takes in an observation $s$ and outputs a vector $\pi(a_1 | s), \ldots \pi(a_n | s)$ representing the probability that the network takes action $a_i$ given observation $s$. We will then adjust the parameters of this network to directly maximize the expected return.

For learning from pixels, a simple convolutional network and some fully connected layers does quite well. Where we have already processed features here, it's even easier: an MLP of this size should be plenty large for any environment today.

Implement the Policy Network using a standard MLP, constructed of alternating Linear and ReLU layers. The size of the input will match the dimensionality of the observation space, and the size of the output will match the number of actions to choose from (associating a probability to each action.) The dimensions of the hidden_sizes are provided.

Here is a diagram of what our particular Policy Network will look like for CartPole (you can open it in a new tab if it's hard to see clearly):


Question - why do we not include a ReLU at the end?

If you end with a ReLU, then the logits are capped below at 0, which means the network can only push down the undesirability of an action so far. We want the output to be unbounded so the network is free to disprefer an action as much as it wants.

Exercise - implement PolicyNetwork

Difficulty: 🔴🔴⚪⚪⚪
Importance: 🔵🔵🔵⚪⚪
You should spend up to 10-15 minutes on this exercise.

Note - in this implementation we can assume that obs_shape is a tuple of length 1 (in the case of CartPole this will be (4,)), so you can treat it as just an integer value above, e.g. your first linear layer should be from obs_shape[0] to 120.

class PolicyNetwork(nn.Module):
    """
    For consistency with your tests, please wrap your modules in a `nn.Sequential` called `layers`.
    """

    layers: nn.Sequential

    def __init__(self, obs_shape: tuple[int], num_actions: int, hidden_sizes: list[int] = [120, 84]):
        super().__init__()
        assert len(obs_shape) == 1, "Expecting a single vector of observations"
        raise NotImplementedError()

    def forward(self, x: Tensor) -> Tensor:
        return self.layers(x)


net = PolicyNetwork(obs_shape=(4,), num_actions=2)
n_params = sum((p.nelement() for p in net.parameters()))
assert isinstance(getattr(net, "layers", None), nn.Sequential)
print(net)
print(f"Total number of parameters: {n_params}")
print("You should manually verify network is Linear-ReLU-Linear-ReLU-Linear")
assert not isinstance(net.layers[-1], nn.ReLU)
assert n_params == 10934
Solution 
class PolicyNetwork(nn.Module):
    """
    For consistency with your tests, please wrap your modules in a `nn.Sequential` called `layers`.
    """

    layers: nn.Sequential

    def __init__(self, obs_shape: tuple[int], num_actions: int, hidden_sizes: list[int] = [120, 84]):
        super().__init__()
        assert len(obs_shape) == 1, "Expecting a single vector of observations"
        in_features_list = [obs_shape[0]] + hidden_sizes
        out_features_list = hidden_sizes + [num_actions]
        layers = []
        for i, (in_features, out_features) in enumerate(zip(in_features_list, out_features_list)):
            layers.append(nn.Linear(in_features, out_features))
            if i < len(in_features_list) - 1:
                layers.append(nn.ReLU())
        self.layers = nn.Sequential(*layers)

    def forward(self, x: Tensor) -> Tensor:
        return self.layers(x)