Value clipping for PPO loss

[wbinventor]

Motivation

It would nice to add a new option to compute the clipped value loss as used in OpenAI's Baselines for PPO:

https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/ppo2/model.py#L66-L75

Currently, the PPOLoss().loss_critic method calls the torchrl.objectives.utils.distance_loss function which supports the "l1", "l2" and "smooth_l1" loss functions. Perhaps this new clipped value loss function could be implemented as a new loss type within distance_loss.

The clipping fraction is also commonly reported as a metric by OpenAI's Baselines and this could be useful to report from PPOLoss for the clipped value loss, as well as for ClipPPOLoss for the loss_objective loss term.

[vmoens]

That makes sense!

@albertbou92 Thoughts? Do you want to take care of this one?

[albertbou92]

Yes I also thought in the past about adding the PPO value clipping.
Not sure yet whether the logic should go in PPO or in distance_loss. My first thought is in PPO, since I have only seen it used in PPO, but I guess could potentially be used in other methods.
In any case, I can definitely start a PR to implement it.

[matteobettini]

Just a note is that in PPO the critic is used without gradients so clipping has no impact on grads. In other AC losses that flow gradients through the critic, clipping will cause to loose the gradient feeback when out of range.

[vmoens]

@matteobettini we could raise a warning if the value requires grad in that case?

In a way what's great about a good lib is that it allows you to do loads of stuff but also tells you when you're off the road

[albertbou92]

Here is the code generally used:

# Compute normal values loss
value_loss = (new_v - return).pow(2)

if self.use_clipped_value_loss:
    value_pred_clipped = old_v + (new_v - old_v).clamp(-self.clip_param, self.clip_param)
    value_loss_clipped = (value_pred_clipped - return).pow(2)
    
    # Chose the most pessimistic value prediction between clipped and non-clipped
    # so the updates are smaller
    value_loss = torch.max(value_loss, value_loss_clipped)

value_loss = 0.5 * value_loss.mean()

where,
old_v : value predicted during collection, without grads
new_v: values predicted during loss computation, with grads

So thinking in terms of PPO, the value_loss and value_pred_clipped will be the same in the first batch update.
In the next iterations, we simply mitigate excessively large updates to the value function.

We are clipping a tensor that has gradients, but this is similar to the PPO actor loss that clips the ratio = torch.exp(new_logp - old_logp).

This make sense for PPO that reuses the data multiple times. Not so sure it make sense for other on-policy methods.
Also, this requires having the original value_pred computed by a relatively recent policy, so not sure if it makes sense for off-policy methods.

In summary, at this point my opinion is that we can implement it just for on-policy methods. And the question is whether it should just be only PPO (for which the clipping make more sense) or for all of them. WDYT?

[wbinventor]

I think your rationale for only supporting on-policy methods makes sense 👍

[albertbou92]

This PR adds the clipping for PPO --> #2005

As mentioned, we could do something similar for other On-Policy Losses. However, the PPO value clipping generally reuses the clipping parameter of PPO. To add it to other losses we would need to define a new clip parameter specifically for the value clipping. Do we want to do that? Maybe for A2C and Reinforce I would only add it if users ask for it, as it is unclear if the clipping will be necessary/useful there.

[wbinventor]

It may be useful to define a separate parameter for this value clipping threshold. It can be useful to tune this different hyperparameter separate from that for the PPO loss.

[vmoens]

It may be useful to define a separate parameter for this value clipping threshold. It can be useful to tune this different hyperparameter separate from that for the PPO loss.

This is the plan! #2005 (comment)

[wbinventor]

Also, it is common and useful to log the "clip fraction" separately for both the PPO loss and the clipped value loss. Perhaps these fractions could be returned as "metrics" similar to how the "entropy" is returned as an item in the loss TensorDict?

[albertbou92]

Make sense, we could add them. You mean returning them after the clipping then right?

clip_fraction_policy = (new_log_prob / current_log_prob).clamp(1 - clip_epsilon, 1 + clip_epsilon).abs()
clip_fraction_value = (new_value_prediction / old_value_prediction).clamp(1 - clip_value, 1 + clip_value).abs()

[wbinventor]

That's right, I think that code implementation would be great to include 👍