3️⃣ Loss
Returns
To compute the REINFORCE loss, we need to compute the return for each step in the trajectory. This gets a little messy as trajectories may be of different lengths, so an episode may have terminated part way through the rollout. You'll need to walk backward through the trajectory, and compute the return for each step.
Note that gen_rollout stores done = terminated | truncated. This is the same done that the DQN section computed when it stepped its environments: it marks an episode boundary.
A real truncated tail does remain, but it comes from the rollout window rather than from either flag. gen_rollout runs for num_steps_per_rollout steps and your backward pass will start from $G = 0$, so an environment which is still alive when the window ends gets no credit at all for the reward it would have gone on to collect. The actor-critic methods in the PPO section fix this: with a critic you can bootstrap $V(s_T)$ at the end of the window instead of pretending the future is worth zero.
Exercise - implement compute_returns
Compute the returns for each trajectory. Easiest to write as a simple reverse for-loop for now, though if you wish later on you can try a vectorized solution.
def compute_returns(
rewards: Float[Tensor, " num_envs num_steps"], done: Bool[Tensor, " num_envs num_steps"], gamma: float = 0.9
):
"""
ARGS:
rewards: The rewards for each trajectory
done: A boolean tensor indicating if an episode finished on the current timestep
gamma: The discount factor
Returns:
The returns G_t for each trajectory.
For example:
- If Rewards = [0, 0, 1, 0, 1]
- And Done = [0, 0, 1, 0, 1]
- Then Returns = [g**2, g, 1, g, 1]
"""
num_envs, num_steps = rewards.shape
returns = t.zeros_like(rewards)
raise NotImplementedError()
tests.test_compute_returns(compute_returns)
Solution
def compute_returns(
rewards: Float[Tensor, " num_envs num_steps"], done: Bool[Tensor, " num_envs num_steps"], gamma: float = 0.9
):
"""
ARGS:
rewards: The rewards for each trajectory
done: A boolean tensor indicating if an episode finished on the current timestep
gamma: The discount factor
Returns:
The returns G_t for each trajectory.
For example:
- If Rewards = [0, 0, 1, 0, 1]
- And Done = [0, 0, 1, 0, 1]
- Then Returns = [g**2, g, 1, g, 1]
"""
num_envs, num_steps = rewards.shape
returns = t.zeros_like(rewards)
G = t.zeros_like(rewards[:, 0]) # (num_envs)
for i in reversed(range(num_steps)):
G = rewards[:, i] + gamma * G * (~done[:, i])
returns[:, i] = G
return returns
tests.test_compute_returns(compute_returns)
Exercise - implement compute_logprobs
Computes the log-probability under the current policy of the action that was taken on each timestep. To pick out the logprob of each taken action from the (num_envs, num_steps, num_actions) log-probs, use gather along the last dimension (as get_actions does), or the eindex library.
def compute_logprobs(tau: RolloutTensors, pi: PolicyNetwork) -> Float[Tensor, " num_envs num_steps"]:
"""
Computes the logprobs of the actions taken on each timestep under the current policy.
"""
raise NotImplementedError()
tests.test_compute_logprobs(compute_logprobs, PolicyNetwork)
Solution
def compute_logprobs(tau: RolloutTensors, pi: PolicyNetwork) -> Float[Tensor, " num_envs num_steps"]:
"""
Computes the logprobs of the actions taken on each timestep under the current policy.
"""
logits = pi(tau.obs)
log_probs = F.log_softmax(logits, dim=-1)
log_probs_taken = log_probs.gather(-1, tau.actions.unsqueeze(-1)).squeeze(-1)
# alternatively with eindex
# log_probs_taken = eindex(log_probs, tau.actions, "env time [env time] -> env time")
return log_probs_taken
tests.test_compute_logprobs(compute_logprobs, PolicyNetwork)
Building up to the loss function
We need to compute the probability ratio $\pi(a_t | s_t) / \pi_{old}(a_t | s_t)$ for each timestep taken in the rollout. This is used to compute the importance weights $\text{iw}_t$, which allows us to learn off-policy. If args.clip_coef is not None, we also clamp the importance weights between 1 - args.clip_coef and 1 + args.clip_coef.
Exercise - implement compute_importance_weights
Keep the result numerically stable by exponentiating the difference between the logprobs.
Gradients should NOT flow through the importance weights. Make sure to use .detach() to prevent this.
def compute_importance_weights(logprobs_taken, tau: RolloutTensors, clip_coef: Optional[float]) -> t.Tensor:
"""
Compute importance weights from log probabilities.
Keeps the result numerically stable by exponentiating the difference between logprobs.
Gradients should NOT flow through the importance weights (uses .detach()).
Optionally clips the weights to [1 - clip_coef, 1 + clip_coef].
"""
raise NotImplementedError()
tests.test_compute_importance_weights(compute_importance_weights)
Solution
def compute_importance_weights(logprobs_taken, tau: RolloutTensors, clip_coef: Optional[float]) -> t.Tensor:
"""
Compute importance weights from log probabilities.
Keeps the result numerically stable by exponentiating the difference between logprobs.
Gradients should NOT flow through the importance weights (uses .detach()).
Optionally clips the weights to [1 - clip_coef, 1 + clip_coef].
"""
iw = t.exp(logprobs_taken - tau.logprobs).detach() # Detach to prevent gradient flow
if clip_coef is not None:
iw = t.clamp(iw, 1 - clip_coef, 1 + clip_coef)
return iw
tests.test_compute_importance_weights(compute_importance_weights)
Exercise - implement normalize_returns
Normalize the returns by ensuring zero mean, unit variance across all trajectories and timesteps. Don't overthink this one, should be a one-liner.
def normalize_returns(returns: Float[Tensor, " num_envs num_steps"]) -> Float[Tensor, " num_envs num_steps"]:
"""
Normalizes the returns by ensuring zero mean, unit variance across all trajectories and timesteps.
"""
raise NotImplementedError()
tests.test_normalize_returns(normalize_returns)
Solution
def normalize_returns(returns: Float[Tensor, " num_envs num_steps"]) -> Float[Tensor, " num_envs num_steps"]:
"""
Normalizes the returns by ensuring zero mean, unit variance across all trajectories and timesteps.
"""
return (returns - returns.mean()) / (returns.std() + 1e-8)
tests.test_normalize_returns(normalize_returns)
Exercise - implement compute_reinforce_loss
This should be easy with everything else you've got. The loss on timestep $t$ is
returns.mean(dim=0, keepdim=True), and subtract it from returns. Remember to .detach() the result so no gradients flow through the advantage. (Separately, compute_loss may have already rescaled the returns to mean-zero unit-variance across all timesteps and trajectories, controlled by args.normalize_returns - that's variance reduction, not the baseline.) PPO uses a learned baseline called a critic, which we will see in the PPO section.
The total loss is the mean of the losses over all timesteps, over all trajectories.
def compute_reinforce_loss(
returns: Float[Tensor, " num_envs num_steps"],
logprobs_taken: Float[Tensor, " num_envs num_steps"],
iw: Float[Tensor, " num_envs num_steps"],
) -> Float[Tensor, ""]:
raise NotImplementedError()
tests.test_compute_reinforce_loss(compute_reinforce_loss)
Solution
def compute_reinforce_loss(
returns: Float[Tensor, " num_envs num_steps"],
logprobs_taken: Float[Tensor, " num_envs num_steps"],
iw: Float[Tensor, " num_envs num_steps"],
) -> Float[Tensor, ""]:
adv = returns - returns.mean(dim=0, keepdim=True) # baseline per timestep across envs
return (iw * logprobs_taken * adv.detach()).mean()
tests.test_compute_reinforce_loss(compute_reinforce_loss)