# 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 abc
import functools
import warnings
from collections.abc import Iterator, Mapping
from copy import deepcopy
from dataclasses import dataclass
from typing import Any, Literal
import torch
from tensordict import is_tensor_collection, NestedKey, TensorDict, TensorDictBase
from tensordict.nn import TensorDictModule, TensorDictModuleBase, TensorDictParams
from tensordict.utils import Buffer
from torch import nn
from torch.nn import Parameter
from torchrl._utils import rl_warnings
from torchrl.envs.utils import ExplorationType, set_exploration_type
from torchrl.modules.tensordict_module.rnn import set_recurrent_mode
from torchrl.objectives.utils import _reduce, default_value_kwargs, ValueEstimators
from torchrl.objectives.value import ValueEstimatorBase
try:
from torch.compiler import is_compiling
except ImportError:
from torch._dynamo import is_compiling
#: Input entries that :attr:`LossModule.loss_mask_key` ``= "auto"`` looks for,
#: in order, ANDing every one it finds. Both are written by TorchRL itself:
#: ``("collector", "mask")`` by :class:`~torchrl.data.SliceSampler` with
#: ``pad_output=True`` (``False`` marks duplicated padding steps), and
#: ``"shifted_valid"`` by the value estimators (``False`` marks positions whose
#: bootstrapped target crosses an episode boundary).
AUTO_LOSS_MASK_KEYS: tuple[NestedKey, ...] = (("collector", "mask"), "shifted_valid")
def _updater_check_forward_prehook(module, *args, **kwargs):
if (
not all(module._has_update_associated.values())
and rl_warnings()
and not is_compiling()
):
warnings.warn(
module.TARGET_NET_WARNING,
category=UserWarning,
)
def _forward_wrapper(func):
@functools.wraps(func)
def new_forward(self, *args, **kwargs):
em = set_exploration_type(self.deterministic_sampling_mode)
em.__enter__()
rm = set_recurrent_mode(True)
rm.__enter__()
try:
return func(self, *args, **kwargs)
finally:
em.__exit__(None, None, None)
rm.__exit__(None, None, None)
return new_forward
class _LossMeta(abc.ABCMeta):
def __init__(cls, name, bases, attr_dict):
super().__init__(name, bases, attr_dict)
cls.forward = _forward_wrapper(cls.forward)
for name, value in cls.__dict__.items():
if not name.startswith("_") and name.endswith("loss"):
setattr(cls, name, _forward_wrapper(value))
# Merge _schedulable_buffers from all bases so __setattr__ can do a
# single O(1) check instead of walking the MRO on every call.
merged = set()
for base in cls.__mro__:
merged |= getattr(base, "_schedulable_buffers", frozenset())
cls._all_schedulable_buffers = frozenset(merged)
[docs]
class LossModule(TensorDictModuleBase, metaclass=_LossMeta):
"""A parent class for RL losses.
LossModule inherits from nn.Module. 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.
:cvar default_value_estimator: The default value type of the class.
Losses that require a value estimation are equipped with a default value
pointer. This class attribute indicates which value estimator will be
used if none other is specified.
The value estimator can be changed using the :meth:`~.make_value_estimator` method.
By default, the forward method is always decorated with a
gh :class:`torchrl.envs.ExplorationType.MEAN`
To utilize the ability configuring the tensordict keys via
:meth:`~.set_keys()` a subclass must define an _AcceptedKeys dataclass.
This dataclass should include all keys that are intended to be configurable.
The default :meth:`~._forward_value_estimator_keys()` implementation forwards
common value-estimator keys when present. Subclasses should override it when
the loss's key names need to be remapped before being forwarded to the
underlying value estimator.
Subclasses can declare a ``_schedulable_buffers`` frozenset to allow direct
scalar assignment (e.g. ``loss.entropy_coeff = 0.003``) for registered
buffers that are commonly scheduled during training. The assignment performs
an in-place update, preserving the buffer's device and dtype.
Padded or otherwise invalid positions are excluded from the reduction
through :attr:`loss_mask_key`. It defaults to ``"auto"``, which discovers
the validity masks TorchRL itself writes (``("collector", "mask")`` from
:class:`~torchrl.data.SliceSampler` with ``pad_output=True``, and
``"shifted_valid"`` from the value estimators); set it to a
:class:`~tensordict.NestedKey` to name a single mask entry, or to ``None``
to reduce over every position:
>>> loss.loss_mask_key = ("my_masks", "valid") # use this entry only
>>> loss.loss_mask_key = None # no masking at all
Examples:
>>> class MyLoss(LossModule):
>>> @dataclass
>>> class _AcceptedKeys:
>>> action = "action"
>>>
>>> def _forward_value_estimator_keys(self, **kwargs) -> None:
>>> pass
>>>
>>> loss = MyLoss()
>>> loss.set_keys(action="action2")
.. note:: When a policy that is wrapped or augmented with an exploration module is passed
to the loss, we want to deactivate the exploration through ``set_exploration_type(<exploration>)`` where
``<exploration>`` is either ``ExplorationType.MEAN``, ``ExplorationType.MODE`` or
``ExplorationType.DETERMINISTIC``. The default value is ``DETERMINISTIC`` and it is set
through the ``deterministic_sampling_mode`` loss attribute. If another
exploration mode is required (or if ``DETERMINISTIC`` is not available), one can
change the value of this attribute which will change the mode.
"""
_schedulable_buffers: frozenset = frozenset()
_loss_mask_key: NestedKey | Literal["auto"] | None = "auto"
@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.
"""
tensor_keys: _AcceptedKeys
_vmap_randomness = None
default_value_estimator: ValueEstimators = None
use_prioritized_weights: str | bool = "auto"
deterministic_sampling_mode: ExplorationType = ExplorationType.DETERMINISTIC
SEP = "."
TARGET_NET_WARNING = (
"No target network updater has been associated "
"with this loss module, but target parameters have been found. "
"While this is supported, it is expected that the target network "
"updates will be manually performed. You can deactivate this warning "
"by turning the RL_WARNINGS env variable to False."
)
@property
def tensor_keys(self) -> _AcceptedKeys:
return self._tensor_keys
def __new__(cls, *args, **kwargs):
self = super().__new__(cls)
return self
def __setattr__(self, name: str, value) -> None:
# Allow direct scalar assignment to schedulable buffers:
# loss.entropy_coeff = 0.003
# performs an in-place copy, preserving device and dtype.
if (
isinstance(value, (int, float))
and name in type(self)._all_schedulable_buffers
and hasattr(self, "_buffers")
and name in self._buffers
and self._buffers[name] is not None
):
self._buffers[name].copy_(
torch.as_tensor(
value,
dtype=self._buffers[name].dtype,
device=self._buffers[name].device,
)
)
return
super().__setattr__(name, value)
def __init__(self):
super().__init__()
self._cache = {}
self._param_maps = {}
self._value_estimator = None
self._has_update_associated = {}
self.value_type = self.default_value_estimator
self._tensor_keys = self._AcceptedKeys()
self.register_forward_pre_hook(_updater_check_forward_prehook)
@property
def functional(self):
"""Whether the module is functional.
Unless it has been specifically designed not to be functional, all losses are functional.
"""
return True
[docs]
def get_stateful_net(self, network_name: str, copy: bool | None = None):
"""Returns a stateful version of the network.
This can be used to initialize parameters.
Such networks will often not be callable out-of-the-box and will require a `vmap` call
to be executable.
Args:
network_name (str): the network name to gather.
copy (bool, optional): if ``True``, a deepcopy of the network is made.
Defaults to ``True``.
.. note:: if the module is not functional, no copy is made.
"""
net = getattr(self, network_name)
if not self.functional:
if copy is not None and copy:
raise RuntimeError("Cannot copy module in non-functional mode.")
return net
copy = True if copy is None else copy
if copy:
net = deepcopy(net)
params = getattr(self, network_name + "_params")
params.to_module(net, preserve_module_state=False)
return net
[docs]
def from_stateful_net(self, network_name: str, stateful_net: nn.Module):
"""Populates the parameters of a model given a stateful version of the network.
See :meth:`~.get_stateful_net` for details on how to gather a stateful version of the network.
Args:
network_name (str): the network name to reset.
stateful_net (nn.Module): the stateful network from which the params should be
gathered.
"""
if not self.functional:
getattr(self, network_name).load_state_dict(stateful_net.state_dict())
return
params = TensorDict.from_module(stateful_net, as_module=True)
keyset0 = set(params.keys(True, True))
self_params = getattr(self, network_name + "_params")
keyset1 = set(self_params.keys(True, True))
if keyset0 != keyset1:
raise RuntimeError(
f"The keys of params and provided module differ: "
f"{keyset1 - keyset0} are in self.params and not in the module, "
f"{keyset0 - keyset1} are in the module but not in self.params."
)
self_params.data.update_(params.data)
def _set_deprecated_ctor_keys(self, **kwargs) -> None:
for key, value in kwargs.items():
if value is not None:
raise RuntimeError(
f"Setting '{key}' via the constructor is deprecated, use .set_keys(<key>='some_key') instead.",
)
@property
def loss_mask_key(self) -> NestedKey | Literal["auto"] | None:
"""Which input entry marks the positions that contribute to the loss.
``"auto"`` (the default) discovers the validity masks TorchRL writes
itself -- see :data:`AUTO_LOSS_MASK_KEYS`. A
:class:`~tensordict.NestedKey` restricts masking to that single entry;
``None`` disables it. To use an entry literally named ``"auto"``, pass
the one-element tuple ``("auto",)``.
"""
return self._loss_mask_key
@loss_mask_key.setter
def loss_mask_key(self, value: NestedKey | Literal["auto"] | None) -> None:
if value is None or (isinstance(value, str) and value == "auto"):
self._loss_mask_key = value
return
error = ValueError(
f"loss_mask_key must be 'auto', None or a NestedKey, got {value!r}."
)
if not isinstance(value, NestedKey):
raise error
self._loss_mask_key = value
def _loss_mask_keys(self) -> tuple[NestedKey, ...]:
"""The input entries :meth:`_reduce_loss` will look for, in order."""
key = self._loss_mask_key
if key is None:
return ()
if key == "auto":
return AUTO_LOSS_MASK_KEYS
return (key,)
@staticmethod
def _expand_loss_mask(mask: torch.Tensor, loss: torch.Tensor) -> torch.Tensor:
# Validity masks conventionally carry a trailing singleton dimension, to
# broadcast against the [..., 1]-shaped rewards they accompany --
# ``("collector", "mask")`` is [B, T, 1] where a per-timestep loss is
# [B, T]. Drop those before broadcasting the other way.
while mask.ndim > loss.ndim and mask.shape[-1] == 1:
mask = mask.squeeze(-1)
if mask.ndim > loss.ndim:
raise ValueError(
f"A mask of shape {tuple(mask.shape)} cannot be applied to a "
f"loss of shape {tuple(loss.shape)}: the mask has more "
"non-singleton dimensions than the loss. Per-element masking "
"requires an elementwise loss, not one whose event dimensions "
"have already been reduced (e.g. a log-prob)."
)
if mask.ndim < loss.ndim:
mask = mask.reshape(mask.shape + (1,) * (loss.ndim - mask.ndim))
return mask.expand_as(loss)
def _reduce_loss(
self,
loss: torch.Tensor,
tensordict: TensorDictBase | None = None,
*,
mask: torch.Tensor | None = None,
reduction: str | None = None,
weights: torch.Tensor | None = None,
) -> torch.Tensor:
if reduction is None:
reduction = self.reduction
# A caller-supplied mask keeps the legacy ``reduction="none"`` contract
# (masked positions are dropped, the output is compacted); masks read
# from the input instead preserve the loss shape, so that per-position
# outputs stay aligned with the input batch.
caller_mask = mask is not None
if mask is not None:
mask = self._expand_loss_mask(mask, loss)
if tensordict is not None:
for mask_key in self._loss_mask_keys():
tensordict_mask = tensordict.get(mask_key, default=None)
if tensordict_mask is not None:
tensordict_mask = self._expand_loss_mask(tensordict_mask, loss)
mask = tensordict_mask if mask is None else mask & tensordict_mask
if mask is not None:
# Select rather than multiply: masked positions may hold non-finite
# values, and ``nan * 0`` is ``nan`` in both the forward and the
# backward pass.
loss = torch.where(mask, loss, torch.zeros_like(loss))
if weights is not None and weights.shape != loss.shape:
weights = self._expand_loss_mask(weights, loss)
if reduction == "none" and not caller_mask:
return loss if weights is None else loss * weights
if weights is None and reduction == "mean":
return loss.sum() / mask.sum().clamp_min(1)
if weights is None and reduction == "sum":
return loss.sum()
return _reduce(loss, reduction=reduction, mask=mask, weights=weights)
[docs]
def set_keys(self, **kwargs) -> None:
"""Set tensordict key names.
Examples:
>>> from torchrl.objectives import DQNLoss
>>> # initialize the DQN loss
>>> actor = torch.nn.Linear(3, 4)
>>> dqn_loss = DQNLoss(actor, action_space="one-hot")
>>> dqn_loss.set_keys(priority_key="td_error", action_value_key="action_value")
"""
for key, value in kwargs.items():
if key not in self._AcceptedKeys.__dataclass_fields__:
raise ValueError(
f"{key} is not an accepted tensordict key. Accepted keys are: {self._AcceptedKeys.__dataclass_fields__}."
)
if value is not None:
setattr(self.tensor_keys, key, value)
else:
setattr(self.tensor_keys, key, getattr(self.default_keys(), key))
try:
self._forward_value_estimator_keys(**kwargs)
except AttributeError as err:
raise AttributeError(
"To utilize `.set_keys(...)` for tensordict key configuration, the subclassed loss module "
"must define an _AcceptedKeys dataclass containing all keys intended for configuration. "
"If the default `._forward_value_estimator_keys()` implementation is insufficient, the "
"subclass must override it to forward modified tensordict keys to the underlying value_estimator."
) from err
[docs]
def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
"""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.
Args:
tensordict: an input tensordict with the values required to compute the loss.
Returns:
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.
"""
raise NotImplementedError
[docs]
def convert_to_functional(
self,
module: TensorDictModule,
module_name: str,
expand_dim: int | None = None,
create_target_params: bool = False,
compare_against: list[Parameter] | None = None,
**kwargs,
) -> None:
"""Converts a module to functional to be used in the loss.
Args:
module (TensorDictModule or compatible): a stateful tensordict module.
Parameters from this module will be isolated in the `<module_name>_params`
attribute and a stateless version of the module will be registered
under the `module_name` attribute.
module_name (str): name where the module will be found.
The parameters of the module will be found under ``loss_module.<module_name>_params``
whereas the module will be found under ``loss_module.<module_name>``.
expand_dim (int, optional): if provided, the parameters of the module
will be expanded ``N`` times, where ``N = expand_dim`` along the
first dimension. This option is to be used whenever a target
network with more than one configuration is to be used.
.. note::
If a ``compare_against`` list of values is provided, the
resulting parameters will simply be a detached expansion
of the original parameters. If ``compare_against`` is not
provided, the value of the parameters will be resampled uniformly
between the minimum and maximum value of the parameter content.
create_target_params (bool, optional): if ``True``, a detached
copy of the parameter will be available to feed a target network
under the name ``loss_module.<module_name>_target_params``.
If ``False`` (default), this attribute will still be available
but it will be a detached instance of the parameters, not a copy.
In other words, any modification of the parameter value
will directly be reflected in the target parameters.
compare_against (iterable of parameters, optional): if provided,
this list of parameters will be used as a comparison set for
the parameters of the module. If the parameters are expanded
(``expand_dim > 0``), the resulting parameters for the module
will be a simple expansion of the original parameter. Otherwise,
the resulting parameters will be a detached version of the
original parameters. If ``None``, the resulting parameters
will carry gradients as expected.
"""
# Walk the MRO so subclasses don't have to redeclare annotations
# introduced by their parents — ``cls.__annotations__`` is *not*
# inherited automatically in Python.
inherited_annotations: set[str] = set()
for base in type(self).__mro__:
inherited_annotations.update(getattr(base, "__annotations__", {}).keys())
for name in (
module_name,
module_name + "_params",
"target_" + module_name + "_params",
):
if name not in inherited_annotations:
warnings.warn(
f"The name {name} wasn't part of the annotations ({sorted(inherited_annotations)}). Make sure it is present in the definition class."
)
if kwargs:
raise TypeError(f"Unrecognised keyword arguments {list(kwargs.keys())}")
# To make it robust to device casting, we must register list of
# tensors as lazy calls to `getattr(self, name_of_tensor)`.
# Otherwise, casting the module to a device will keep old references
# to uncast tensors
sep = self.SEP
if isinstance(module, (list, tuple)):
if len(module) != expand_dim:
raise RuntimeError(
"The ``expand_dim`` value must match the length of the module list/tuple "
"if a single module isn't provided."
)
params = TensorDict.from_modules(
*module, as_module=True, expand_identical=True
)
# Use the first module as the functional forward reference.
module = module[0]
else:
params = TensorDict.from_module(module, as_module=True)
for key in params.keys(True):
if sep in key:
raise KeyError(
f"The key {key} contains the '_sep_' pattern which is prohibited. Consider renaming the parameter / buffer."
)
if compare_against is not None:
compare_against = set(compare_against)
else:
compare_against = set()
if expand_dim:
# Expands the dims of params and buffers.
# If the param already exist in the module, we return a simple expansion of the
# original one. Otherwise, we expand and resample it.
# For buffers, a cloned expansion (or equivalently a repeat) is returned.
def _compare_and_expand(param):
if is_tensor_collection(param):
return param._apply_nest(
_compare_and_expand,
batch_size=[expand_dim, *param.shape],
filter_empty=False,
call_on_nested=True,
)
if not isinstance(param, nn.Parameter):
buffer = param.expand(expand_dim, *param.shape).clone()
return buffer
if param in compare_against:
expanded_param = param.data.expand(expand_dim, *param.shape)
# the expanded parameter must be sent to device when to()
# is called:
return expanded_param
else:
p_out = param.expand(expand_dim, *param.shape).clone()
p_out = nn.Parameter(
p_out.uniform_(
p_out.data.min().item(), p_out.data.max().item()
).requires_grad_()
)
return p_out
params = TensorDictParams(
params.apply(
_compare_and_expand,
batch_size=[expand_dim, *params.shape],
filter_empty=False,
call_on_nested=True,
),
no_convert=True,
)
param_name = module_name + "_params"
prev_set_params = set(self.parameters())
# register parameters and buffers
for key, parameter in list(params.items(True, True)):
if parameter not in prev_set_params:
pass
elif compare_against is not None and parameter in compare_against:
params.set(key, parameter.data)
setattr(self, param_name, params)
# Set the module in the __dict__ directly to avoid listing its params
# A deepcopy with meta device could be used but that assumes that the model is copyable!
self.__dict__[module_name] = module
name_params_target = "target_" + module_name
if create_target_params:
# if create_target_params:
# we create a TensorDictParams to keep the target params as Buffer instances
target_params = TensorDictParams(
params.apply(
_make_target_param(clone=create_target_params), filter_empty=False
),
no_convert=True,
)
target_params.uninitialized_keys = {
key
for key, value in params.items(True, True)
if torch.nn.parameter.is_lazy(value)
}
setattr(self, name_params_target + "_params", target_params)
self._has_update_associated[module_name] = not create_target_params
def _clear_weakrefs(self, *tds):
if is_compiling():
# Waiting for weakrefs reconstruct to be supported by compile
for td in tds:
if isinstance(td, str):
td = getattr(self, td, None)
if not is_tensor_collection(td):
continue
td.clear_refs_for_compile_()
def __getattr__(self, item):
if item.startswith("target_") and item.endswith("_params"):
params = self._modules.get(item, None)
if params is not None:
pending_keys = getattr(params, "uninitialized_keys", None)
if pending_keys:
source_params = getattr(self, item[7:])
for key in pending_keys.copy():
source = source_params.get(key)
if not torch.nn.parameter.is_lazy(source):
params.get(key).data = source.data.clone()
pending_keys.remove(key)
if params is None:
# no target param, take detached data
params = getattr(self, item[7:])
params = params.data
elif (
not self._has_update_associated[item[7:-7]]
and rl_warnings()
and not is_compiling()
):
# no updater associated
warnings.warn(
self.TARGET_NET_WARNING,
category=UserWarning,
)
return params
return super().__getattr__(item)
def _apply(self, fn):
# any call to apply erases the cache: the reason is that detached
# params will fail to be cast so we need to get the cache back
self._erase_cache()
return super()._apply(fn)
def _erase_cache(self):
for key in list(self.__dict__):
if key.startswith("_cache"):
delattr(self, key)
def _networks(self) -> Iterator[nn.Module]:
for item in self.__dir__():
if isinstance(item, nn.Module):
yield item
[docs]
def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
for _, param in self.named_parameters(recurse=recurse):
yield param
[docs]
def named_parameters(
self, prefix: str = "", recurse: bool = True
) -> Iterator[tuple[str, Parameter]]:
for name, param in super().named_parameters(prefix=prefix, recurse=recurse):
if not name.startswith("_target"):
yield name, param
def reset(self) -> None:
# mainly used for PPO with KL target
pass
def _maybe_get_priority_weight(
self, tensordict: TensorDictBase
) -> torch.Tensor | None:
"""Extract priority weights from tensordict if prioritized replay is enabled.
Args:
tensordict (TensorDictBase): The input tensordict that may contain priority weights.
Returns:
torch.Tensor | None: The priority weights if available and enabled, None otherwise.
"""
weights = None
if (
self.use_prioritized_weights in (True, "auto")
and self.tensor_keys.priority_weight in tensordict.keys()
):
weights = tensordict.get(self.tensor_keys.priority_weight)
return weights
def _reset_module_parameters(self, module_name, module):
params_name = f"{module_name}_params"
target_name = f"target_{module_name}_params"
params = self._modules.get(params_name, None)
target = self._modules.get(target_name, None)
if params is not None:
with params.to_module(module, preserve_module_state=False):
module.reset_parameters_recursive()
else:
module.reset_parameters_recursive()
if target is not None:
with target.to_module(module, preserve_module_state=False):
module.reset_parameters_recursive()
[docs]
def reset_parameters_recursive(
self,
):
"""Reset the parameters of the module."""
for key, item in self.__dict__.items():
if isinstance(item, nn.Module):
self._reset_module_parameters(key, item)
@property
def value_estimator(self) -> ValueEstimatorBase:
"""The value function blends in the reward and value estimate(s) from upcoming state(s)/state-action pair(s) into a target value estimate for the value network."""
out = self._value_estimator
if out is None:
self._default_value_estimator()
return self._value_estimator
return out
@value_estimator.setter
def value_estimator(self, value):
self._value_estimator = value
[docs]
def load_state_dict(self, state_dict: Mapping[str, Any], *args, **kwargs):
"""Load a state dict, building the default value estimator first when the state carries one.
The value estimator is created on first use, so a loss saved after its
first forward pass holds ``_value_estimator.*`` entries that a freshly
constructed loss has no target for. Losses configured with a
non-default estimator must call :meth:`make_value_estimator` before
loading, as they would before training.
"""
if self._value_estimator is None and any(
str(key).startswith("_value_estimator.") for key in state_dict
):
self._default_value_estimator()
return super().load_state_dict(state_dict, *args, **kwargs)
def _default_value_estimator(self):
"""A value-function constructor when none is provided.
No kwarg should be present as default parameters should be retrieved
from :obj:`torchrl.objectives.utils.DEFAULT_VALUE_FUN_PARAMS`.
"""
self.make_value_estimator(
self.default_value_estimator, device=self._default_device
)
@property
def _default_device(self) -> torch.device | None:
"""A util to find the default device.
Returns ``None`` if parameters are spread across multiple devices.
"""
devices = set()
for p in self.parameters():
devices.add(p.device)
if len(devices) == 1:
return list(devices)[0]
return None
def _prepare_value_estimator_kwargs(self, value_type, **hyperparams):
"""Common preamble for make_value_estimator overrides.
Handles three boilerplate steps that every subclass repeats:
defaulting ``value_type``, delegating the instance/subclass path to
the base-class handler, and building the ``hp`` dict from
:func:`~torchrl.objectives.utils.default_value_kwargs` merged with
any caller-supplied overrides.
Returns:
``(resolved_value_type, hp)`` — the caller should continue with
its own if/elif value-type dispatch.
``(None, None)`` — the call was fully handled (instance or
subclass path); the caller should ``return self``.
"""
if value_type is None:
value_type = self.default_value_estimator
if isinstance(value_type, ValueEstimatorBase) or (
isinstance(value_type, type) and issubclass(value_type, ValueEstimatorBase)
):
LossModule.make_value_estimator(self, value_type, **hyperparams)
return None, None
self.value_type = value_type
hp = dict(default_value_kwargs(value_type))
if hasattr(self, "gamma"):
hp["gamma"] = self.gamma
hp.update(hyperparams)
return value_type, hp
[docs]
def register_coeff_buffer(
self,
name: str,
value: float | int | torch.Tensor | None,
*,
device: torch.device | str | None = None,
dtype: torch.dtype | None = None,
) -> None:
"""Register a scalar coefficient as a buffer, converting it to a tensor.
Eliminates the recurring ``if not isinstance(value, Tensor): value =
torch.tensor(value); self.register_buffer(name, value)`` boilerplate in
loss ``__init__`` methods.
If ``value`` is ``None`` the attribute is set to ``None`` instead of a
buffer being registered, matching the common optional-coefficient idiom
(e.g. ``critic_coeff`` / ``clip_value``).
Args:
name (str): the buffer / attribute name.
value (float, int, Tensor or None): the coefficient. ``None`` sets
the attribute to ``None``.
device (torch.device, optional): device for the buffer.
dtype (torch.dtype, optional): dtype for the buffer.
"""
if value is None:
setattr(self, name, None)
return
if isinstance(value, bool) or not isinstance(value, (float, int, torch.Tensor)):
raise ValueError(f"{name} must be a float or a scalar tensor, got {value}.")
value = torch.as_tensor(value, device=device, dtype=dtype)
if value.numel() != 1:
raise ValueError(f"{name} must be a float or a scalar tensor, got {value}.")
self.register_buffer(name, value)
# Value-estimator keys forwarded by the default
# :meth:`_forward_value_estimator_keys`. These six are accepted by every
# built-in value estimator. Keys that only some estimators accept (e.g.
# ``sample_log_prob``) are intentionally excluded; losses that forward them
# should override :meth:`_forward_value_estimator_keys`.
_value_estimator_default_keys = (
"advantage",
"value_target",
"value",
"reward",
"done",
"terminated",
)
def _forward_value_estimator_keys(self, **kwargs) -> None:
"""Default forwarding of tensordict keys to the value estimator.
Forwards every key in :attr:`_value_estimator_default_keys` that is
present on this loss's ``tensor_keys`` to the underlying value
estimator, then refreshes the loss input keys via ``_set_in_keys`` when
that method exists.
Losses whose value-estimator key *names* differ from their own
``tensor_keys`` names -- e.g. mapping the estimator's ``value`` to a
``state_action_value`` / ``global_value`` key -- or that forward
estimator-specific keys such as ``sample_log_prob`` must override this
method.
"""
value_estimator = getattr(self, "_value_estimator", None)
if value_estimator is not None:
keys = {
name: getattr(self.tensor_keys, name)
for name in self._value_estimator_default_keys
if hasattr(self.tensor_keys, name)
}
if keys:
value_estimator.set_keys(**keys)
set_in_keys = getattr(self, "_set_in_keys", None)
if callable(set_in_keys):
set_in_keys()
[docs]
def make_value_estimator(self, value_type: ValueEstimators = None, **hyperparams):
"""Value-function constructor.
If the non-default value function is wanted, it must be built using
this method.
Args:
value_type (ValueEstimators, ValueEstimatorBase, or type): The value
estimator to use. This can be one of the following:
- A :class:`~torchrl.objectives.utils.ValueEstimators` enum type
indicating which value function to use. If none is provided,
the default stored in the ``default_value_estimator``
attribute will be used.
- A :class:`~torchrl.objectives.value.ValueEstimatorBase` instance,
which will be used directly as the value estimator.
- A :class:`~torchrl.objectives.value.ValueEstimatorBase` subclass,
which will be instantiated with the provided ``hyperparams``.
The resulting value estimator class will be registered in
``self.value_type``, allowing future refinements.
**hyperparams: hyperparameters to use for the value function.
If not provided, the value indicated by
:func:`~torchrl.objectives.utils.default_value_kwargs` will be
used. When passing a ``ValueEstimatorBase`` subclass, these
hyperparameters are passed directly to the class constructor.
Returns:
self: Returns the loss module for method chaining.
Examples:
>>> from torchrl.objectives import DQNLoss
>>> # initialize the DQN loss
>>> actor = torch.nn.Linear(3, 4)
>>> dqn_loss = DQNLoss(actor, action_space="one-hot")
>>> # updating the parameters of the default value estimator
>>> dqn_loss.make_value_estimator(gamma=0.9)
>>> dqn_loss.make_value_estimator(
... ValueEstimators.TD1,
... gamma=0.9)
>>> # if we want to change the gamma value
>>> dqn_loss.make_value_estimator(dqn_loss.value_type, gamma=0.9)
Using a :class:`~torchrl.objectives.value.ValueEstimatorBase` subclass:
>>> from torchrl.objectives.value import TD0Estimator
>>> dqn_loss.make_value_estimator(TD0Estimator, gamma=0.99, value_network=value_net)
Using a :class:`~torchrl.objectives.value.ValueEstimatorBase` instance:
>>> from torchrl.objectives.value import GAE
>>> gae = GAE(gamma=0.99, lmbda=0.95, value_network=value_net)
>>> ppo_loss.make_value_estimator(gae)
"""
if value_type is None:
value_type = self.default_value_estimator
if isinstance(value_type, ValueEstimatorBase):
self._value_estimator = value_type
self.value_type = type(value_type)
return self
if isinstance(value_type, type) and issubclass(value_type, ValueEstimatorBase):
if "device" not in hyperparams:
device = self._default_device
if device is not None:
hyperparams["device"] = device
self._value_estimator = value_type(**hyperparams)
self.value_type = value_type
return self
self.value_type = value_type
if value_type == ValueEstimators.TD1:
raise NotImplementedError(
f"Value type {value_type} it not implemented for loss {type(self)}."
)
elif value_type == ValueEstimators.TD0:
raise NotImplementedError(
f"Value type {value_type} it not implemented for loss {type(self)}."
)
elif value_type == ValueEstimators.GAE:
raise NotImplementedError(
f"Value type {value_type} it not implemented for loss {type(self)}."
)
elif value_type == ValueEstimators.VTrace:
raise NotImplementedError(
f"Value type {value_type} it not implemented for loss {type(self)}."
)
elif value_type == ValueEstimators.TDLambda:
raise NotImplementedError(
f"Value type {value_type} it not implemented for loss {type(self)}."
)
else:
raise NotImplementedError(f"Unknown value type {value_type}")
return self
@property
def vmap_randomness(self):
"""Vmap random mode.
The vmap randomness mode controls what :func:`~torch.vmap` should do when dealing with
functions with a random outcome such as :func:`~torch.randn` and :func:`~torch.rand`.
If `"error"`, any random function will raise an exception indicating that `vmap` does not
know how to handle the random call.
If `"different"`, every element of the batch along which vmap is being called will
behave differently. If `"same"`, vmaps will copy the same result across all elements.
``vmap_randomness`` defaults to `"error"` if no random module is detected, and to `"different"` in
other cases. By default, only a limited number of modules are listed as random, but the list can be extended
using the :func:`~torchrl.objectives.common.add_random_module` function.
This property supports setting its value.
"""
if self._vmap_randomness is None:
import torchrl.objectives.utils
main_modules = list(self.__dict__.values()) + list(self.children())
modules = (
module
for main_module in main_modules
if isinstance(main_module, nn.Module)
for module in main_module.modules()
)
for val in modules:
if isinstance(val, torchrl.objectives.utils.RANDOM_MODULE_LIST):
self._vmap_randomness = "different"
break
else:
self._vmap_randomness = "error"
return self._vmap_randomness
def set_vmap_randomness(self, value):
if value not in ("error", "same", "different"):
raise ValueError(
"Wrong vmap randomness, should be one of 'error', 'same' or 'different'."
)
self._vmap_randomness = value
self._make_vmap()
@staticmethod
def _make_meta_params(param):
is_param = isinstance(param, nn.Parameter)
pd = param.detach().to("meta")
if is_param:
pd = nn.Parameter(pd, requires_grad=False)
return pd
def _make_vmap(self):
"""Caches thevmap callers to reduce the overhead at runtime."""
raise NotImplementedError(
f"_make_vmap has been called but is not implemented for loss of type {type(self).__name__}."
)
class _make_target_param:
def __init__(self, clone):
self.clone = clone
def __call__(self, x):
x = x.data.clone() if self.clone else x.data
if isinstance(x, nn.Parameter):
return Buffer(x)
return x
[docs]
def add_random_module(module):
"""Adds a random module to the list of modules that will be detected by :meth:`~torchrl.objectives.LossModule.vmap_randomness` as random."""
import torchrl.objectives.utils
torchrl.objectives.utils.RANDOM_MODULE_LIST = (
torchrl.objectives.utils.RANDOM_MODULE_LIST + (module,)
)
