- class torchrl.objectives.KLPENPPOLoss(*args, **kwargs)¶
KL Penalty PPO loss.
- The KL penalty loss has the following formula:
loss = loss - beta * KL(old_policy, new_policy)
The “beta” parameter is adapted on-the-fly to match a target KL divergence between the new and old policy, thus favouring a certain level of distancing between the two while still preventing them to be too much apart.
actor (ProbabilisticTensorDictSequential) – policy operator.
critic (ValueOperator) – value operator.
- Keyword Arguments:
dtarg (scalar, optional) – target KL divergence. Defaults to
samples_mc_kl (int, optional) – number of samples used to compute the KL divergence if no analytical formula can be found. Defaults to
beta (scalar, optional) – initial KL divergence multiplier. Defaults to
decrement (scalar, optional) – how much beta should be decremented if KL < dtarg. Valid range: decrement <= 1.0 default:
increment (scalar, optional) – how much beta should be incremented if KL > dtarg. Valid range: increment >= 1.0 default:
entropy_bonus (bool, optional) – if
True, an entropy bonus will be added to the loss to favour exploratory policies. Defaults to
samples_mc_entropy (int, optional) – if the distribution retrieved from the policy operator does not have a closed form formula for the entropy, a Monte-Carlo estimate will be used.
samples_mc_entropywill control how many samples will be used to compute this estimate. Defaults to
entropy_coef (scalar, optional) – entropy multiplier when computing the total loss. Defaults to
critic_coef (scalar, optional) – critic loss multiplier when computing the total loss. Defaults to
loss_critic_type (str, optional) – loss function for the value discrepancy. Can be one of “l1”, “l2” or “smooth_l1”. Defaults to
normalize_advantage (bool, optional) – if
True, the advantage will be normalized before being used. Defaults to
separate_losses (bool, optional) – if
True, shared parameters between policy and critic will only be trained on the policy loss. Defaults to
False, ie. gradients are propagated to shared parameters for both policy and critic losses.
advantage_key (str, optional) – [Deprecated, use set_keys(advantage_key=advantage_key) instead] The input tensordict key where the advantage is expected to be written. Defaults to
value_target_key (str, optional) – [Deprecated, use set_keys(value_target_key=value_target_key) instead] The input tensordict key where the target state value is expected to be written. Defaults to
value_key (str, optional) – [Deprecated, use set_keys(value_key) instead] The input tensordict key where the state value is expected to be written. Defaults to
If the actor and the value function share parameters, one can avoid calling the common module multiple times by passing only the head of the value network to the PPO loss module:
>>> common = SomeModule(in_keys=["observation"], out_keys=["hidden"]) >>> actor_head = SomeActor(in_keys=["hidden"]) >>> value_head = SomeValue(in_keys=["hidden"]) >>> # first option, with 2 calls on the common module >>> model = ActorCriticOperator(common, actor_head, value_head) >>> loss_module = PPOLoss(model.get_policy_operator(), model.get_value_operator()) >>> # second option, with a single call to the common module >>> loss_module = PPOLoss(ProbabilisticTensorDictSequential(model, actor_head), value_head)
This will work regardless of whether separate_losses is activated or not.
- forward(tensordict: TensorDictBase) TensorDict ¶
It is designed to read an input TensorDict and return another tensordict with loss keys named “loss*”.
Splitting the loss in its component can then be used by the trainer to log the various loss values throughout training. Other scalars present in the output tensordict will be logged too.
tensordict – an input tensordict with the values required to compute the loss.
A new tensordict with no batch dimension containing various loss scalars which will be named “loss*”. It is essential that the losses are returned with this name as they will be read by the trainer before backpropagation.