# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import annotations
import contextlib
import warnings
from collections.abc import Mapping
from dataclasses import dataclass
import torch
from tensordict import (
is_tensor_collection,
TensorDict,
TensorDictBase,
TensorDictParams,
)
from tensordict.nn import (
composite_lp_aggregate,
CompositeDistribution,
dispatch,
ProbabilisticTensorDictModule,
ProbabilisticTensorDictSequential,
set_composite_lp_aggregate,
TensorDictModule,
)
from tensordict.utils import NestedKey
from torch import distributions as d
from torchrl._utils import _standardize
from torchrl.modules.distributions.utils import (
_warn_mc_entropy,
_warn_mc_kl,
composite_entropy,
has_analytic_entropy,
has_analytic_kl,
is_dynamo_compiling,
sample_and_log_prob,
)
from torchrl.modules.value_norm import ValueNorm
from torchrl.objectives.common import LossModule
from torchrl.objectives.utils import (
_cache_values,
_clip_value_loss,
_GAMMA_LMBDA_DEPREC_ERROR,
_maybe_add_or_extend_key,
_maybe_get_or_select,
_sum_td_features,
_valid_value_target_rows,
_validate_clip_epsilon,
dispatch_value_estimator,
distance_loss,
ValueEstimators,
)
def _check_advantage_broadcast(
advantage: torch.Tensor, log_weight: torch.Tensor
) -> None:
# A flat [B] advantage against a [B, 1] log-weight silently broadcasts to
# a [B, B] outer product inside the surrogate objective. The log-weight is
# always shaped [*advantage.shape[:-1], 1], so any broadcast producing
# more elements than both inputs is that bug.
if not isinstance(advantage, torch.Tensor) or not isinstance(
log_weight, torch.Tensor
):
return
try:
broadcast = torch.broadcast_shapes(advantage.shape, log_weight.shape)
except RuntimeError:
# incompatible shapes fail with a shape error downstream
return
if broadcast.numel() > max(advantage.numel(), log_weight.numel()):
raise ValueError(
f"The advantage (shape={tuple(advantage.shape)}) and the log-weight "
f"(shape={tuple(log_weight.shape)}) would broadcast to "
f"{tuple(broadcast)}, turning the PPO objective into an outer "
"product over samples. This usually means the advantage is missing "
"its trailing value dimension: reshape it to batch shape + (1,) "
"(e.g. advantage.unsqueeze(-1))."
)
def _broadcast_advantage_to_log_weight(
advantage: torch.Tensor, log_weight: torch.Tensor
) -> torch.Tensor:
# Per-token objectives (e.g. token-level DAPO / GRPO on a chunked VLA head)
# carry one log-weight per action token: log_weight is [*B, *token_dims, 1]
# while the advantage is one value per decision, [*B, 1]. Broadcast the
# decision advantage across the token dims so the surrogate stays
# elementwise -- the alternative is forcing the caller to pre-expand the
# advantage by hand. No-op when the shapes already match (the standard
# one-ratio-per-sample case), so the outer-product guard still fires on a
# genuinely missing value dimension.
if not isinstance(advantage, torch.Tensor) or not isinstance(
log_weight, torch.Tensor
):
return advantage
if (
advantage.shape == log_weight.shape
or advantage.ndim < 2
or advantage.shape[-1] != 1
or log_weight.shape[-1] != 1
or log_weight.ndim <= advantage.ndim
):
return advantage
batch = advantage.shape[:-1]
# the advantage batch must be a strict prefix of the log-weight's leading
# dims; the extra middle dims are the per-token axes to broadcast over
if log_weight.shape[: len(batch)] != batch:
return advantage
token_dims = log_weight.shape[len(batch) : -1]
return advantage.reshape(*batch, *(1,) * len(token_dims), 1).expand(
*batch, *token_dims, 1
)
[docs]
class PPOLoss(LossModule):
"""A parent PPO loss class.
PPO (Proximal Policy Optimization) is a model-free, online RL algorithm
that makes use of a recorded (batch of)
trajectories to perform several optimization steps, while actively
preventing the updated policy to deviate too
much from its original parameter configuration.
PPO loss can be found in different flavors, depending on the way the
constrained optimization is implemented: ClipPPOLoss and KLPENPPOLoss.
Unlike its subclasses, this class does not implement any regularization
and should therefore be used cautiously.
For more details regarding PPO, refer to: "Proximal Policy Optimization Algorithms",
https://arxiv.org/abs/1707.06347
Args:
actor_network (ProbabilisticTensorDictSequential): policy operator.
Typically, a :class:`~tensordict.nn.ProbabilisticTensorDictSequential` subclass taking observations
as input and outputting an action (or actions) as well as its log-probability value.
critic_network (ValueOperator): value operator. The critic will usually take the observations as input
and return a scalar value (``state_value`` by default) in the output keys.
.. note::
While this loss module does not enforce any specific model mode (train/eval), it is highly recommended
to keep your model in eval mode during RL training to ensure deterministic behavior.
A failure to learn due to a train/eval mode mismatch is often observed when the Effective Sample Size (ESS)
drops or increases significantly (see note below).
.. note::
The PPO loss exposes a couple of additional metrics that can be used to monitor the training process:
- The clip fraction is the ratio of the number of clipped weights in the PPO loss (i.e. the ratio of the number of weights that were clipped to the total number of weights).
- The Effective Sample Size (ESS) is a measure of the effective number of samples in the batch, computed as the inverse of the sum of the squared importance weights.
A value of 1 indicates that the importance weights are all equal to 1 (i.e., the samples are equally weighted).
Any value below 1 indicates that the samples are not equally weighted, and the ESS is a measure of the effective number of samples.
If the value drops or increases significantly, it often indicates issues with the model configuration (such as a train/eval mode mismatch, or a large policy update).
Keyword Args:
entropy_bonus (bool, optional): if ``True``, an entropy bonus will be added to the
loss to favour exploratory policies.
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_entropy`` will control how many
samples will be used to compute this estimate.
Defaults to ``1``.
entropy_coeff: scalar | Mapping[NestedKey, scalar], optional): entropy multiplier when computing the total loss.
* **Scalar**: one value applied to the summed entropy of every action head.
* **Mapping** ``{head_name: coeff}`` gives an individual coefficient for each action-head's entropy.
Defaults to ``0.01``.
See :ref:`ppo_entropy_coefficients` for detailed usage examples and troubleshooting.
log_explained_variance (bool, optional): if ``True``, the explained variance of the critic
predictions w.r.t. value targets will be computed and logged as ``"explained_variance"``.
This can help monitor critic quality during training. Best possible score is 1.0, lower values are worse. Defaults to ``True``.
critic_coeff (scalar, optional): critic loss multiplier when computing the total
loss. Defaults to ``1.0``. Set ``critic_coeff`` to ``None`` to exclude the value
loss from the forward outputs.
loss_critic_type (str, optional): loss function for the value discrepancy.
Can be one of "l1", "l2" or "smooth_l1". Defaults to ``"smooth_l1"``.
normalize_advantage (bool, optional): if ``True``, the advantage will be normalized
before being used. Defaults to ``False``.
normalize_advantage_exclude_dims (Tuple[int], optional): dimensions to exclude from the advantage
standardization. Negative dimensions are valid. This is useful in multiagent (or multiobjective) settings
where the agent (or objective) dimension may be excluded from the reductions. Default: ().
advantage_norm (ValueNorm, optional): a stateful
:class:`~torchrl.modules.ValueNorm` used to rescale the advantage,
e.g. :class:`~torchrl.modules.PercentileValueNorm` for
DreamerV3-style return normalization. The advantage is divided by
``advantage_norm.scale()`` without re-centering (the advantage is
already centred by the value baseline). In training mode, the
statistics are updated with the fresh value targets when the loss
runs its own value estimator (i.e. when the input tensordict
carries no advantage entry), so repeated forwards over the same
precomputed rollout do not skew the moving average; rows marked
invalid by the validity masks (see
:attr:`~torchrl.objectives.LossModule.loss_mask_key`) are excluded
from the update. When the advantage is precomputed outside the
loss, call ``advantage_norm.update(value_target)`` once per rollout
yourself. Mutually exclusive with ``normalize_advantage``.
Defaults to ``None``.
separate_losses (bool, optional): if ``True``, shared parameters between
policy and critic will only be trained on the policy loss.
Defaults to ``False``, i.e., 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 ``"advantage"``.
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_target"``.
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 ``"state_value"``.
functional (bool, optional): whether modules should be functionalized.
Functionalizing permits features like meta-RL, but makes it
impossible to use distributed models (DDP, FSDP, ...) and comes
with a little cost. Defaults to ``True``.
reduction (str, optional): Specifies the reduction to apply to the output:
``"none"`` | ``"mean"`` | ``"sum"``. ``"none"``: no reduction will be applied,
``"mean"``: the sum of the output will be divided by the number of
elements in the output, ``"sum"``: the output will be summed. Default: ``"mean"``.
clip_value (:obj:`float`, optional): If provided, it will be used to compute a clipped version of the value
prediction with respect to the input tensordict value estimate and use it to calculate the value loss.
The purpose of clipping is to limit the impact of extreme value predictions, helping stabilize training
and preventing large updates. However, it will have no impact if the value estimate was done by the current
version of the value estimator. Defaults to ``None``.
delay_actor (bool, optional): if ``True``, a detached copy of the actor parameters is kept under
``target_actor_network_params`` and used as the *proximal policy* of the objective, decoupled from
the *behavior policy* that collected the data (whose log-probabilities are read from the
``sample_log_prob`` entry). The surrogate stays weighted by the behavior ratio ``pi_theta / pi_behav``
(so the gradient estimate remains unbiased for the data at hand) while the proximal policy is used
by the trust-region term of the subclasses (clipping in :class:`~torchrl.objectives.ClipPPOLoss`,
KL penalty in :class:`~torchrl.objectives.KLPENPPOLoss`) and by the ``kl_approx`` diagnostic.
Updating the target parameters with :class:`~torchrl.objectives.SoftUpdate` after every optimizer
step makes the proximal policy an exponentially-weighted moving average of the policy, i.e. PPO-EWMA
("Batch size-invariance for policy optimization", Hilton et al., 2021,
https://arxiv.org/abs/2110.00641). Requires ``functional=True``. Defaults to ``False``.
max_importance_ratio (:obj:`float`, optional): if provided, the behavior ratio ``pi_theta / pi_behav`` is
capped at this value. Stale data, or a proximal policy that drifted away from the behavior policy,
can otherwise produce arbitrarily large ratios. The cap lifts the (gradient-free) behavior
log-probability rather than clamping the ratio, so capped samples keep a rescaled policy gradient,
as in the reference PPO-EWMA implementation (which uses ``100.0``). Defaults to ``None`` (no cap).
device (torch.device, optional): device of the buffers. Defaults to ``None``.
.. note:: Parameters and buffers from the policy / critic will not be cast to that device to ensure that
the storages match the ones that are passed to other components, such as data collectors.
.. note::
The advantage (typically GAE) can be computed by the loss function or
in the training loop. The latter option is usually preferred, but this is
up to the user to choose which option is to be preferred.
If the advantage key (``"advantage`` by default) is not present in the
input tensordict, the advantage will be computed by the :meth:`~.forward`
method.
>>> ppo_loss = PPOLoss(actor, critic)
>>> advantage = GAE(critic)
>>> data = next(datacollector)
>>> losses = ppo_loss(data)
>>> # equivalent
>>> advantage(data)
>>> losses = ppo_loss(data)
A custom advantage module can be built using :meth:`~.make_value_estimator`.
The default is :class:`~torchrl.objectives.value.GAE` with hyperparameters
dictated by :func:`~torchrl.objectives.utils.default_value_kwargs`.
>>> ppo_loss = PPOLoss(actor, critic)
>>> ppo_loss.make_value_estimator(ValueEstimators.TDLambda)
>>> data = next(datacollector)
>>> losses = ppo_loss(data)
.. note::
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 = ActorValueOperator(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.
Examples:
>>> import torch
>>> from torch import nn
>>> from torchrl.data.tensor_specs import Bounded
>>> from torchrl.modules.distributions import NormalParamExtractor, TanhNormal
>>> from torchrl.modules.tensordict_module.actors import ProbabilisticActor, ValueOperator
>>> from torchrl.modules.tensordict_module.common import SafeModule
>>> from torchrl.objectives.ppo import PPOLoss
>>> from tensordict import TensorDict
>>> n_act, n_obs = 4, 3
>>> spec = Bounded(-torch.ones(n_act), torch.ones(n_act), (n_act,))
>>> base_layer = nn.Linear(n_obs, 5)
>>> net = nn.Sequential(base_layer, nn.Linear(5, 2 * n_act), NormalParamExtractor())
>>> module = SafeModule(net, in_keys=["observation"], out_keys=["loc", "scale"])
>>> actor = ProbabilisticActor(
... module=module,
... distribution_class=TanhNormal,
... in_keys=["loc", "scale"],
... spec=spec)
>>> module = nn.Sequential(base_layer, nn.Linear(5, 1))
>>> value = ValueOperator(
... module=module,
... in_keys=["observation"])
>>> loss = PPOLoss(actor, value)
>>> batch = [2, ]
>>> action = spec.rand(batch)
>>> data = TensorDict({"observation": torch.randn(*batch, n_obs),
... "action": action,
... "action_log_prob": torch.randn_like(action[..., 1]),
... ("next", "done"): torch.zeros(*batch, 1, dtype=torch.bool),
... ("next", "terminated"): torch.zeros(*batch, 1, dtype=torch.bool),
... ("next", "reward"): torch.randn(*batch, 1),
... ("next", "observation"): torch.randn(*batch, n_obs),
... }, batch)
>>> loss(data)
TensorDict(
fields={
entropy: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
explained_variance: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
kl_approx: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
loss_critic: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
loss_entropy: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
loss_objective: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([]),
device=None,
is_shared=False)
This class is compatible with non-tensordict based modules too and can be
used without recurring to any tensordict-related primitive. In this case,
the expected keyword arguments are:
``["action", "sample_log_prob", "next_reward", "next_done", "next_terminated"]`` + in_keys of the actor and value network.
The return value is a tuple of tensors in the following order:
``["loss_objective"]`` + ``["entropy", "loss_entropy"]`` if entropy_bonus is set + ``"loss_critic"`` if critic_coeff is not ``None``.
The output keys can also be filtered using :meth:`PPOLoss.select_out_keys` method.
Examples:
>>> import torch
>>> from torch import nn
>>> from torchrl.data.tensor_specs import Bounded
>>> from torchrl.modules.distributions import NormalParamExtractor, TanhNormal
>>> from torchrl.modules.tensordict_module.actors import ProbabilisticActor, ValueOperator
>>> from torchrl.modules.tensordict_module.common import SafeModule
>>> from torchrl.objectives.ppo import PPOLoss
>>> n_act, n_obs = 4, 3
>>> spec = Bounded(-torch.ones(n_act), torch.ones(n_act), (n_act,))
>>> base_layer = nn.Linear(n_obs, 5)
>>> net = nn.Sequential(base_layer, nn.Linear(5, 2 * n_act), NormalParamExtractor())
>>> module = SafeModule(net, in_keys=["observation"], out_keys=["loc", "scale"])
>>> actor = ProbabilisticActor(
... module=module,
... distribution_class=TanhNormal,
... in_keys=["loc", "scale"],
... spec=spec)
>>> module = nn.Sequential(base_layer, nn.Linear(5, 1))
>>> value = ValueOperator(
... module=module,
... in_keys=["observation"])
>>> loss = PPOLoss(actor, value)
>>> loss.set_keys(sample_log_prob="sampleLogProb")
>>> _ = loss.select_out_keys("loss_objective")
>>> batch = [2, ]
>>> action = spec.rand(batch)
>>> loss_objective = loss(
... observation=torch.randn(*batch, n_obs),
... action=action,
... sampleLogProb=torch.randn_like(action[..., 1]) / 10,
... next_done=torch.zeros(*batch, 1, dtype=torch.bool),
... next_terminated=torch.zeros(*batch, 1, dtype=torch.bool),
... next_reward=torch.randn(*batch, 1),
... next_observation=torch.randn(*batch, n_obs))
>>> loss_objective.backward()
**Simple Entropy Coefficient Examples**:
>>> # Scalar entropy coefficient (default behavior)
>>> loss = PPOLoss(actor, critic, entropy_coeff=0.01)
>>>
>>> # Per-head entropy coefficients (for composite action spaces)
>>> entropy_coeff = {
... ("agent0", "action_log_prob"): 0.01, # Low exploration
... ("agent1", "action_log_prob"): 0.05, # High exploration
... }
>>> loss = PPOLoss(actor, critic, entropy_coeff=entropy_coeff)
.. note::
There is an exception regarding compatibility with non-tensordict-based modules.
If the actor network is probabilistic and uses a :class:`~tensordict.nn.distributions.CompositeDistribution`,
this class must be used with tensordicts and cannot function as a tensordict-independent module.
This is because composite action spaces inherently rely on the structured representation of data provided by
tensordicts to handle their actions.
.. _ppo_entropy_coefficients:
.. note::
**Entropy Bonus and Coefficient Management**
The entropy bonus encourages exploration by adding the negative entropy of the policy to the loss.
This can be configured in two ways:
**Scalar Coefficient (Default)**: Use a single coefficient for all action heads:
>>> loss = PPOLoss(actor, critic, entropy_coeff=0.01)
**Per-Head Coefficients**: Use different coefficients for different action components:
>>> # For a robot with movement and gripper actions
>>> entropy_coeff = {
... ("agent0", "action_log_prob"): 0.01, # Movement: low exploration
... ("agent1", "action_log_prob"): 0.05, # Gripper: high exploration
... }
>>> loss = PPOLoss(actor, critic, entropy_coeff=entropy_coeff)
**Key Requirements**: When using per-head coefficients, you must provide the full nested key
path to each action head's log probability (e.g., `("agent0", "action_log_prob")`).
**Monitoring Entropy Loss**:
When using composite action spaces, the loss output includes:
- `"entropy"`: Summed entropy across all action heads (for logging)
- `"composite_entropy"`: Individual entropy values for each action head
- `"loss_entropy"`: The weighted entropy loss term
Example output:
>>> result = loss(data)
>>> print(result["entropy"]) # Total entropy: 2.34
>>> print(result["composite_entropy"]) # Per-head: {"movement": 1.2, "gripper": 1.14}
>>> print(result["loss_entropy"]) # Weighted loss: -0.0234
**Common Issues**:
**KeyError: "Missing entropy coeff for head 'head_name'"**:
- Ensure you provide coefficients for ALL action heads
- Use full nested keys: `("head_name", "action_log_prob")`
- Check that your action space structure matches the coefficient mapping
**Incorrect Entropy Calculation**:
- Call `set_composite_lp_aggregate(False).set()` before creating your policy
- Verify that your action space uses :class:`~tensordict.nn.distributions.CompositeDistribution`
"""
@dataclass
class _AcceptedKeys:
"""Maintains default values for all configurable tensordict keys.
This class defines which tensordict keys can be set using '.set_keys(key_name=key_value)' and their
default values
Attributes:
advantage (NestedKey): The input tensordict key where the advantage is expected.
Will be used for the underlying value estimator. Defaults to ``"advantage"``.
value_target (NestedKey): The input tensordict key where the target state value is expected.
Will be used for the underlying value estimator Defaults to ``"value_target"``.
value (NestedKey): The input tensordict key where the state value is expected.
Will be used for the underlying value estimator. Defaults to ``"state_value"``.
sample_log_prob (NestedKey or list of nested keys): The input tensordict key where the
sample log probability is expected.
Defaults to ``"sample_log_prob"`` when :func:`~tensordict.nn.composite_lp_aggregate` returns `True`,
`"action_log_prob"` otherwise.
action (NestedKey or list of nested keys): The input tensordict key where the action is expected.
Defaults to ``"action"``.
reward (NestedKey or list of nested keys): The input tensordict key where the reward is expected.
Will be used for the underlying value estimator. Defaults to ``"reward"``.
done (NestedKey or list of nested keys): The key in the input TensorDict that indicates
whether a trajectory is done. Will be used for the underlying value estimator.
Defaults to ``"done"``.
terminated (NestedKey or list of nested keys): The key in the input TensorDict that indicates
whether a trajectory is terminated. Will be used for the underlying value estimator.
Defaults to ``"terminated"``.
"""
advantage: NestedKey = "advantage"
value_target: NestedKey = "value_target"
value: NestedKey = "state_value"
sample_log_prob: NestedKey | list[NestedKey] | None = None
action: NestedKey | list[NestedKey] = "action"
reward: NestedKey | list[NestedKey] = "reward"
done: NestedKey | list[NestedKey] = "done"
terminated: NestedKey | list[NestedKey] = "terminated"
def __post_init__(self):
if self.sample_log_prob is None:
if composite_lp_aggregate(nowarn=True):
self.sample_log_prob = "sample_log_prob"
else:
self.sample_log_prob = "action_log_prob"
default_keys = _AcceptedKeys
tensor_keys: _AcceptedKeys
default_value_estimator = ValueEstimators.GAE
_schedulable_buffers = frozenset(
{"entropy_coeff", "critic_coeff", "clip_value", "max_importance_ratio"}
)
actor_network: ProbabilisticTensorDictModule
critic_network: TensorDictModule
actor_network_params: TensorDictParams
critic_network_params: TensorDictParams
target_actor_network_params: TensorDictParams
target_critic_network_params: TensorDictParams
def __init__(
self,
actor_network: ProbabilisticTensorDictSequential | None = None,
critic_network: TensorDictModule | None = None,
*,
entropy_bonus: bool = True,
samples_mc_entropy: int = 1,
entropy_coeff: float | Mapping[NestedKey, float] | None = None,
log_explained_variance: bool = True,
critic_coeff: float | None = None,
loss_critic_type: str = "smooth_l1",
normalize_advantage: bool = False,
normalize_advantage_exclude_dims: tuple[int] = (),
advantage_norm: ValueNorm | None = None,
gamma: float | None = None,
separate_losses: bool = False,
advantage_key: str | None = None,
value_target_key: str | None = None,
value_key: str | None = None,
functional: bool = True,
actor: ProbabilisticTensorDictSequential = None,
critic: ProbabilisticTensorDictSequential = None,
reduction: str | None = None,
clip_value: float | None = None,
delay_actor: bool = False,
max_importance_ratio: float | None = None,
device: torch.device | None = None,
**kwargs,
):
if actor is not None:
actor_network = actor
del actor
if critic is not None:
critic_network = critic
del critic
# critic_coef has been removed in v0.11
if "critic_coef" in kwargs:
raise TypeError(
"'critic_coef' has been removed in torchrl v0.11. Please use 'critic_coeff' instead."
)
if critic_coeff is None and critic_network is not None:
critic_coeff = 1.0
elif critic_coeff in (None, 0) and critic_network is not None:
critic_coeff = None
if actor_network is None or (
critic_network is None and critic_coeff not in (None, 0.0)
):
raise TypeError(
"Missing positional arguments actor_network or critic_network."
)
if reduction is None:
reduction = "mean"
self._functional = functional
self._in_keys = None
self._out_keys = None
super().__init__()
if functional:
self.convert_to_functional(
actor_network, "actor_network", create_target_params=delay_actor
)
elif delay_actor:
raise ValueError(
"delay_actor=True requires functional=True: the proximal policy is "
"held in target_actor_network_params, which only exists for "
"functionalized modules."
)
else:
self.actor_network = actor_network
self.actor_network_params = None
self.target_actor_network_params = None
self.delay_actor = delay_actor
if separate_losses:
# we want to make sure there are no duplicates in the params: the
# params of critic must be refs to actor if they're shared
policy_params = list(actor_network.parameters())
else:
policy_params = None
if functional and critic_network is not None:
self.convert_to_functional(
critic_network, "critic_network", compare_against=policy_params
)
else:
self.critic_network = critic_network
self.critic_network_params = None
self.target_critic_network_params = None
self.log_explained_variance = log_explained_variance
self.samples_mc_entropy = samples_mc_entropy
self.entropy_bonus = entropy_bonus
self.separate_losses = separate_losses
self.reduction = reduction
if device is None:
try:
device = next(self.parameters()).device
except (AttributeError, StopIteration):
device = getattr(
torch, "get_default_device", lambda: torch.device("cpu")
)()
# entropy_coef has been removed in v0.11
if "entropy_coef" in kwargs:
raise TypeError(
"'entropy_coef' has been removed in torchrl v0.11. Please use 'entropy_coeff' instead."
)
# Set default value if None
if entropy_coeff is None:
entropy_coeff = 0.01
if isinstance(entropy_coeff, Mapping):
# Store the mapping for per-head coefficients
self._entropy_coeff_map = {k: float(v) for k, v in entropy_coeff.items()}
# Register an empty buffer for compatibility
self.register_coeff_buffer("entropy_coeff", 0.0, device=device)
elif isinstance(entropy_coeff, (float, int, torch.Tensor)):
self.register_coeff_buffer("entropy_coeff", entropy_coeff, device=device)
self._entropy_coeff_map = None
else:
raise TypeError("entropy_coeff must be a float or a Mapping[str, float]")
self.register_coeff_buffer("critic_coeff", critic_coeff, device=device)
self._has_critic = bool(self.critic_coeff is not None and self.critic_coeff > 0)
self.loss_critic_type = loss_critic_type
if advantage_norm is not None and normalize_advantage:
raise ValueError(
"normalize_advantage=True and advantage_norm are mutually "
"exclusive: pass either the mean/std standardization flag or "
"a ValueNorm instance, not both."
)
self.normalize_advantage = normalize_advantage
self.normalize_advantage_exclude_dims = normalize_advantage_exclude_dims
self.advantage_norm = advantage_norm
if gamma is not None:
raise TypeError(_GAMMA_LMBDA_DEPREC_ERROR)
self._set_deprecated_ctor_keys(
advantage=advantage_key,
value_target=value_target_key,
value=value_key,
)
self.register_coeff_buffer("clip_value", clip_value, device=device)
if max_importance_ratio is not None and max_importance_ratio <= 0:
raise ValueError(
f"max_importance_ratio must be > 0, got {max_importance_ratio}."
)
self.register_coeff_buffer(
"max_importance_ratio", max_importance_ratio, device=device
)
try:
log_prob_keys = self.actor_network.log_prob_keys
action_keys = self.actor_network.dist_sample_keys
if len(log_prob_keys) > 1:
self.set_keys(sample_log_prob=log_prob_keys, action=action_keys)
else:
self.set_keys(sample_log_prob=log_prob_keys[0], action=action_keys[0])
except AttributeError:
pass
@property
def functional(self):
return self._functional
def _set_in_keys(self):
keys = []
_maybe_add_or_extend_key(keys, self.actor_network.in_keys)
_maybe_add_or_extend_key(keys, self.actor_network.in_keys, "next")
if self.critic_network is not None:
_maybe_add_or_extend_key(keys, self.critic_network.in_keys)
_maybe_add_or_extend_key(keys, self.tensor_keys.action)
_maybe_add_or_extend_key(keys, self.tensor_keys.sample_log_prob)
_maybe_add_or_extend_key(keys, self.tensor_keys.reward, "next")
_maybe_add_or_extend_key(keys, self.tensor_keys.done, "next")
_maybe_add_or_extend_key(keys, self.tensor_keys.terminated, "next")
self._in_keys = list(set(keys))
@property
def in_keys(self):
if self._in_keys is None:
self._set_in_keys()
return self._in_keys
@in_keys.setter
def in_keys(self, values):
self._in_keys = values
@property
def out_keys(self):
if self._out_keys is None:
keys = ["loss_objective"]
if self.entropy_bonus:
keys.extend(["entropy", "loss_entropy"])
if self.loss_critic:
keys.append("loss_critic")
if self.clip_value:
keys.append("value_clip_fraction")
self._out_keys = keys
return self._out_keys
@out_keys.setter
def out_keys(self, values):
self._out_keys = values
def _forward_value_estimator_keys(self, **kwargs) -> None:
if hasattr(self, "_value_estimator") and self._value_estimator is not None:
self._value_estimator.set_keys(
advantage=self.tensor_keys.advantage,
value_target=self.tensor_keys.value_target,
value=self.tensor_keys.value,
reward=self.tensor_keys.reward,
done=self.tensor_keys.done,
terminated=self.tensor_keys.terminated,
sample_log_prob=self.tensor_keys.sample_log_prob,
)
self._set_in_keys()
def reset(self) -> None:
pass
def _get_entropy(
self, dist: d.Distribution, adv_shape: torch.Size
) -> torch.Tensor | TensorDict:
is_composite = isinstance(dist, CompositeDistribution)
if is_composite:
can_compute_component_entropy = all(
has_analytic_entropy(component) or component.has_rsample
for component in dist.dists.values()
)
if can_compute_component_entropy:
entropy = composite_entropy(dist, self.samples_mc_entropy)
else:
if not is_dynamo_compiling():
_warn_mc_entropy(dist)
with set_composite_lp_aggregate(False):
_, log_prob = sample_and_log_prob(
dist,
(self.samples_mc_entropy,),
reparameterize=getattr(dist, "has_rsample", False),
)
if isinstance(self.tensor_keys.sample_log_prob, NestedKey):
log_prob = log_prob.get(self.tensor_keys.sample_log_prob)
else:
log_prob = log_prob.select(*self.tensor_keys.sample_log_prob)
entropy = -log_prob.mean(0)
else:
analytic_entropy = has_analytic_entropy(dist)
if analytic_entropy:
entropy = dist.entropy()
compiling = is_dynamo_compiling()
needs_mc = not analytic_entropy or compiling
if analytic_entropy and not compiling and not entropy.isfinite().all():
needs_mc = True
if needs_mc:
if not compiling:
_warn_mc_entropy(dist)
_, log_prob = sample_and_log_prob(
dist,
(self.samples_mc_entropy,),
reparameterize=dist.has_rsample,
)
sampled_entropy = -log_prob.mean(0)
if analytic_entropy and compiling:
entropy = torch.where(entropy.isfinite(), entropy, sampled_entropy)
else:
entropy = sampled_entropy
if is_composite:
if is_tensor_collection(entropy) and entropy.batch_size != adv_shape:
entropy.batch_size = adv_shape
return entropy.unsqueeze(-1)
def _get_cur_log_prob(
self, tensordict: TensorDictBase, params: TensorDictParams | None = None
):
# ``params`` selects the parameter set the actor runs with (the current
# parameters by default; the target parameters for the proximal policy).
if isinstance(
self.actor_network,
(ProbabilisticTensorDictSequential, ProbabilisticTensorDictModule),
) or hasattr(self.actor_network, "get_dist"):
# assert tensordict['log_probs'].requires_grad
# assert tensordict['logits'].requires_grad
if params is None:
params = self.actor_network_params
with (
params.to_module(self.actor_network, preserve_module_state=False)
if self.functional
else contextlib.nullcontext()
):
dist = self.actor_network.get_dist(tensordict)
is_composite = isinstance(dist, CompositeDistribution)
if is_composite:
action = tensordict.select(
*(
(self.tensor_keys.action,)
if isinstance(self.tensor_keys.action, NestedKey)
else self.tensor_keys.action
)
)
else:
action = _maybe_get_or_select(tensordict, self.tensor_keys.action)
if action.requires_grad:
raise RuntimeError(
f"tensordict stored {self.tensor_keys.action} requires grad."
)
log_prob = dist.log_prob(action)
else:
raise NotImplementedError(
"Only probabilistic modules from tensordict.nn are currently supported. "
"If you need to implement a custom logic to retrieve the log-probs (to compute "
"the PPO objective) or the distribution (for the PPO entropy), please augment "
f"the {type(self).__class__} by implementing your own logic in _get_cur_log_prob."
)
# with self.actor_network_params.to_module(
# self.actor_network
# ) if self.functional else contextlib.nullcontext():
# td = self.actor_network(tensordict)
# log_prob = td.get(self.tensor_keys.sample_log_prob)
# dist = torch.distributions.Categorical(td.get("logits"))
# is_composite = False
return log_prob, dist, is_composite
def _log_weight(
self, tensordict: TensorDictBase, adv_shape: torch.Size
) -> tuple[torch.Tensor, d.Distribution, torch.Tensor, torch.Tensor | None]:
# Returns
# log_weight: log(pi_theta / pi_prox), the ratio the trust region acts on;
# dist: the current policy distribution;
# kl_approx: log(pi_prox / pi_theta), a diagnostic;
# log_is_weight: log(pi_prox / pi_behav), the gradient-free correction
# that turns log_weight into the behavior ratio
# log(pi_theta / pi_behav) the surrogate must be weighted by.
# ``None`` when the proximal policy is the behavior policy and no
# ratio cap is set, in which case log_weight already is that ratio.
prev_log_prob = _maybe_get_or_select(
tensordict,
self.tensor_keys.sample_log_prob,
adv_shape,
)
if prev_log_prob is None:
raise KeyError(
f"Couldn't find the log-prob {self.tensor_keys.sample_log_prob} in the input data."
)
if prev_log_prob.requires_grad:
raise RuntimeError(
f"tensordict stored {self.tensor_keys.sample_log_prob} requires grad."
)
log_prob, dist, is_composite = self._get_cur_log_prob(tensordict)
if self.delay_actor:
with torch.no_grad():
prox_log_prob, _, _ = self._get_cur_log_prob(
tensordict, params=self.target_actor_network_params
)
else:
prox_log_prob = prev_log_prob
if is_composite:
with set_composite_lp_aggregate(False):
if not is_tensor_collection(prev_log_prob):
# this isn't great: in general, multi-head actions should have a composite log-prob too
warnings.warn(
"You are using a composite distribution, yet your log-probability is a tensor. "
"Make sure you have called tensordict.nn.set_composite_lp_aggregate(False).set() at "
"the beginning of your script to get a proper composite log-prob.",
category=UserWarning,
)
if is_tensor_collection(log_prob):
log_prob = _sum_td_features(log_prob)
log_prob.view_as(prev_log_prob)
if is_tensor_collection(prox_log_prob):
prox_log_prob = _sum_td_features(prox_log_prob)
if is_tensor_collection(log_prob) and log_prob.batch_size != adv_shape:
log_prob.batch_size = adv_shape
if (
is_tensor_collection(prox_log_prob)
and prox_log_prob.batch_size != adv_shape
):
prox_log_prob.batch_size = adv_shape
log_weight = (log_prob - prox_log_prob).unsqueeze(-1)
if is_tensor_collection(log_weight):
log_weight = _sum_td_features(log_weight)
log_weight = log_weight.view(adv_shape).unsqueeze(-1)
kl_approx = (prox_log_prob - log_prob).unsqueeze(-1)
if is_tensor_collection(kl_approx):
kl_approx = _sum_td_features(kl_approx)
log_is_weight = None
if self.delay_actor:
log_is_weight = (prox_log_prob - prev_log_prob).unsqueeze(-1)
if is_tensor_collection(log_is_weight):
log_is_weight = _sum_td_features(log_is_weight)
log_is_weight = log_is_weight.view(adv_shape).unsqueeze(-1)
if self.max_importance_ratio is not None:
if log_is_weight is None:
log_is_weight = torch.zeros_like(log_weight)
# Cap pi_theta / pi_behav at max_importance_ratio by lifting the
# (gradient-free) behavior log-prob, as the reference PPO-EWMA
# implementation does: a capped sample keeps a rescaled policy
# gradient instead of the zero gradient a clamp on the ratio gives.
log_is_weight = torch.minimum(
log_is_weight, self.max_importance_ratio.log() - log_weight.detach()
)
return log_weight, dist, kl_approx, log_is_weight
def _critic_loss_inputs(
self,
target_return: torch.Tensor,
state_value: torch.Tensor,
old_state_value: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
"""Transform the critic-loss inputs before distance / clip.
Default implementation is identity. Subclasses (notably
:class:`~torchrl.objectives.multiagent.MAPPOLoss`) override this to
push the three tensors through a running
:class:`~torchrl.modules.ValueNorm` so the MSE / smooth-L1 distance
lives on a fixed scale -- and so the PPO value-clip radius stays
meaningful when reward scales drift.
"""
return target_return, state_value, old_state_value
[docs]
def loss_critic(
self, tensordict: TensorDictBase
) -> tuple[torch.Tensor | TensorDict, ...]:
"""Returns the critic loss multiplied by ``critic_coeff``, if it is not ``None``."""
# TODO: if the advantage is gathered by forward, this introduces an
# overhead that we could easily reduce.
if self.separate_losses:
tensordict = tensordict.detach()
target_return = tensordict.get(
self.tensor_keys.value_target, None
) # TODO: None soon to be removed
if target_return is None:
raise KeyError(
f"the key {self.tensor_keys.value_target} was not found in the input tensordict. "
f"Make sure you provided the right key and the value_target (i.e. the target "
f"return) has been retrieved accordingly. Advantage classes such as GAE, "
f"TDLambdaEstimate and TDEstimate all return a 'value_target' entry that "
f"can be used for the value loss."
)
old_state_value = None
if self.clip_value:
old_state_value = tensordict.get(self.tensor_keys.value)
if old_state_value is None:
raise KeyError(
f"clip_value is set to {self.clip_value}, but "
f"the key {self.tensor_keys.value} was not found in the input tensordict. "
f"Make sure that the 'value_key' passed to PPO exists in the input tensordict."
)
with (
self.critic_network_params.to_module(
self.critic_network, preserve_module_state=False
)
if self.functional
else contextlib.nullcontext()
):
state_value_td = self.critic_network(tensordict)
state_value = state_value_td.get(self.tensor_keys.value)
if state_value is None:
raise KeyError(
f"the key {self.tensor_keys.value} was not found in the critic output tensordict. "
f"Make sure that the 'value_key' passed to PPO is accurate."
)
# Subclass hook: lets e.g. MAPPOLoss inject PopArt-style value
# normalisation uniformly across target_return / state_value /
# old_state_value so the MSE and the clip radius both live in
# normalised space.
target_return, state_value, old_state_value = self._critic_loss_inputs(
target_return, state_value, old_state_value
)
loss_value = distance_loss(
target_return,
state_value,
loss_function=self.loss_critic_type,
)
clip_fraction = None
if self.clip_value:
loss_value, clip_fraction = _clip_value_loss(
old_state_value,
state_value,
self.clip_value,
target_return,
loss_value,
self.loss_critic_type,
)
explained_variance = None
if self.log_explained_variance:
with torch.no_grad(): # <‑‑ break grad‐flow
tgt = target_return.detach()
pred = state_value.detach()
eps = torch.finfo(tgt.dtype).eps
resid = torch.var(tgt - pred, correction=0, dim=0)
total = torch.var(tgt, correction=0, dim=0)
explained_variance = 1.0 - resid / (total + eps)
self._clear_weakrefs(
tensordict,
"actor_network_params",
"critic_network_params",
"target_actor_network_params",
"target_critic_network_params",
)
if self._has_critic:
return self.critic_coeff * loss_value, clip_fraction, explained_variance
return loss_value, clip_fraction, explained_variance
@property
@_cache_values
def _cached_critic_network_params_detached(self):
if not self.functional:
return None
return self.critic_network_params.detach()
def _standardize_advantage(
self, advantage: torch.Tensor, tensordict: TensorDictBase
) -> torch.Tensor:
mask = tensordict.get("shifted_valid", default=None)
if mask is None:
return _standardize(advantage, self.normalize_advantage_exclude_dims)
while mask.ndim < advantage.ndim:
mask = mask.unsqueeze(-1)
mask = mask.expand_as(advantage).to(advantage.dtype)
exclude_dims = {
dim if dim >= 0 else advantage.ndim + dim
for dim in self.normalize_advantage_exclude_dims
}
reduce_dims = [dim for dim in range(advantage.ndim) if dim not in exclude_dims]
count = mask.sum(dim=reduce_dims, keepdim=True).clamp_min(1)
loc = (advantage * mask).sum(dim=reduce_dims, keepdim=True) / count
scale = (
(
((advantage - loc).pow(2) * mask).sum(dim=reduce_dims, keepdim=True)
/ count
)
.sqrt()
.clamp_min(1e-4)
)
return (advantage - loc) / scale
def _maybe_normalize_advantage(
self,
advantage: torch.Tensor,
tensordict: TensorDictBase,
*,
update_norm: bool,
) -> torch.Tensor:
if self.advantage_norm is not None:
return self._scale_advantage(advantage, tensordict, update=update_norm)
if self.normalize_advantage and advantage.numel() > 1:
if advantage.numel() > tensordict.batch_size.numel() and not len(
self.normalize_advantage_exclude_dims
):
warnings.warn(
"You requested advantage normalization and the advantage key has more dimensions"
" than the tensordict batch. Make sure to pass `normalize_advantage_exclude_dims` "
"if you want to keep any dimension independent while computing normalization statistics. "
"If you are working in multi-agent/multi-objective settings this is highly suggested."
)
advantage = self._standardize_advantage(advantage, tensordict)
return advantage
def _scale_advantage(
self,
advantage: torch.Tensor,
tensordict: TensorDictBase,
*,
update: bool,
) -> torch.Tensor:
if update and self.training:
value_target = tensordict.get(self.tensor_keys.value_target, None)
if value_target is None:
raise KeyError(
f"advantage_norm requires the value target "
f"({self.tensor_keys.value_target!r}) in the input "
"tensordict to update the return statistics; the value "
"estimator is expected to write it."
)
value_target = _valid_value_target_rows(
value_target, tensordict, self._loss_mask_keys()
)
self.advantage_norm.update(value_target)
return advantage / self.advantage_norm.scale()
[docs]
@dispatch
def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
tensordict = tensordict.clone(False)
advantage = tensordict.get(self.tensor_keys.advantage, None)
update_norm = advantage is None
if advantage is None:
self.value_estimator(
tensordict,
params=self._cached_critic_network_params_detached,
target_params=self.target_critic_network_params,
)
advantage = tensordict.get(self.tensor_keys.advantage)
advantage = self._maybe_normalize_advantage(
advantage, tensordict, update_norm=update_norm
)
log_weight, dist, kl_approx, log_is_weight = self._log_weight(
tensordict, adv_shape=advantage.shape[:-1]
)
advantage = _broadcast_advantage_to_log_weight(advantage, log_weight)
_check_advantage_broadcast(advantage, log_weight)
if log_is_weight is not None:
# the surrogate is weighted by the behavior ratio pi_theta / pi_behav
log_weight = log_weight + log_is_weight
neg_loss = log_weight.exp() * advantage
td_out = TensorDict({"loss_objective": -neg_loss})
td_out.set("kl_approx", kl_approx.detach().mean()) # for logging
if self.entropy_bonus:
entropy = self._get_entropy(dist, adv_shape=advantage.shape[:-1])
if is_tensor_collection(entropy):
# Reports the entropy of each action head.
td_out.set("composite_entropy", entropy.detach())
td_out.set(
"entropy", _sum_td_features(entropy).detach().mean()
) # for logging
else:
td_out.set("entropy", entropy.detach().mean()) # for logging
td_out.set("loss_entropy", self._weighted_loss_entropy(entropy))
if self._has_critic:
loss_critic, value_clip_fraction, explained_variance = self.loss_critic(
tensordict
)
td_out.set("loss_critic", loss_critic)
if value_clip_fraction is not None:
td_out.set("value_clip_fraction", value_clip_fraction)
if explained_variance is not None:
td_out.set("explained_variance", explained_variance)
td_out = td_out.named_apply(
lambda name, value: self._reduce_loss(value, tensordict).squeeze(-1)
if name.startswith("loss_")
else value,
)
self._clear_weakrefs(
tensordict,
td_out,
"actor_network_params",
"critic_network_params",
"target_actor_network_params",
"target_critic_network_params",
)
return td_out
SUPPORTED_VALUE_ESTIMATORS = (
ValueEstimators.TD0,
ValueEstimators.TD1,
ValueEstimators.TDLambda,
ValueEstimators.GAE,
ValueEstimators.MAGAE,
ValueEstimators.VTrace,
)
[docs]
def make_value_estimator(self, value_type: ValueEstimators = None, **hyperparams):
value_type, hp = self._prepare_value_estimator_kwargs(value_type, **hyperparams)
if value_type is None:
return self
dispatch_value_estimator(
self,
value_type,
supported=self.SUPPORTED_VALUE_ESTIMATORS,
tensor_keys={
"advantage": self.tensor_keys.advantage,
"value": self.tensor_keys.value,
"value_target": self.tensor_keys.value_target,
"reward": self.tensor_keys.reward,
"done": self.tensor_keys.done,
"terminated": self.tensor_keys.terminated,
"sample_log_prob": self.tensor_keys.sample_log_prob,
},
**hp,
)
def _weighted_loss_entropy(
self, entropy: torch.Tensor | TensorDictBase
) -> torch.Tensor:
"""Compute the weighted entropy loss.
If `self._entropy_coeff_map` is provided, apply per-head entropy coefficients.
Otherwise, use the scalar `self.entropy_coeff`.
The entries in self._entropy_coeff_map require the full nested key to the entropy head.
"""
# Mode 1: Use scalar entropy coefficient (default behavior)
if self._entropy_coeff_map is None:
# If entropy is a TensorDict (composite action space), sum all entropy values
if is_tensor_collection(entropy):
entropy = _sum_td_features(entropy)
# Apply scalar coefficient: loss = -coeff * entropy (negative for maximization)
return -self.entropy_coeff * entropy
# Mode 2: Use per-head entropy coefficients (for composite action spaces)
loss_term = None # Initialize running sum over action heads
coeff = 0 # Placeholder for coefficient value
# Iterate through all entropy heads in the composite action space
for head_name, entropy_head in entropy.items(
include_nested=True, leaves_only=True
):
try:
# Look up the coefficient for this specific action head
coeff = self._entropy_coeff_map[head_name]
except KeyError as exc:
# Provide clear error message if coefficient mapping is incomplete
raise KeyError(f"Missing entropy coeff for head '{head_name}'") from exc
# Convert coefficient to tensor with matching dtype and device
coeff_t = torch.as_tensor(
coeff, dtype=entropy_head.dtype, device=entropy_head.device
)
# Compute weighted loss for this head: -coeff * entropy
head_loss_term = -coeff_t * entropy_head
# Accumulate loss terms across all heads
loss_term = (
head_loss_term if loss_term is None else loss_term + head_loss_term
)
return loss_term
[docs]
class ClipPPOLoss(PPOLoss):
"""Clipped PPO loss.
The clipped importance weighted loss is computed as follows:
loss = -min( weight * advantage, min(max(weight, 1-eps), 1+eps) * advantage)
Args:
actor_network (ProbabilisticTensorDictSequential): policy operator.
critic_network (ValueOperator): value operator.
.. note::
While this loss module does not enforce any specific model mode (train/eval), it is highly recommended
to keep your model in eval mode during RL training to ensure deterministic behavior.
A failure to learn due to a train/eval mode mismatch is often observed when the Effective Sample Size (ESS)
drops or increases significantly (see note below).
.. note::
The PPO loss exposes a couple of additional metrics that can be used to monitor the training process:
- The clip fraction is the ratio of the number of clipped weights in the PPO loss (i.e. the ratio of the number of weights that were clipped to the total number of weights).
- The Effective Sample Size (ESS) is a measure of the effective number of samples in the batch, computed as the inverse of the sum of the squared importance weights.
A value of 1 indicates that the importance weights are all equal to 1 (i.e., the samples are equally weighted).
Any value below 1 indicates that the samples are not equally weighted, and the ESS is a measure of the effective number of samples.
If the value drops or increases significantly, it often indicates issues with the model configuration (such as a train/eval mode mismatch, or a large policy update).
Keyword Args:
clip_epsilon (scalar or tuple of scalars, optional): weight clipping threshold(s) in the clipped
PPO loss equation.
- float ``x``: symmetric clipping ``[1 - x, 1 + x]``. Default: ``0.2``.
- tuple ``(eps_low, eps_high)``: asymmetric clipping ``[1 - eps_low, 1 + eps_high]`` as in
DAPO Clip-Higher (recommended values ``(0.20, 0.28)``; see Eq. (10) of the
`DAPO paper <https://arxiv.org/html/2503.14476>`_). With a tuple, the thresholds are
exposed (and schedulable) as the ``clip_epsilon_low`` / ``clip_epsilon_high`` buffers
instead of ``clip_epsilon``, and ``clip_value=True`` is not allowed (pass an explicit
float threshold instead).
entropy_bonus (bool, optional): if ``True``, an entropy bonus will be added to the
loss to favour exploratory policies.
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_entropy`` will control how many
samples will be used to compute this estimate.
Defaults to ``1``.
entropy_coeff: (scalar | Mapping[NestedKey, scalar], optional): entropy multiplier when computing the total loss.
* **Scalar**: one value applied to the summed entropy of every action head.
* **Mapping** ``{head_name: coeff}`` gives an individual coefficient for each action-head's entropy.
Defaults to ``0.01``.
See :ref:`ppo_entropy_coefficients` for detailed usage examples and troubleshooting.
critic_coeff (scalar, optional): critic loss multiplier when computing the total
loss. Defaults to ``1.0``. Set ``critic_coeff`` to ``None`` to exclude the value
loss from the forward outputs.
loss_critic_type (str, optional): loss function for the value discrepancy.
Can be one of "l1", "l2" or "smooth_l1". Defaults to ``"smooth_l1"``.
normalize_advantage (bool, optional): if ``True``, the advantage will be normalized
before being used. Defaults to ``False``.
normalize_advantage_exclude_dims (Tuple[int], optional): dimensions to exclude from the advantage
standardization. Negative dimensions are valid. This is useful in multiagent (or multiobjective) settings
where the agent (or objective) dimension may be excluded from the reductions. Default: ().
advantage_norm (ValueNorm, optional): a stateful
:class:`~torchrl.modules.ValueNorm` used to rescale the advantage,
e.g. :class:`~torchrl.modules.PercentileValueNorm` for
DreamerV3-style return normalization. The advantage is divided by
``advantage_norm.scale()`` without re-centering (the advantage is
already centred by the value baseline). In training mode, the
statistics are updated with the fresh value targets when the loss
runs its own value estimator (i.e. when the input tensordict
carries no advantage entry), so repeated forwards over the same
precomputed rollout do not skew the moving average; rows marked
invalid by the validity masks (see
:attr:`~torchrl.objectives.LossModule.loss_mask_key`) are excluded
from the update. When the advantage is precomputed outside the
loss, call ``advantage_norm.update(value_target)`` once per rollout
yourself. Mutually exclusive with ``normalize_advantage``.
Defaults to ``None``.
separate_losses (bool, optional): if ``True``, shared parameters between
policy and critic will only be trained on the policy loss.
Defaults to ``False``, i.e., 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 ``"advantage"``.
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_target"``.
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 ``"state_value"``.
functional (bool, optional): whether modules should be functionalized.
Functionalizing permits features like meta-RL, but makes it
impossible to use distributed models (DDP, FSDP, ...) and comes
with a little cost. Defaults to ``True``.
reduction (str, optional): Specifies the reduction to apply to the output:
``"none"`` | ``"mean"`` | ``"sum"``. ``"none"``: no reduction will be applied,
``"mean"``: the sum of the output will be divided by the number of
elements in the output, ``"sum"``: the output will be summed. Default: ``"mean"``.
clip_value (bool or float, optional): If a ``float`` is provided, it will be used to compute a clipped
version of the value prediction with respect to the input tensordict value estimate and use it to
calculate the value loss. The purpose of clipping is to limit the impact of extreme value predictions,
helping stabilize training and preventing large updates. However, it will have no impact if the value
estimate was done by the current version of the value estimator. If instead ``True`` is provided, the
``clip_epsilon`` parameter will be used as the clipping threshold (this is only compatible with a
scalar ``clip_epsilon``; with an asymmetric ``(low, high)`` tuple, pass an explicit float threshold
instead). If not provided or ``False``, no clipping will be performed. Defaults to ``False``.
delay_actor (bool, optional): if ``True``, a detached copy of the actor parameters is kept under
``target_actor_network_params`` and used as the *proximal policy*: the clipping acts on
``pi_theta / pi_prox`` while the surrogate is re-weighted by ``pi_prox / pi_behav``, where
``pi_behav`` is the *behavior policy* that collected the data (``sample_log_prob`` entry), so the
gradient estimate remains unbiased for the data at hand. This decouples the strength of the trust
region from how (and how recently) the data was collected. Updating the target parameters with
:class:`~torchrl.objectives.SoftUpdate` after every optimizer step makes the proximal policy an
exponentially-weighted moving average of the policy, i.e. PPO-EWMA ("Batch size-invariance for
policy optimization", Hilton et al., 2021, https://arxiv.org/abs/2110.00641); see the note below.
Requires ``functional=True``. Defaults to ``False``.
max_importance_ratio (:obj:`float`, optional): if provided, the behavior ratio ``pi_theta / pi_behav`` is
capped at this value. Stale data, or a proximal policy that drifted away from the behavior policy,
can otherwise produce arbitrarily large ratios. The cap lifts the (gradient-free) behavior
log-probability rather than clamping the ratio, so capped samples keep a rescaled policy gradient,
as in the reference PPO-EWMA implementation (which uses ``100.0``). Defaults to ``None`` (no cap).
device (torch.device, optional): device of the buffers. Defaults to ``None``.
.. note:: Parameters and buffers from the policy / critic will not be cast to that device to ensure that
the storages match the ones that are passed to other components, such as data collectors.
.. note:
The advantage (typically GAE) can be computed by the loss function or
in the training loop. The latter option is usually preferred, but this is
up to the user to choose which option is to be preferred.
If the advantage key (``"advantage`` by default) is not present in the
input tensordict, the advantage will be computed by the :meth:`~.forward`
method.
>>> ppo_loss = ClipPPOLoss(actor, critic)
>>> advantage = GAE(critic)
>>> data = next(datacollector)
>>> losses = ppo_loss(data)
>>> # equivalent
>>> advantage(data)
>>> losses = ppo_loss(data)
A custom advantage module can be built using :meth:`~.make_value_estimator`.
The default is :class:`~torchrl.objectives.value.GAE` with hyperparameters
dictated by :func:`~torchrl.objectives.utils.default_value_kwargs`.
>>> ppo_loss = ClipPPOLoss(actor, critic)
>>> ppo_loss.make_value_estimator(ValueEstimators.TDLambda)
>>> data = next(datacollector)
>>> losses = ppo_loss(data)
.. note::
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 = ActorValueOperator(common, actor_head, value_head)
>>> loss_module = ClipPPOLoss(model.get_policy_operator(), model.get_value_operator())
>>> # second option, with a single call to the common module
>>> loss_module = ClipPPOLoss(ProbabilisticTensorDictSequential(model, actor_head), value_head)
This will work regardless of whether separate_losses is activated or not.
.. note::
**Decoupled proximal policy (PPO-EWMA).** With ``delay_actor=True`` the loss follows the decoupled
clipped objective of Hilton et al. (2021),
loss = -(pi_prox / pi_behav) * min(r * advantage, clip(r, 1 - eps, 1 + eps) * advantage),
r = pi_theta / pi_prox,
where ``pi_prox`` runs on ``target_actor_network_params``. Any :class:`~torchrl.objectives.TargetNetUpdater`
controls how old the proximal policy is: a :class:`~torchrl.objectives.SoftUpdate` stepped after every
optimizer step gives the EWMA of the paper (``eps`` being the decay rate ``beta_prox``), a
:class:`~torchrl.objectives.HardUpdate` refreshes it every ``value_network_update_interval`` steps.
``clip_fraction`` and ``kl_approx`` then describe ``r``, the ratio the trust region acts on, whereas
``ESS``, ``max_ratio`` and ``mean_ratio`` describe the behavior ratio ``pi_theta / pi_behav`` the data
is actually weighted by.
>>> loss_module = ClipPPOLoss(actor, critic, delay_actor=True, max_importance_ratio=100.0)
>>> updater = SoftUpdate(loss_module, eps=0.889)
>>> for batch in replay_buffer:
... losses = loss_module(batch)
... loss = losses["loss_objective"] + losses["loss_critic"] + losses["loss_entropy"]
... loss.backward()
... optimizer.step()
... optimizer.zero_grad()
... updater.step()
"""
_schedulable_buffers = frozenset(
{"clip_epsilon", "clip_epsilon_low", "clip_epsilon_high"}
)
actor_network: TensorDictModule
critic_network: TensorDictModule
actor_network_params: TensorDictParams
critic_network_params: TensorDictParams
target_actor_network_params: TensorDictParams
target_critic_network_params: TensorDictParams
def __init__(
self,
actor_network: ProbabilisticTensorDictSequential | None = None,
critic_network: TensorDictModule | None = None,
*,
clip_epsilon: float | tuple[float, float] = 0.2,
entropy_bonus: bool = True,
samples_mc_entropy: int = 1,
entropy_coeff: float | Mapping[NestedKey, float] | None = None,
critic_coeff: float | None = None,
loss_critic_type: str = "smooth_l1",
normalize_advantage: bool = False,
normalize_advantage_exclude_dims: tuple[int] = (),
advantage_norm: ValueNorm | None = None,
gamma: float | None = None,
separate_losses: bool = False,
reduction: str | None = None,
clip_value: bool | float | None = None,
delay_actor: bool = False,
max_importance_ratio: float | None = None,
device: torch.device | None = None,
**kwargs,
):
# Define clipping of the value loss
if isinstance(clip_value, bool):
if clip_value and isinstance(clip_epsilon, (tuple, list)):
raise ValueError(
"clip_value=True cannot infer the value-clip threshold from "
"an asymmetric clip_epsilon tuple. Pass an explicit float "
"threshold instead, e.g. clip_value=0.2."
)
clip_value = clip_epsilon if clip_value else None
super().__init__(
actor_network,
critic_network,
entropy_bonus=entropy_bonus,
samples_mc_entropy=samples_mc_entropy,
entropy_coeff=entropy_coeff,
critic_coeff=critic_coeff,
loss_critic_type=loss_critic_type,
normalize_advantage=normalize_advantage,
normalize_advantage_exclude_dims=normalize_advantage_exclude_dims,
advantage_norm=advantage_norm,
gamma=gamma,
separate_losses=separate_losses,
reduction=reduction,
clip_value=clip_value,
delay_actor=delay_actor,
max_importance_ratio=max_importance_ratio,
device=device,
**kwargs,
)
if device is None:
try:
device = next(self.parameters()).device
except (AttributeError, StopIteration):
device = getattr(
torch, "get_default_device", lambda: torch.device("cpu")
)()
self._asymmetric_clip = isinstance(clip_epsilon, (tuple, list))
if self._asymmetric_clip:
# Asymmetric (DAPO Clip-Higher) thresholds, as in GRPOLoss
eps_low, eps_high = _validate_clip_epsilon(clip_epsilon)
self.register_buffer(
"clip_epsilon_low", torch.tensor(eps_low, device=device)
)
self.register_buffer(
"clip_epsilon_high", torch.tensor(eps_high, device=device)
)
else:
# the symmetric float keeps the legacy (schedulable) buffer and
# is deliberately not validated, for backward compatibility
self.register_buffer(
"clip_epsilon", torch.tensor(clip_epsilon, device=device)
)
_CLIP_EPSILON_BUFFER_NAMES = (
"clip_epsilon",
"clip_epsilon_low",
"clip_epsilon_high",
)
def __setattr__(self, name: str, value) -> None:
# Scheduling the wrong clip-threshold flavor (e.g. annealing
# `loss.clip_epsilon` on a loss built with an asymmetric tuple) would
# silently create a shadow attribute while clipping keeps using the
# registered buffers - fail loudly instead.
if (
name in self._CLIP_EPSILON_BUFFER_NAMES
and isinstance(value, (int, float))
and getattr(self, "_buffers", None) is not None
and name not in self._buffers
and any(n in self._buffers for n in self._CLIP_EPSILON_BUFFER_NAMES)
):
if self._asymmetric_clip:
raise AttributeError(
f"Cannot schedule {name!r}: this {type(self).__name__} was built "
"with an asymmetric clip_epsilon tuple; schedule "
"'clip_epsilon_low' / 'clip_epsilon_high' instead."
)
raise AttributeError(
f"Cannot schedule {name!r}: this {type(self).__name__} was built "
"with a symmetric float clip_epsilon; schedule 'clip_epsilon' instead."
)
super().__setattr__(name, value)
@property
def _clip_bounds(self):
if self._asymmetric_clip:
return (
(-self.clip_epsilon_low).log1p(),
self.clip_epsilon_high.log1p(),
)
return (
(-self.clip_epsilon).log1p(),
self.clip_epsilon.log1p(),
)
@property
def out_keys(self):
if self._out_keys is None:
keys = ["loss_objective", "clip_fraction"]
if self.entropy_bonus:
keys.extend(["entropy", "loss_entropy"])
if self.loss_critic:
keys.append("loss_critic")
if self.clip_value:
keys.append("value_clip_fraction")
keys.extend(["ESS", "kl_approx", "max_ratio", "mean_ratio"])
self._out_keys = keys
return self._out_keys
@out_keys.setter
def out_keys(self, values):
self._out_keys = values
[docs]
@dispatch
def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
tensordict = tensordict.clone(False)
advantage = tensordict.get(
self.tensor_keys.advantage, None, as_padded_tensor=True
)
update_norm = advantage is None
if advantage is None:
if self.critic_network is None:
raise RuntimeError(
"Critic network is not specified, cannot compute advantage within forward."
)
self.value_estimator(
tensordict,
params=self._cached_critic_network_params_detached,
target_params=self.target_critic_network_params,
)
advantage = tensordict.get(self.tensor_keys.advantage)
advantage = self._maybe_normalize_advantage(
advantage, tensordict, update_norm=update_norm
)
log_weight, dist, kl_approx, log_is_weight = self._log_weight(
tensordict, adv_shape=advantage.shape[:-1]
)
advantage = _broadcast_advantage_to_log_weight(advantage, log_weight)
_check_advantage_broadcast(advantage, log_weight)
# log_weight is the ratio the clipping acts on (pi_theta / pi_prox); the
# surrogate itself is weighted by the behavior ratio
# pi_theta / pi_behav = exp(log_weight + log_is_weight). Both coincide
# when the proximal policy is the behavior policy (log_is_weight=None).
log_behav_weight = (
log_weight if log_is_weight is None else log_weight + log_is_weight
)
# ESS for logging
with torch.no_grad():
# In theory, ESS should be computed on particles sampled from the same source. Here we sample according
# to different, unrelated trajectories, which is not standard. Still, it can give an idea of the weights'
# dispersion.
lw = log_behav_weight.squeeze(-1)
ess = (2 * lw.logsumexp(0) - (2 * lw).logsumexp(0)).exp()
batch = log_weight.shape[0]
gain1 = log_behav_weight.exp() * advantage
log_weight_clip = log_weight.clamp(*self._clip_bounds)
clip_fraction = (log_weight_clip != log_weight).to(log_weight.dtype).mean()
if log_is_weight is not None:
log_weight_clip = log_weight_clip + log_is_weight
ratio = log_weight_clip.exp()
gain2 = ratio * advantage
gain = torch.stack([gain1, gain2], -1).min(dim=-1).values
td_out = TensorDict({"loss_objective": -gain})
td_out.set("clip_fraction", clip_fraction)
td_out.set("kl_approx", kl_approx.detach().mean()) # for logging
if self.entropy_bonus:
entropy = self._get_entropy(dist, adv_shape=advantage.shape[:-1])
if is_tensor_collection(entropy):
# Reports the entropy of each action head.
td_out.set("composite_entropy", entropy.detach())
td_out.set(
"entropy", _sum_td_features(entropy).detach().mean()
) # for logging
else:
td_out.set("entropy", entropy.detach().mean()) # for logging
td_out.set("loss_entropy", self._weighted_loss_entropy(entropy))
if self._has_critic:
loss_critic, value_clip_fraction, explained_variance = self.loss_critic(
tensordict
)
td_out.set("loss_critic", loss_critic)
if value_clip_fraction is not None:
td_out.set("value_clip_fraction", value_clip_fraction)
if explained_variance is not None:
td_out.set("explained_variance", explained_variance)
td_out.set("ESS", ess / batch)
with torch.no_grad():
ratio = log_behav_weight.exp()
td_out.set("max_ratio", ratio.max())
td_out.set("mean_ratio", ratio.mean())
td_out = td_out.named_apply(
lambda name, value: self._reduce_loss(value, tensordict).squeeze(-1)
if name.startswith("loss_")
else value,
)
self._clear_weakrefs(
tensordict,
td_out,
"actor_network_params",
"critic_network_params",
"target_actor_network_params",
"target_critic_network_params",
)
return td_out
[docs]
class KLPENPPOLoss(PPOLoss):
"""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.
Args:
actor_network (ProbabilisticTensorDictSequential): policy operator.
critic_network (ValueOperator): value operator.
Keyword Args:
dtarg (scalar, optional): target KL divergence. Defaults to ``0.01``.
samples_mc_kl (int, optional): number of samples used to compute the KL divergence
if no analytical formula can be found. Defaults to ``1``.
beta (scalar, optional): initial KL divergence multiplier.
Defaults to ``1.0``.
decrement (scalar, optional): how much beta should be decremented if KL < dtarg. Valid range: decrement <= 1.0
default: ``0.5``.
increment (scalar, optional): how much beta should be incremented if KL > dtarg. Valid range: increment >= 1.0
default: ``2.0``.
entropy_bonus (bool, optional): if ``True``, an entropy bonus will be added to the
loss to favour exploratory policies. Defaults to ``True``.
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_entropy`` will control how many
samples will be used to compute this estimate.
Defaults to ``1``.
entropy_coeff: scalar | Mapping[NestedKey, scalar], optional): entropy multiplier when computing the total loss.
* **Scalar**: one value applied to the summed entropy of every action head.
* **Mapping** ``{head_name: coeff}`` gives an individual coefficient for each action-head's entropy.
Defaults to ``0.01``.
See :ref:`ppo_entropy_coefficients` for detailed usage examples and troubleshooting.
critic_coeff (scalar, optional): critic loss multiplier when computing the total
loss. Defaults to ``1.0``.
loss_critic_type (str, optional): loss function for the value discrepancy.
Can be one of "l1", "l2" or "smooth_l1". Defaults to ``"smooth_l1"``.
normalize_advantage (bool, optional): if ``True``, the advantage will be normalized
before being used. Defaults to ``False``.
normalize_advantage_exclude_dims (Tuple[int], optional): dimensions to exclude from the advantage
standardization. Negative dimensions are valid. This is useful in multiagent (or multiobjective) settings
where the agent (or objective) dimension may be excluded from the reductions. Default: ().
advantage_norm (ValueNorm, optional): a stateful
:class:`~torchrl.modules.ValueNorm` used to rescale the advantage,
e.g. :class:`~torchrl.modules.PercentileValueNorm` for
DreamerV3-style return normalization. The advantage is divided by
``advantage_norm.scale()`` without re-centering (the advantage is
already centred by the value baseline). In training mode, the
statistics are updated with the fresh value targets when the loss
runs its own value estimator (i.e. when the input tensordict
carries no advantage entry), so repeated forwards over the same
precomputed rollout do not skew the moving average; rows marked
invalid by the validity masks (see
:attr:`~torchrl.objectives.LossModule.loss_mask_key`) are excluded
from the update. When the advantage is precomputed outside the
loss, call ``advantage_norm.update(value_target)`` once per rollout
yourself. Mutually exclusive with ``normalize_advantage``.
Defaults to ``None``.
separate_losses (bool, optional): if ``True``, shared parameters between
policy and critic will only be trained on the policy loss.
Defaults to ``False``, i.e., 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 ``"advantage"``.
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_target"``.
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 ``"state_value"``.
functional (bool, optional): whether modules should be functionalized.
Functionalizing permits features like meta-RL, but makes it
impossible to use distributed models (DDP, FSDP, ...) and comes
with a little cost. Defaults to ``True``.
reduction (str, optional): Specifies the reduction to apply to the output:
``"none"`` | ``"mean"`` | ``"sum"``. ``"none"``: no reduction will be applied,
``"mean"``: the sum of the output will be divided by the number of
elements in the output, ``"sum"``: the output will be summed. Default: ``"mean"``.
clip_value (:obj:`float`, optional): If provided, it will be used to compute a clipped version of the value
prediction with respect to the input tensordict value estimate and use it to calculate the value loss.
The purpose of clipping is to limit the impact of extreme value predictions, helping stabilize training
and preventing large updates. However, it will have no impact if the value estimate was done by the current
version of the value estimator. Defaults to ``None``.
delay_actor (bool, optional): if ``True``, a detached copy of the actor parameters is kept under
``target_actor_network_params`` and used as the *proximal policy* of the objective, decoupled from
the *behavior policy* that collected the data (whose log-probabilities are read from the
``sample_log_prob`` entry). The surrogate stays weighted by the behavior ratio ``pi_theta / pi_behav``
(so the gradient estimate remains unbiased for the data at hand) while the proximal policy is used
by the trust-region term of the subclasses (clipping in :class:`~torchrl.objectives.ClipPPOLoss`,
KL penalty in :class:`~torchrl.objectives.KLPENPPOLoss`) and by the ``kl_approx`` diagnostic.
Updating the target parameters with :class:`~torchrl.objectives.SoftUpdate` after every optimizer
step makes the proximal policy an exponentially-weighted moving average of the policy, i.e. PPO-EWMA
("Batch size-invariance for policy optimization", Hilton et al., 2021,
https://arxiv.org/abs/2110.00641). Requires ``functional=True``. Defaults to ``False``.
max_importance_ratio (:obj:`float`, optional): if provided, the behavior ratio ``pi_theta / pi_behav`` is
capped at this value. Stale data, or a proximal policy that drifted away from the behavior policy,
can otherwise produce arbitrarily large ratios. The cap lifts the (gradient-free) behavior
log-probability rather than clamping the ratio, so capped samples keep a rescaled policy gradient,
as in the reference PPO-EWMA implementation (which uses ``100.0``). Defaults to ``None`` (no cap).
device (torch.device, optional): device of the buffers. Defaults to ``None``.
.. note:: Parameters and buffers from the policy / critic will not be cast to that device to ensure that
the storages match the ones that are passed to other components, such as data collectors.
.. note:
The advantage (typically GAE) can be computed by the loss function or
in the training loop. The latter option is usually preferred, but this is
up to the user to choose which option is to be preferred.
If the advantage key (``"advantage`` by default) is not present in the
input tensordict, the advantage will be computed by the :meth:`~.forward`
method.
>>> ppo_loss = KLPENPPOLoss(actor, critic)
>>> advantage = GAE(critic)
>>> data = next(datacollector)
>>> losses = ppo_loss(data)
>>> # equivalent
>>> advantage(data)
>>> losses = ppo_loss(data)
A custom advantage module can be built using :meth:`~.make_value_estimator`.
The default is :class:`~torchrl.objectives.value.GAE` with hyperparameters
dictated by :func:`~torchrl.objectives.utils.default_value_kwargs`.
>>> ppo_loss = KLPENPPOLoss(actor, critic)
>>> ppo_loss.make_value_estimator(ValueEstimators.TDLambda)
>>> data = next(datacollector)
>>> losses = ppo_loss(data)
.. note::
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 = ActorValueOperator(common, actor_head, value_head)
>>> loss_module = KLPENPPOLoss(model.get_policy_operator(), model.get_value_operator())
>>> # second option, with a single call to the common module
>>> loss_module = KLPENPPOLoss(ProbabilisticTensorDictSequential(model, actor_head), value_head)
This will work regardless of whether separate_losses is activated or not.
"""
_schedulable_buffers = frozenset({"beta"})
actor_network: TensorDictModule
critic_network: TensorDictModule
actor_network_params: TensorDictParams
critic_network_params: TensorDictParams
target_actor_network_params: TensorDictParams
target_critic_network_params: TensorDictParams
def __init__(
self,
actor_network: ProbabilisticTensorDictSequential | None = None,
critic_network: TensorDictModule | None = None,
*,
dtarg: float = 0.01,
beta: float = 1.0,
increment: float = 2,
decrement: float = 0.5,
samples_mc_kl: int = 1,
entropy_bonus: bool = True,
samples_mc_entropy: int = 1,
entropy_coeff: float | Mapping[NestedKey, float] | None = None,
critic_coeff: float | None = None,
loss_critic_type: str = "smooth_l1",
normalize_advantage: bool = False,
normalize_advantage_exclude_dims: tuple[int] = (),
advantage_norm: ValueNorm | None = None,
gamma: float | None = None,
separate_losses: bool = False,
reduction: str | None = None,
clip_value: float | None = None,
delay_actor: bool = False,
max_importance_ratio: float | None = None,
device: torch.device | None = None,
**kwargs,
):
super().__init__(
actor_network,
critic_network,
entropy_bonus=entropy_bonus,
samples_mc_entropy=samples_mc_entropy,
entropy_coeff=entropy_coeff,
critic_coeff=critic_coeff,
loss_critic_type=loss_critic_type,
normalize_advantage=normalize_advantage,
normalize_advantage_exclude_dims=normalize_advantage_exclude_dims,
advantage_norm=advantage_norm,
gamma=gamma,
separate_losses=separate_losses,
reduction=reduction,
clip_value=clip_value,
delay_actor=delay_actor,
max_importance_ratio=max_importance_ratio,
device=device,
**kwargs,
)
if device is None:
try:
device = next(self.parameters()).device
except (AttributeError, StopIteration):
device = getattr(
torch, "get_default_device", lambda: torch.device("cpu")
)()
self.dtarg = dtarg
self._beta_init = beta
self.register_buffer("beta", torch.tensor(beta, device=device))
if increment < 1.0:
raise ValueError(
f"increment should be >= 1.0 in KLPENPPOLoss, got {increment:4.4f}"
)
self.increment = increment
if decrement > 1.0:
raise ValueError(
f"decrement should be <= 1.0 in KLPENPPOLoss, got {decrement:4.4f}"
)
self.decrement = decrement
self.samples_mc_kl = samples_mc_kl
def _set_in_keys(self):
keys = []
_maybe_add_or_extend_key(keys, self.actor_network.in_keys)
_maybe_add_or_extend_key(keys, self.actor_network.in_keys, "next")
if self.critic_network is not None:
_maybe_add_or_extend_key(keys, self.critic_network.in_keys)
_maybe_add_or_extend_key(keys, self.tensor_keys.action)
_maybe_add_or_extend_key(keys, self.tensor_keys.sample_log_prob)
_maybe_add_or_extend_key(keys, self.tensor_keys.reward, "next")
_maybe_add_or_extend_key(keys, self.tensor_keys.done, "next")
_maybe_add_or_extend_key(keys, self.tensor_keys.terminated, "next")
# Get the parameter keys from the actor dist
actor_dist_module = None
for module in self.actor_network.modules():
# Ideally we should combine them if there is more than one
if isinstance(module, ProbabilisticTensorDictModule):
if actor_dist_module is not None:
raise RuntimeError(
"Actors with one and only one distribution are currently supported "
f"in {type(self).__name__}. If you need to use more than one "
f"distributions over the action space please submit an issue "
f"on github."
)
actor_dist_module = module
if actor_dist_module is None:
raise RuntimeError("Could not find the probabilistic module in the actor.")
keys += list(actor_dist_module.in_keys)
self._in_keys = list(set(keys))
@property
def out_keys(self):
if self._out_keys is None:
keys = ["loss_objective", "kl"]
if self.entropy_bonus:
keys.extend(["entropy", "loss_entropy"])
if self.loss_critic:
keys.append("loss_critic")
if self.clip_value:
keys.append("value_clip_fraction")
self._out_keys = keys
return self._out_keys
@out_keys.setter
def out_keys(self, values):
self._out_keys = values
[docs]
@dispatch
def forward(self, tensordict: TensorDictBase) -> TensorDict:
tensordict_copy = tensordict.copy()
if self.delay_actor:
# the KL penalty pulls the policy towards the proximal policy held
# in the target parameters, not towards the behavior policy whose
# distribution parameters are stored in the data
with torch.no_grad(), self.target_actor_network_params.to_module(
self.actor_network, preserve_module_state=False
):
previous_dist = self.actor_network.get_dist(tensordict_copy)
else:
try:
previous_dist = self.actor_network.build_dist_from_params(tensordict)
except KeyError as err:
raise KeyError(
"The parameters of the distribution were not found. "
f"Make sure they are provided to {type(self).__name__}."
) from err
advantage = tensordict_copy.get(self.tensor_keys.advantage, None)
update_norm = advantage is None
if advantage is None:
self.value_estimator(
tensordict_copy,
params=self._cached_critic_network_params_detached,
target_params=self.target_critic_network_params,
)
advantage = tensordict_copy.get(self.tensor_keys.advantage)
advantage = self._maybe_normalize_advantage(
advantage, tensordict_copy, update_norm=update_norm
)
log_weight, dist, kl_approx, log_is_weight = self._log_weight(
tensordict_copy, adv_shape=advantage.shape[:-1]
)
advantage = _broadcast_advantage_to_log_weight(advantage, log_weight)
_check_advantage_broadcast(advantage, log_weight)
if log_is_weight is not None:
# the surrogate is weighted by the behavior ratio pi_theta / pi_behav
log_weight = log_weight + log_is_weight
neg_loss = log_weight.exp() * advantage
with (
self.actor_network_params.to_module(
self.actor_network, preserve_module_state=False
)
if self.functional
else contextlib.nullcontext()
):
current_dist = self.actor_network.get_dist(tensordict_copy)
is_composite = isinstance(current_dist, CompositeDistribution)
if has_analytic_kl(previous_dist, current_dist):
kl = torch.distributions.kl.kl_divergence(previous_dist, current_dist)
else:
_warn_mc_kl(previous_dist, current_dist)
with (
set_composite_lp_aggregate(False)
if is_composite
else contextlib.nullcontext()
):
x, previous_log_prob = sample_and_log_prob(
previous_dist, (self.samples_mc_kl,)
)
current_log_prob = current_dist.log_prob(x)
if is_tensor_collection(previous_log_prob):
if previous_log_prob.batch_size != advantage.shape[:-1]:
previous_log_prob.batch_size = (
self.samples_mc_kl,
) + advantage.shape[:-1]
current_log_prob.batch_size = (
self.samples_mc_kl,
) + advantage.shape[:-1]
previous_log_prob = _sum_td_features(previous_log_prob)
# Both dists have presumably the same params
current_log_prob = _sum_td_features(current_log_prob)
kl = (previous_log_prob - current_log_prob).mean(0)
kl = kl.unsqueeze(-1)
neg_loss = neg_loss - self.beta * kl
if kl.mean() > self.dtarg * 1.5:
self.beta.data *= self.increment
elif kl.mean() < self.dtarg / 1.5:
self.beta.data *= self.decrement
td_out = TensorDict(
{
"loss_objective": -neg_loss,
"kl": kl.detach(),
"kl_approx": kl_approx.detach().mean(),
},
)
if self.entropy_bonus:
entropy = self._get_entropy(dist, adv_shape=advantage.shape[:-1])
if is_tensor_collection(entropy):
# Reports the entropy of each action head.
td_out.set("composite_entropy", entropy.detach())
td_out.set(
"entropy", _sum_td_features(entropy).detach().mean()
) # for logging
else:
td_out.set("entropy", entropy.detach().mean()) # for logging
td_out.set("loss_entropy", self._weighted_loss_entropy(entropy))
if self._has_critic:
loss_critic, value_clip_fraction, explained_variance = self.loss_critic(
tensordict_copy
)
td_out.set("loss_critic", loss_critic)
if value_clip_fraction is not None:
td_out.set("value_clip_fraction", value_clip_fraction)
if explained_variance is not None:
td_out.set("explained_variance", explained_variance)
td_out = td_out.named_apply(
lambda name, value: self._reduce_loss(value, tensordict_copy).squeeze(-1)
if name.startswith("loss_")
else value,
)
self._clear_weakrefs(
tensordict,
td_out,
"actor_network_params",
"critic_network_params",
"target_actor_network_params",
"target_critic_network_params",
)
return td_out
def reset(self) -> None:
self.beta = self._beta_init
