# Copyright (c) Meta Plobs_dictnc. 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 functools
import math
from collections.abc import Sequence
from contextlib import nullcontext
from copy import copy
from enum import IntEnum
from textwrap import indent
from typing import Any, Literal, TYPE_CHECKING
import torch
from tensordict import TensorDict, TensorDictBase
from tensordict.nn import TensorDictModuleBase
from tensordict.utils import _zip_strict, expand_as_right, NestedKey, unravel_key
from torch import nn
from torchrl._utils import _replace_last
from torchrl.data.tensor_specs import (
Bounded,
Categorical,
Composite,
ContinuousBox,
MultiCategorical,
MultiOneHot,
OneHot,
TensorSpec,
Unbounded,
)
if TYPE_CHECKING:
from torchrl.data.vla import RobotDatasetMetadata
from torchrl.envs.common import EnvBase
if TYPE_CHECKING:
from typing import Self
else:
Self = Any
from torchrl.data.vla.schema import (
ACTION_CHUNK_KEY,
ACTION_IS_PAD_KEY,
ACTION_KEY,
ACTION_TOKENS_KEY,
)
from torchrl.data.vla.tokenizers import ActionTokenizerBase
from torchrl.envs.transforms._base import (
Compose,
FORWARD_NOT_IMPLEMENTED,
Transform,
TransformedEnv,
)
from torchrl.envs.transforms._observation import CatFrames, UnsqueezeTransform
from torchrl.envs.transforms.utils import _get_reset
from torchrl.envs.utils import ExplorationType, set_exploration_type
from torchrl.modules.tensordict_module.controllers import LowLevelController
from torchrl.modules.utils import get_env_transforms_from_module
__all__ = [
"ActionChunkTransform",
"ActionDiscretizer",
"ActionMask",
"ActionScaling",
"ActionTokenizerTransform",
"DiscreteActionProjection",
"FlattenAction",
"LastAction",
"MultiAction",
"ClosedLoopMultiAction",
]
[docs]
class DiscreteActionProjection(Transform):
"""Projects discrete actions from a high dimensional space to a low dimensional space.
Given a discrete action (from 1 to N) encoded as a one-hot vector and a
maximum action index num_actions (with num_actions < N), transforms the action such that
action_out is at most num_actions.
If the input action is > num_actions, it is being replaced by a random value
between 0 and num_actions-1. Otherwise the same action is kept.
This is intended to be used with policies applied over multiple discrete
control environments with different action space.
A call to DiscreteActionProjection.forward (eg from a replay buffer or in a
sequence of nn.Modules) will call the transform num_actions_effective -> max_actions
on the :obj:`"in_keys"`, whereas a call to _call will be ignored. Indeed,
transformed envs are instructed to update the input keys only for the inner
base_env, but the original input keys will remain unchanged.
Args:
num_actions_effective (int): max number of action considered.
max_actions (int): maximum number of actions that this module can read.
action_key (NestedKey, optional): key name of the action. Defaults to "action".
include_forward (bool, optional): if ``True``, a call to forward will also
map the action from one domain to the other when the module is called
by a replay buffer or an nn.Module chain. Defaults to `True`.
Examples:
>>> torch.manual_seed(0)
>>> N = 3
>>> M = 2
>>> action = torch.zeros(N, dtype=torch.long)
>>> action[-1] = 1
>>> td = TensorDict({"action": action}, [])
>>> transform = DiscreteActionProjection(num_actions_effective=M, max_actions=N)
>>> _ = transform.inv(td)
>>> print(td.get("action"))
tensor([1])
"""
def __init__(
self,
num_actions_effective: int,
max_actions: int,
action_key: NestedKey = "action",
include_forward: bool = True,
):
in_keys_inv = [action_key]
if include_forward:
in_keys = in_keys_inv
else:
in_keys = []
if in_keys_inv is None:
in_keys_inv = []
super().__init__(
in_keys=in_keys,
out_keys=copy(in_keys),
in_keys_inv=in_keys_inv,
out_keys_inv=copy(in_keys_inv),
)
self.num_actions_effective = num_actions_effective
self.max_actions = max_actions
if max_actions < num_actions_effective:
raise RuntimeError(
"The `max_actions` int must be greater or equal to `num_actions_effective`."
)
def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase:
# We don't do anything here because the action is modified by the inv
# method but we don't need to map it back as it won't be updated in the original
# tensordict
return next_tensordict
def _apply_transform(self, action: torch.Tensor) -> torch.Tensor:
# We still need to code the forward transform for replay buffers and models
action = action.argmax(-1) # bool to int
action = nn.functional.one_hot(action, self.max_actions)
return action
def _inv_apply_transform(self, action: torch.Tensor) -> torch.Tensor:
if action.shape[-1] != self.max_actions:
raise RuntimeError(
f"action.shape[-1]={action.shape[-1]} must match self.max_actions={self.max_actions}."
)
action = action.long().argmax(-1) # bool to int
idx = action >= self.num_actions_effective
if idx.any():
action[idx] = torch.randint(self.num_actions_effective, (idx.sum(),))
action = nn.functional.one_hot(action, self.num_actions_effective)
return action
def __repr__(self) -> str:
return (
f"{self.__class__.__name__}(num_actions_effective={self.num_actions_effective}, max_actions={self.max_actions}, "
f"in_keys_inv={self.in_keys_inv})"
)
[docs]
class ActionMask(Transform):
"""An adaptive action masker.
This transform is useful to ensure that randomly generated actions
respect legal actions, by masking the action specs.
It reads the mask from the input tensordict after the step is executed,
and adapts the mask of the finite action spec.
.. note:: This transform will fail when used without an environment.
.. note:: **MultiDiscrete action spaces with 2D masks (e.g., board games)**
When wrapping a Gym environment with a ``MultiDiscrete`` action space
(e.g., ``MultiDiscrete([5, 5])``) and an ``action_mask`` observation whose
shape matches the ``nvec`` (e.g., shape ``(5, 5)``), the :class:`~torchrl.envs.GymWrapper`
automatically converts the action space to a flattened ``Categorical(n=25)``
or ``OneHot(n=25)``. This allows the mask to represent all possible action
combinations (25 in this example) rather than independent sub-actions.
This is particularly useful for grid-based games where the mask indicates
which (row, column) positions are valid moves.
Args:
action_key (NestedKey, optional): the key where the action tensor can be found.
Defaults to ``"action"``.
mask_key (NestedKey, optional): the key where the action mask can be found.
Defaults to ``"action_mask"``.
Examples:
>>> import torch
>>> from torchrl.data.tensor_specs import Categorical, Binary, Unbounded, Composite
>>> from torchrl.envs.transforms import ActionMask, TransformedEnv
>>> from torchrl.envs.common import EnvBase
>>> class MaskedEnv(EnvBase):
... def __init__(self, *args, **kwargs):
... super().__init__(*args, **kwargs)
... self.action_spec = Categorical(4)
... self.state_spec = Composite(action_mask=Binary(4, dtype=torch.bool))
... self.observation_spec = Composite(obs=Unbounded(3))
... self.reward_spec = Unbounded(1)
...
... def _reset(self, tensordict=None):
... td = self.observation_spec.rand()
... td.update(torch.ones_like(self.state_spec.rand()))
... return td
...
... def _step(self, data):
... td = self.observation_spec.rand()
... mask = data.get("action_mask")
... action = data.get("action")
... mask = mask.scatter(-1, action.unsqueeze(-1), 0)
...
... td.set("action_mask", mask)
... td.set("reward", self.reward_spec.rand())
... td.set("done", ~mask.any().view(1))
... return td
...
... def _set_seed(self, seed) -> None:
... pass
...
>>> torch.manual_seed(0)
>>> base_env = MaskedEnv()
>>> env = TransformedEnv(base_env, ActionMask())
>>> r = env.rollout(10)
>>> r["action_mask"]
tensor([[ True, True, True, True],
[ True, True, False, True],
[ True, True, False, False],
[ True, False, False, False]])
"""
ACCEPTED_SPECS = (
OneHot,
Categorical,
MultiOneHot,
MultiCategorical,
)
SPEC_TYPE_ERROR = "The action spec must be one of {}. Got {} instead."
def __init__(
self, action_key: NestedKey = "action", mask_key: NestedKey = "action_mask"
):
if not isinstance(action_key, (tuple, str)):
raise ValueError(
f"The action key must be a nested key. Got {type(action_key)} instead."
)
if not isinstance(mask_key, (tuple, str)):
raise ValueError(
f"The mask key must be a nested key. Got {type(mask_key)} instead."
)
super().__init__(
in_keys=[action_key, mask_key], out_keys=[], in_keys_inv=[], out_keys_inv=[]
)
[docs]
def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
raise RuntimeError(FORWARD_NOT_IMPLEMENTED.format(type(self)))
@property
def action_spec(self) -> TensorSpec:
action_spec = self.container.full_action_spec[self.in_keys[0]]
if not isinstance(action_spec, self.ACCEPTED_SPECS):
raise ValueError(
self.SPEC_TYPE_ERROR.format(self.ACCEPTED_SPECS, type(action_spec))
)
return action_spec
def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase:
if self.parent is None:
raise RuntimeError(
f"{type(self)}.parent cannot be None: make sure this transform is executed within an environment."
)
mask = next_tensordict.get(self.in_keys[1])
self.action_spec.update_mask(mask.to(self.action_spec.device))
return next_tensordict
def _reset(
self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
) -> TensorDictBase:
return self._call(tensordict_reset)
[docs]
class ActionDiscretizer(Transform):
"""A transform to discretize a continuous action space.
This transform makes it possible to use an algorithm designed for discrete
action spaces such as DQN over environments with a continuous action space.
Args:
num_intervals (int or torch.Tensor): the number of discrete values
for each element of the action space. If a single integer is provided,
all action items are sliced with the same number of elements.
If a tensor is provided, it must have the same number of elements
as the action space (ie, the length of the ``num_intervals`` tensor
must match the last dimension of the action space).
action_key (NestedKey, optional): the action key to use. Points to
the action of the parent env (the floating point action).
Defaults to ``"action"``.
out_action_key (NestedKey, optional): the key where the discrete
action should be written. If ``None`` is provided, it defaults to
the value of ``action_key``. If both keys do not match, the
continuous action_spec is moved from the ``full_action_spec``
environment attribute to the ``full_state_spec`` container,
as only the discrete action should be sampled for an action to
be taken. Providing ``out_action_key`` can ensure that the
floating point action is available to be recorded.
sampling (ActionDiscretizer.SamplingStrategy, optinoal): an element
of the ``ActionDiscretizer.SamplingStrategy`` ``IntEnum`` object
(``MEDIAN``, ``LOW``, ``HIGH`` or ``RANDOM``). Indicates how the
continuous action should be sampled in the provided interval.
categorical (bool, optional): if ``False``, one-hot encoding is used.
Defaults to ``True``.
Examples:
>>> from torchrl.envs import GymEnv, check_env_specs
>>> import torch
>>> base_env = GymEnv("HalfCheetah-v4")
>>> num_intervals = torch.arange(5, 11)
>>> categorical = True
>>> sampling = ActionDiscretizer.SamplingStrategy.MEDIAN
>>> t = ActionDiscretizer(
... num_intervals=num_intervals,
... categorical=categorical,
... sampling=sampling,
... out_action_key="action_disc",
... )
>>> env = base_env.append_transform(t)
TransformedEnv(
env=GymEnv(env=HalfCheetah-v4, batch_size=torch.Size([]), device=cpu),
transform=ActionDiscretizer(
num_intervals=tensor([ 5, 6, 7, 8, 9, 10]),
action_key=action,
out_action_key=action_disc,,
sampling=0,
categorical=True))
>>> check_env_specs(env)
>>> # Produce a rollout
>>> r = env.rollout(4)
>>> print(r)
TensorDict(
fields={
action: Tensor(shape=torch.Size([4, 6]), device=cpu, dtype=torch.float32, is_shared=False),
action_disc: Tensor(shape=torch.Size([4, 6]), device=cpu, dtype=torch.int64, is_shared=False),
done: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.bool, is_shared=False),
next: TensorDict(
fields={
done: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.bool, is_shared=False),
observation: Tensor(shape=torch.Size([4, 17]), device=cpu, dtype=torch.float64, is_shared=False),
reward: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.float32, is_shared=False),
terminated: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.bool, is_shared=False),
truncated: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
batch_size=torch.Size([4]),
device=cpu,
is_shared=False),
observation: Tensor(shape=torch.Size([4, 17]), device=cpu, dtype=torch.float64, is_shared=False),
terminated: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.bool, is_shared=False),
truncated: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
batch_size=torch.Size([4]),
device=cpu,
is_shared=False)
>>> assert r["action"].dtype == torch.float
>>> assert r["action_disc"].dtype == torch.int64
>>> assert (r["action"] < base_env.action_spec.high).all()
>>> assert (r["action"] > base_env.action_spec.low).all()
.. note:: Custom Sampling Strategies
To implement a custom sampling strategy beyond the built-in options
(``MEDIAN``, ``LOW``, ``HIGH``, ``RANDOM``), subclass ``ActionDiscretizer``
and override the :meth:`~ActionDiscretizer.custom_arange` method. This
method computes the normalized interval positions (values in ``[0, 1)``)
that determine where each discrete action maps within the continuous
action interval.
Example:
>>> class LogSpacedActionDiscretizer(ActionDiscretizer):
... def custom_arange(self, nint, device):
... # Use logarithmic spacing instead of linear
... return torch.logspace(-2, 0, nint, device=device) - 0.01
.. seealso:: :class:`~torchrl.envs.transforms.ActionTokenizerTransform` -- a
bidirectional action <-> token codec built around an explicit
:class:`~torchrl.data.vla.ActionTokenizerBase`. Prefer
``ActionDiscretizer`` when the binning should be derived from the env's
bounded ``action_spec`` (with configurable in-bin sampling); prefer
``ActionTokenizerTransform`` when the binning is owned by a tokenizer
that must be shared between offline encoding (replay buffer) and online
decoding (env), e.g. for an autoregressive token VLA policy.
"""
[docs]
class SamplingStrategy(IntEnum):
"""The sampling strategies for ActionDiscretizer."""
MEDIAN = 0
LOW = 1
HIGH = 2
RANDOM = 3
def __init__(
self,
num_intervals: int | torch.Tensor,
action_key: NestedKey = "action",
out_action_key: NestedKey = None,
sampling=None,
categorical: bool = True,
):
if out_action_key is None:
out_action_key = action_key
super().__init__(in_keys_inv=[action_key], out_keys_inv=[out_action_key])
self.action_key = action_key
self.out_action_key = out_action_key
if not isinstance(num_intervals, torch.Tensor):
self.num_intervals = num_intervals
else:
self.register_buffer("num_intervals", num_intervals)
if sampling is None:
sampling = self.SamplingStrategy.MEDIAN
self.sampling = sampling
self.categorical = categorical
def __repr__(self) -> str:
def _indent(s):
return indent(s, 4 * " ")
num_intervals = f"num_intervals={self.num_intervals}"
action_key = f"action_key={self.action_key}"
out_action_key = f"out_action_key={self.out_action_key}"
sampling = f"sampling={self.sampling}"
categorical = f"categorical={self.categorical}"
return (
f"{type(self).__name__}(\n{_indent(num_intervals)},\n{_indent(action_key)},"
f"\n{_indent(out_action_key)},\n{_indent(sampling)},\n{_indent(categorical)})"
)
[docs]
def custom_arange(self, nint, device):
"""Compute the normalized interval positions for discretization.
This method generates values in the range [0, 1) that determine where
each discrete action maps within the continuous action interval.
Override this method in a subclass to implement custom sampling
strategies beyond the built-in ``MEDIAN``, ``LOW``, ``HIGH``, and
``RANDOM`` strategies.
Args:
nint (int): the number of intervals (discrete actions) for this
action dimension.
device (torch.device): the device on which to create the tensor.
Returns:
torch.Tensor: a 1D tensor of shape ``(nint,)`` with values in
``[0, 1)`` representing the normalized positions within each
interval.
Example:
>>> class CustomActionDiscretizer(ActionDiscretizer):
... def custom_arange(self, nint, device):
... # Custom sampling: use logarithmic spacing
... return torch.logspace(-2, 0, nint, device=device) - 0.01
"""
result = torch.arange(
start=0.0,
end=1.0,
step=1 / nint,
dtype=self.dtype,
device=device,
)
result_ = result
if self.sampling in (
self.SamplingStrategy.HIGH,
self.SamplingStrategy.MEDIAN,
):
result_ = (1 - result).flip(0)
if self.sampling == self.SamplingStrategy.MEDIAN:
result = (result + result_) / 2
else:
result = result_
return result
def _init(self):
# We just need to access the action spec for everything to be initialized
try:
_ = self.container.full_action_spec
except AttributeError:
raise RuntimeError(
f"Cannot execute transform {type(self).__name__} without a parent env."
)
[docs]
def inv(self, tensordict):
if self.out_keys_inv[0] == self.in_keys_inv[0]:
return super().inv(tensordict)
# We re-write this because we don't want to clone the TD here
return self._inv_call(tensordict)
def _inv_call(self, tensordict):
# action is categorical, map it to desired dtype
intervals = getattr(self, "intervals", None)
if intervals is None:
self._init()
return self._inv_call(tensordict)
action = tensordict.get(self.out_keys_inv[0])
if self.categorical:
action = action.unsqueeze(-1)
if isinstance(intervals, torch.Tensor):
shape = action.shape[: -intervals.ndim]
intervals = intervals.expand(shape + intervals.shape)
action = intervals.gather(index=action, dim=-1).squeeze(-1)
else:
action = torch.stack(
[
interval.gather(index=action, dim=-1).squeeze(-1)
for interval, action in zip(intervals, action.unbind(-2))
],
-1,
)
else:
nvec = self.nvec
empty_shape = not nvec.ndim
if not empty_shape:
nvec = nvec.tolist()
if isinstance(intervals, torch.Tensor):
shape = action.shape[: (-intervals.ndim + 1)]
intervals = intervals.expand(shape + intervals.shape)
intervals = intervals.unbind(-2)
action = action.split(nvec, dim=-1)
action = torch.stack(
[
intervals[action].view(action.shape[:-1])
for (intervals, action) in zip(intervals, action)
],
-1,
)
else:
shape = action.shape[: -intervals.ndim]
intervals = intervals.expand(shape + intervals.shape)
action = intervals[action].squeeze(-1)
if self.sampling == self.SamplingStrategy.RANDOM:
action = action + self.jitters * torch.rand_like(self.jitters)
return tensordict.set(self.in_keys_inv[0], action)
[docs]
class MultiAction(Transform):
"""A transform to execute multiple actions in the parent environment.
This transform unbinds the actions along a specific dimension and passes each action independently.
The returned transform can be either a stack of the observations gathered during the steps or only the
last observation (and similarly for the rewards, see args below).
By default, the actions must be stacked along the first dimension after the root tensordict batch-dims, i.e.
>>> td = policy(td)
>>> actions = td.select(*env.action_keys)
>>> # Adapt the batch-size
>>> actions = actions.auto_batch_size_(td.ndim + 1)
>>> # Step-wise actions
>>> actions = actions.unbind(-1)
If a `"done"` entry is encountered, the next steps are skipped for the env that has reached that state.
.. note:: If a transform is appended before the MultiAction, it will be called multiple times. If it is appended
after, it will be called once per macro-step.
.. note:: Extra entries written by the policy alongside the actions (e.g. the action tokens and
log-probabilities of a token-head policy) are left untouched on the root tensordict and therefore
ride along on the outer (macro-step) transition: each outer step of a rollout carries the policy
outputs of the chunk decided at that step.
.. note:: When a done state fires inside the chunk (with ``stack_rewards=True``), the reward stack of
that outer step holds the executed steps' rewards followed by a single zero-filled slot for the
skipped remainder of the chunk. Its length therefore differs from a full chunk's, and stacking
such outer steps in a rollout yields a lazy stack with ragged reward entries. If the per-chunk
reward is computed from the outer transition anyway (e.g. with
:class:`~torchrl.envs.transforms.SuccessReward` appended after this transform), pass
``stack_rewards=False`` to keep the outer transition dense and uniform.
.. note:: Skipping the remaining steps after a done state relies on the ``"_step"`` partial-step
entry. Single (unbatched) environments and batched environments
(:class:`~torchrl.envs.SerialEnv` / :class:`~torchrl.envs.ParallelEnv`) handle it natively; for a
batch-locked vectorized environment, the base environment's ``_step`` is trusted to honor the
mask itself (see :meth:`~torchrl.envs.EnvBase.step`) and environments that ignore it will keep
stepping every sub-environment until the end of the chunk.
Keyword Args:
dim (int, optional): the stack dimension with respect to the tensordict ``ndim`` attribute.
Must be greater than 0. Defaults to ``1`` (the first dimension after the batch-dims).
stack_rewards (bool, optional): if ``True``, each step's reward will be stack in the output tensordict.
If ``False``, only the last reward will be returned. The reward spec is adapted accordingly. The
stack dimension is the same as the action stack dimension. Defaults to ``True``.
stack_observations (bool, optional): if ``True``, each step's observation will be stack in the output tensordict.
If ``False``, only the last observation will be returned. The observation spec is adapted accordingly. The
stack dimension is the same as the action stack dimension. Defaults to ``False``.
action_key (NestedKey, optional): the one-step action key consumed by
the base environment. Defaults to the parent environment action key.
chunk_key (NestedKey, optional): the policy-facing key that holds the
stacked actions. Defaults to ``action_key`` for backward
compatibility. Set this to values such as
``("vla_action", "chunk")`` when a chunk policy should act through
:class:`MultiAction` without re-keying its output. See also
:meth:`from_vla`.
reward_aggregation (str, optional): "last", "stack", "sum", or "mean".
An explicit value overrides stack_rewards. Sum and mean reduce
only executed steps without allocating a reward stack. Defaults
to None, preserving the stack_rewards behavior. See also
:class:`~torchrl.trainers.algorithms.configs.MultiActionConfig`.
.. seealso:: :class:`~torchrl.envs.transforms.ActionChunkTransform` -- when
the stacked actions are a chunk policy's *prediction* (overlapping
per-step training targets) rather than a macro action to replay
verbatim. The chunk transform builds the training targets on the data
path and, attached to an env, executes only the first action of each
predicted chunk (re-planning at every step) instead of stepping the
base env once per action.
"""
def __init__(
self,
*,
dim: int = 1,
stack_rewards: bool = True,
stack_observations: bool = False,
action_key: NestedKey | None = None,
chunk_key: NestedKey | None = None,
reward_aggregation: Literal["last", "stack", "sum", "mean"] | None = None,
):
if action_key is None and chunk_key is not None:
action_key = "action"
if action_key is not None and chunk_key is None:
chunk_key = action_key
in_keys_inv = None if action_key is None else [action_key]
out_keys_inv = None if chunk_key is None else [chunk_key]
super().__init__(in_keys_inv=in_keys_inv, out_keys_inv=out_keys_inv)
if reward_aggregation not in (None, "last", "stack", "sum", "mean"):
raise ValueError(
"reward_aggregation must be last, stack, sum, mean, or None."
)
if isinstance(dim, bool) or not isinstance(dim, int) or dim < 1:
raise ValueError("dim must be a positive integer.")
self.reward_aggregation = reward_aggregation
self._reduce_rewards = reward_aggregation in ("sum", "mean", "last")
self.stack_rewards = (
stack_rewards
if reward_aggregation is None
else reward_aggregation == "stack"
)
self.stack_observations = stack_observations
self.dim = dim
[docs]
@classmethod
def from_vla(cls, *, action_key: NestedKey = ACTION_KEY, **kwargs) -> MultiAction:
"""Build a :class:`MultiAction` that consumes the default VLA chunk key.
Args:
action_key (NestedKey): the one-step action key consumed by the base
environment. Defaults to ``"action"``.
Keyword Args:
Additional :class:`MultiAction` keyword arguments.
Examples:
>>> from torchrl.envs.transforms import MultiAction
>>> transform = MultiAction.from_vla(stack_rewards=False)
>>> transform.out_keys_inv
[('vla_action', 'chunk')]
"""
return cls(action_key=action_key, chunk_key=ACTION_CHUNK_KEY, **kwargs)
def _stack_tds(self, td_list, next_tensordict, keys):
td = torch.stack(td_list + [next_tensordict.select(*keys)], -1)
if self.dim != 1:
d = td.ndim - 1
td.auto_batch_size_(d + self.dim)
td = td.transpose(d, d + self.dim)
return td
def _step(
self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
) -> TensorDictBase:
if self._reduce_rewards:
if tensordict is not None:
self._accumulate_rewards(
next_tensordict, self._final_active, self._final_global_idx
)
reward = self._reward_total
if self.reward_aggregation == "mean":
reward = reward / self._reward_count.clamp_min(1)
if tensordict is not None and self._final_global_idx is not None:
reward = reward[self._final_global_idx]
next_tensordict.update(reward)
# Collect the stacks if needed
if self.stack_rewards:
reward_td = self.rewards
reward_td = self._stack_tds(
reward_td, next_tensordict, self.parent.reward_keys
)
next_tensordict.update(reward_td)
if self.stack_observations:
obs_td = self.obs
obs_td = self._stack_tds(
obs_td, next_tensordict, self.parent.observation_keys
)
next_tensordict.update(obs_td)
return next_tensordict
def _reset(
self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
) -> TensorDictBase:
return tensordict_reset
def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase:
# Get the actions
parent = self.parent
action_keys = self.in_keys_inv or parent.action_keys
chunk_keys = self.out_keys_inv or action_keys
if len(action_keys) != len(chunk_keys):
raise ValueError(
"action_key and chunk_key lists must have the same length, got "
f"{len(action_keys)} and {len(chunk_keys)}."
)
actions = tensordict.empty()
for action_key, chunk_key in zip(action_keys, chunk_keys):
action = tensordict.get(chunk_key, None)
if action is None:
raise KeyError(
f"{type(self).__name__} expected stacked actions at key "
f"{chunk_key!r} before env.step, but the key was missing. "
"For VLA policies, use MultiAction.from_vla() or pass "
"chunk_key=('vla_action', 'chunk'). Available keys are "
f"{list(tensordict.keys(True, True))}."
)
actions.set(action_key, action)
actions = actions.auto_batch_size_(batch_dims=tensordict.ndim + self.dim)
actions = actions.unbind(-1)
return self._execute_steps(
tensordict,
len(actions),
functools.partial(self._write_chunk_action, actions=actions),
)
def _write_chunk_action(self, td, index, global_idx, active, *, actions):
action = actions[index]
if global_idx is not None:
action = action[global_idx]
return td.replace(action)
def _accumulate_rewards(self, next_td, active, global_idx):
reward_td = next_td.select(*self.parent.reward_keys)
if self._reward_total is None:
self._reward_total = reward_td.new_zeros(self._reward_batch_size)
if self.reward_aggregation == "mean":
self._reward_count = self._reward_total.clone()
for key, reward in reward_td.items(True, True):
live = active.to(reward.device)
live = live.reshape(*live.shape, *([1] * (reward.ndim - live.ndim)))
updates = [(self._reward_total, torch.where(live, reward, 0))]
if self.reward_aggregation == "mean":
updates.append((self._reward_count, live.expand_as(reward)))
for accumulator, increment in updates:
old = accumulator.get(key)
selected = old if global_idx is None else old[global_idx]
value = (
torch.where(live, reward, selected)
if self.reward_aggregation == "last"
else selected + increment
)
if global_idx is not None:
updated = old.clone()
updated[global_idx] = value
value = updated
accumulator.set(key, value)
def _execute_steps(self, tensordict, steps, write_action):
parent = self.parent
td = tensordict
if self._reduce_rewards:
self._reward_batch_size = td.batch_size
self._reward_total = self._reward_count = None
idx = None
global_idx = None
reset = False
if self.stack_rewards:
self.rewards = rewards = []
if self.stack_observations:
self.obs = obs = []
for index in range(steps - 1):
active = td.get(
"_step", torch.ones(td.shape, dtype=torch.bool, device=td.device)
)
td = write_action(td, index, global_idx, active)
td = parent.step(td)
if self._reduce_rewards:
self._accumulate_rewards(td["next"], active, global_idx)
# Save rewards and done states
if self.stack_rewards:
reward_td = td["next"].select(*self.parent.reward_keys)
if global_idx is not None:
reward_td_expand = reward_td.new_zeros(
global_idx.shape + reward_td.shape[global_idx.ndim :]
)
reward_td_expand[global_idx] = reward_td
else:
reward_td_expand = reward_td
rewards.append(reward_td_expand)
if self.stack_observations:
obs_td = td["next"].select(*self.parent.observation_keys)
# obs_td = td.select("next", *self.parent.observation_keys).set("next", obs_td)
if global_idx is not None:
expanded = obs_td.new_zeros(
global_idx.shape + obs_td.shape[global_idx.ndim :]
)
expanded[global_idx] = obs_td
obs_td = expanded
obs.append(obs_td)
td = parent.step_mdp(td)
if self.stack_rewards:
td.update(reward_td)
any_done = parent.any_done(td)
if any_done:
# Intersect the resets to avoid making any step after reset has been called
reset = reset | td.pop("_reset").view(td.shape)
if reset.all():
# Skip step for all
td["_step"] = ~reset
break
elif parent.batch_locked:
td["_step"] = ~reset
else:
# we can simply index the tensordict
idx = ~reset.view(td.shape)
if global_idx is None:
global_idx = idx.clone()
td_out = td
# td_out's root reward/observation tensors alias the
# entries just appended to the stacks: de-alias them so
# the masked writes into td_out below do not corrupt
# the stacked history
keys = []
if self.stack_rewards:
keys += list(self.parent.reward_keys)
if self.stack_observations:
keys += list(self.parent.observation_keys)
for key in keys:
td_out.set(key, td_out.get(key).clone())
else:
td_out[global_idx] = td
global_idx = torch.masked_scatter(global_idx, global_idx, idx)
td = td[idx]
reset = reset[idx] # Should be all False
active = td.get(
"_step", torch.ones(td.shape, dtype=torch.bool, device=td.device)
)
self._final_active = active
self._final_global_idx = global_idx
if global_idx is None:
td_out = write_action(td, steps - 1, None, active)
if (
self.stack_rewards or self.stack_observations or self._reduce_rewards
) and not td_out.get("_step", torch.ones((), dtype=torch.bool)).any():
if self.stack_rewards:
# the final outer step is skipped for every env (done fired
# inside the chunk): its slot in the reward stack would
# otherwise carry the stale reward of the last executed
# step - zero it, matching the zero-fill of the other
# skipped slots
for key in self.parent.reward_keys:
td_out.set(key, torch.zeros_like(td_out.get(key)))
td_out = self._step(None, td_out)
else:
td_out[global_idx] = write_action(td, steps - 1, global_idx, active)
if self.stack_rewards:
# zero the trailing reward slot of the envs that finished
# early: their final outer step is skipped, so it would
# otherwise carry the stale reward of their last executed step
for key in self.parent.reward_keys:
reward = td_out.get(key).clone()
reward[~global_idx] = 0
td_out.set(key, reward)
if self.stack_rewards or self.stack_observations or self._reduce_rewards:
td_out = self._step(None, td_out)
if self.stack_rewards:
self.rewards = list(
torch.stack(self.rewards, -1)[global_idx].unbind(-1)
)
if self.stack_observations:
self.obs = list(torch.stack(self.obs, -1)[global_idx].unbind(-1))
td_out["_step"] = global_idx
return td_out
def _transform_reward_spec(self, reward_spec: TensorSpec, ndim) -> TensorSpec:
if not self.stack_rewards:
return reward_spec
for _ in range(self.dim):
reward_spec = reward_spec.unsqueeze(ndim)
# Make the dim dynamic
reward_spec = reward_spec.expand(
tuple(
d if i != (ndim + self.dim - 1) else -1
for i, d in enumerate(reward_spec.shape)
)
)
return reward_spec
def _transform_observation_spec(
self, observation_spec: TensorSpec, ndim
) -> TensorSpec:
if not self.stack_observations:
return observation_spec
for _ in range(self.dim):
observation_spec = observation_spec.unsqueeze(ndim)
# Make the dim dynamic
observation_spec = observation_spec.expand(
tuple(
d if i != (ndim + self.dim - 1) else -1
for i, d in enumerate(observation_spec.shape)
)
)
return observation_spec
[docs]
class ClosedLoopMultiAction(MultiAction):
"""Execute a controller against fresh observations for one high-level decision.
Unlike an action chunk, the low-level action is recomputed at every physical
step. High-level decisions and their log probabilities remain on the outer
transition. Finished environments stop executing the controller.
Args:
controller (TensorDictModuleBase): low-level policy, typically
:class:`~torchrl.modules.LowLevelController`.
Keyword Args:
steps (int): positive number of physical steps per decision.
decision_spec (Composite, optional): complete policy-facing action spec,
including environment batch dimensions. Defaults to None, inferring
the spec from LowLevelController and preserving unrelated actions.
reward_aggregation (str, optional): "sum", "mean", "last", or "stack".
Defaults to "sum". Mean counts only executed steps; last returns
the last executed reward. Stack uses MultiAction's ragged convention.
exploration_type (ExplorationType, optional): controller sampling mode.
Defaults to DETERMINISTIC; the caller's exploration mode is restored.
no_grad (bool, optional): disable gradients during controller inference.
Defaults to True. This does not freeze the policy's parameters.
dim (int, optional): stack dimension relative to each leaf's containing
TensorDict batch dimensions. Defaults to 1, keeping agent dimensions
before the stack dimension.
stack_observations (bool, optional): return stacked inner observations.
Defaults to False (the final observation). Persistent state remains
unstacked. The controller uses the latest observation on its next call.
Use :meth:`from_env` to install controller primers before this transform.
The base environment must honor partial-step masks, as for MultiAction.
Discount factors on the resulting environment count high-level decisions.
Examples:
>>> import torch
>>> from tensordict.nn import TensorDictModule
>>> from torchrl.data import Bounded, Composite
>>> from torchrl.modules import LowLevelController
>>> from torchrl.testing.mocking_classes import CountingEnv
>>> policy = TensorDictModule(
... torch.nn.Identity(), in_keys=["command"], out_keys=["action"])
>>> controller = LowLevelController(
... policy, Composite(command=Bounded(0, 1, shape=(1,))))
>>> env = ClosedLoopMultiAction.from_env(CountingEnv(), controller, steps=3)
>>> td = env.reset().set("command", torch.ones(1))
>>> env.step(td)["next", "observation"]
tensor([3], dtype=torch.int32)
>>> env.close()
.. seealso::
:class:`~torchrl.modules.LowLevelController` provides independent
recurrent state for each controlled instance;
:class:`~torchrl.envs.MicroDuckSkillEnv` uses this transform to expose
skill decisions as environment actions; and
:class:`~torchrl.trainers.algorithms.configs.ClosedLoopMultiActionConfig`
exposes this class through Hydra configuration.
"""
def __init__(
self,
controller: TensorDictModuleBase,
*,
steps: int,
decision_spec: Composite | None = None,
reward_aggregation: Literal["last", "stack", "sum", "mean"] = "sum",
exploration_type: ExplorationType = ExplorationType.DETERMINISTIC,
no_grad: bool = True,
dim: int = 1,
stack_observations: bool = False,
):
if isinstance(steps, bool) or not isinstance(steps, int) or steps < 1:
raise ValueError("steps must be a positive integer.")
super().__init__(
dim=dim,
reward_aggregation=reward_aggregation,
stack_observations=stack_observations,
)
if decision_spec is None and not isinstance(controller, LowLevelController):
raise ValueError(
"Pass decision_spec when controller is not a LowLevelController."
)
self.controller = controller
self.steps = steps
self.decision_spec = None if decision_spec is None else decision_spec.clone()
self.exploration_type = exploration_type
self.no_grad = no_grad
self._decision_keys = None
[docs]
@classmethod
def from_env(
cls,
env: EnvBase,
controller: TensorDictModuleBase,
*,
steps: int,
init_key: str = "is_init",
**kwargs: Any,
) -> TransformedEnv:
"""Wrap an environment, automatically installing controller state.
Args:
env (EnvBase): physical environment.
controller (TensorDictModuleBase): controller to execute.
Keyword Args:
steps (int): positive number of controller steps per decision.
init_key (str, optional): episode-start marker. Defaults to "is_init".
**kwargs: additional ClosedLoopMultiAction constructor arguments.
Returns:
TransformedEnv: environment exposing high-level actions.
"""
pending = [get_env_transforms_from_module(controller, init_key=init_key)]
transforms = []
while pending:
transform = pending.pop(0)
if isinstance(transform, Compose):
pending[0:0] = list(transform.transforms)
else:
transform.reset_parent()
transforms.append(transform)
transforms.append(cls(controller, steps=steps, **kwargs))
return TransformedEnv(env, Compose(*transforms))
def _write_controller_action(self, td, index, global_idx, active, *, decisions):
if not active.any():
return td
if global_idx is not None:
decisions = decisions[global_idx]
all_active = active.all()
active_td = td if all_active else td[active]
active_td.update(decisions if all_active else decisions[active])
with (
torch.no_grad() if self.no_grad else nullcontext(),
set_exploration_type(self.exploration_type),
):
result = self.controller(active_td)
if all_active:
return result
# Indexed assignment allocates zero-filled tensors for new keys.
# Seed next-state entries from their current values so stopped rows
# retain their controller state when only live rows produce outputs.
next_state = result.get("next", None)
if next_state is not None:
for key in next_state.keys(True, True):
path = ("next", key) if isinstance(key, str) else ("next", *key)
current = td.get(key, None)
if current is not None and td.get(path, None) is None:
td.set(path, current.clone())
# active can be td["_step"], which indexed assignment also writes.
# Keep the indexing mask independent of the destination tensors.
td[active.clone()] = result
return td
def _stack_tds(self, td_list, next_tensordict, keys):
# Keep group batch dimensions stable: rewards can be reduced while
# observations are stacked in the same agent TensorDict.
result = next_tensordict.select(*keys).clone(recurse=False)
state_keys = self.parent.full_state_spec.keys(True, True)
for key, value in result.items(True, True):
if key in state_keys:
continue
node = result if isinstance(key, str) else result[key[:-1]]
result.set(
key,
torch.stack(
[td.get(key) for td in td_list] + [value], node.ndim + self.dim - 1
),
)
return result
def _stack_spec(self, spec):
state_keys = self.parent.full_state_spec.keys(True, True)
for key, leaf in list(spec.items(True, True)):
if key in state_keys:
continue
node = spec if isinstance(key, str) else spec[key[:-1]]
dim = node.ndim + self.dim - 1
leaf = leaf.unsqueeze(dim)
shape = list(leaf.shape)
shape[dim] = -1
spec[key] = leaf.expand(shape)
return spec
def _transform_reward_spec(self, reward_spec, ndim):
return self._stack_spec(reward_spec) if self.stack_rewards else reward_spec
def _transform_observation_spec(self, observation_spec, ndim):
return (
self._stack_spec(observation_spec)
if self.stack_observations
else observation_spec
)
def _reset(self, tensordict, tensordict_reset):
if self.stack_observations:
tensordict_reset.update(
self._stack_tds([], tensordict_reset, self.parent.observation_keys)
)
return tensordict_reset
def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase:
if self.stack_observations:
for key, leaf in self.parent.observation_spec.items(True, True):
value = tensordict.get(key, None)
if value is None:
continue
extra_batch_dims = tensordict.ndim - self.parent.ndim
if value.ndim == leaf.ndim + extra_batch_dims + 1:
node = tensordict if isinstance(key, str) else tensordict[key[:-1]]
tensordict.set(key, value.select(node.ndim + self.dim - 1, -1))
if self._decision_keys is None:
self.transform_input_spec(self.parent.input_spec)
decisions = tensordict.select(*self._decision_keys).clone()
return self._execute_steps(
tensordict,
self.steps,
functools.partial(self._write_controller_action, decisions=decisions),
)
[docs]
class ActionScaling(Transform):
r"""Affine-scale a continuous action using the bounds of the action spec.
Given a bounded action spec with bounds ``[low, high]``, this transform exposes
a normalized action space to the policy and rescales actions back to the
original env range before they are passed to the environment.
The ``loc`` and ``scale`` are derived from the spec:
.. math::
loc = \frac{high + low}{2}, \quad scale = \frac{high - low}{2}.
When ``standard_normal=True`` (default) the normalized action space is
``[-1, 1]`` and the inverse mapping (policy action -> env action) is
.. math::
a_{env} = a_{norm} \cdot scale + loc.
The forward mapping (env action -> normalized action, used by replay buffer
transforms) is the inverse:
.. math::
a_{norm} = (a_{env} - loc) / scale.
When ``standard_normal=False`` the normalized space is ``[0, 1]`` and the
mapping is rescaled accordingly so that ``0`` maps to ``low`` and ``1`` to
``high``.
Args:
in_keys_inv (sequence of NestedKey, optional): keys read during the
``inv`` direction (policy -> env). Defaults to ``["action"]``. A
single key per :class:`ActionScaling` instance is supported; compose
several instances to scale several actions. Pass an empty list for
a forward-only transform (normalize raw dataset actions on the
replay-buffer sample path while leaving ``extend`` and the env-side
action interface untouched); this requires explicit ``loc`` and
``scale``.
out_keys_inv (sequence of NestedKey, optional): keys written during the
``inv`` direction. Defaults to ``in_keys_inv``.
in_keys (sequence of NestedKey, optional): keys read during the forward
direction (env action -> normalized action, used by replay buffers
and inside :class:`~torch.nn.Module` chains). Defaults to
``in_keys_inv``, or ``["action"]`` when ``in_keys_inv=[]``
(forward-only mode).
out_keys (sequence of NestedKey, optional): keys written during the
forward direction. Defaults to ``in_keys``.
Keyword Args:
loc (torch.Tensor or float, optional): explicit location of the affine
transform. If both ``loc`` and ``scale`` are provided the values are
used as-is and no derivation from the spec is performed (useful when
no parent environment is available, e.g. inside a replay buffer).
Defaults to ``None``.
scale (torch.Tensor or float, optional): explicit scale of the affine
transform. Must be provided together with ``loc``.
Defaults to ``None``.
standard_normal (bool, optional): if ``True`` (default), the normalized
action space is ``[-1, 1]``. If ``False``, the normalized action
space is ``[0, 1]``.
Raises:
RuntimeError: if ``loc`` and ``scale`` are derived from the spec (no
explicit values passed) and the action spec is unbounded or
partially unbounded (any bound is non-finite). With explicit
``loc``/``scale``, a bounded spec is mapped through the affine
transform and an unbounded (or partially unbounded) spec is
advertised as ``Unbounded`` instead of raising.
With explicit ``loc`` and ``scale`` the transform is fully spec-independent
-- the standard workflow when training on dataset action statistics, e.g.
for VLA policies. Use :meth:`from_stats` (``mean``/``std`` or
``low``/``high``) or :meth:`from_metadata` to build such an instance from
dataset statistics. Attached to an environment, it denormalizes the
policy's actions on the inverse path: a bounded action spec is mapped
through the affine transform (and an unbounded action spec stays
unbounded), so the advertised normalized space reflects the actual
statistics rather than being assumed ``[-1, 1]``. Appended to a replay
buffer, it normalizes actions on the ``sample`` path; beware that
``ReplayBuffer.extend`` applies the *inverse* transform, so when raw
(env-scale) data is written through ``extend``, use a forward-only
instance (``in_keys_inv=[]``) to leave the stored data untouched -- the
default bidirectional keys suit the env side and pre-populated dataset
storages.
Examples:
>>> import torch
>>> from torchrl.data.tensor_specs import Bounded
>>> from torchrl.envs.transforms import ActionScaling, TransformedEnv
>>> from torchrl.testing.mocking_classes import ContinuousActionVecMockEnv
>>> base_env = ContinuousActionVecMockEnv(
... action_spec=Bounded(low=-2.0, high=4.0, shape=(7,))
... )
>>> env = TransformedEnv(base_env, ActionScaling())
>>> env.action_spec.space.low
tensor([-1., -1., -1., -1., -1., -1., -1.])
>>> env.action_spec.space.high
tensor([1., 1., 1., 1., 1., 1., 1.])
>>> # dataset-statistics-driven normalization (no env required): the
>>> # forward pass maps raw actions to the normalized space
>>> from tensordict import TensorDict
>>> t = ActionScaling.from_stats(
... mean=torch.tensor([1.0, 2.0]), std=torch.tensor([2.0, 4.0])
... )
>>> td = TensorDict({"action": torch.tensor([[3.0, 6.0]])}, batch_size=[1])
>>> t(td)["action"]
tensor([[1., 1.]])
>>> # on a replay buffer, a forward-only instance (in_keys_inv=[])
>>> # normalizes on sample and leaves data written through extend
>>> # untouched (extend applies the inverse pass)
>>> from torchrl.data import LazyTensorStorage, TensorDictReplayBuffer
>>> t = ActionScaling.from_stats(
... mean=torch.tensor([1.0, 2.0]),
... std=torch.tensor([2.0, 4.0]),
... in_keys_inv=[],
... )
>>> rb = TensorDictReplayBuffer(
... storage=LazyTensorStorage(10), transform=t, batch_size=2
... )
>>> raw = TensorDict(
... {"action": torch.tensor([[3.0, 6.0]]).expand(10, 2)}, batch_size=[10]
... )
>>> indices = rb.extend(raw) # stored as-is
>>> rb.sample()["action"] # normalized with the dataset statistics
tensor([[1., 1.],
[1., 1.]])
>>> # the same affine map is exposed on raw tensors for execution-time
>>> # use, e.g. mapping a policy's normalized prediction to the robot
>>> t.denormalize(torch.tensor([[1.0, 1.0]]))
tensor([[3., 6.]])
"""
invertible = True
def __init__(
self,
in_keys_inv: Sequence[NestedKey] | None = None,
out_keys_inv: Sequence[NestedKey] | None = None,
in_keys: Sequence[NestedKey] | None = None,
out_keys: Sequence[NestedKey] | None = None,
*,
loc: torch.Tensor | float | None = None,
scale: torch.Tensor | float | None = None,
standard_normal: bool = True,
):
if in_keys_inv is None:
in_keys_inv = ["action"]
if not isinstance(in_keys_inv, (list, tuple)):
in_keys_inv = [in_keys_inv]
if len(in_keys_inv) > 1:
raise ValueError(
"ActionScaling only supports a single action key per instance. "
"Compose several ActionScaling transforms to scale multiple actions."
)
if out_keys_inv is None:
out_keys_inv = copy(in_keys_inv)
if in_keys is None:
# Forward-only mode (``in_keys_inv=[]``) still normalizes "action"
# on the forward (sample) path by default.
in_keys = copy(in_keys_inv) if in_keys_inv else ["action"]
if out_keys is None:
out_keys = copy(in_keys)
super().__init__(
in_keys=in_keys,
out_keys=out_keys,
in_keys_inv=in_keys_inv,
out_keys_inv=out_keys_inv,
)
self.standard_normal = bool(standard_normal)
if (loc is None) != (scale is None):
raise ValueError(
"loc and scale must either both be provided or both be None."
)
self._explicit = loc is not None
if not in_keys_inv and not self._explicit:
raise ValueError(
"in_keys_inv=[] (forward-only mode) requires explicit loc and "
"scale: without an inverse action key there is no action spec "
"to derive them from."
)
if loc is not None:
loc = torch.as_tensor(loc)
scale = torch.as_tensor(scale)
if not loc.dtype.is_floating_point:
loc = loc.to(torch.get_default_dtype())
if not scale.dtype.is_floating_point:
scale = scale.to(torch.get_default_dtype())
if (scale == 0).any():
raise ValueError(
"scale must not contain zero entries (would cause division by zero)."
)
self.register_buffer("loc", loc)
self.register_buffer("scale", scale)
else:
self.register_buffer("loc", nn.UninitializedBuffer())
self.register_buffer("scale", nn.UninitializedBuffer())
@property
def initialized(self) -> bool:
return not isinstance(self.loc, nn.UninitializedBuffer)
def _ensure_initialized(self) -> None:
# Lazily populate ``loc`` and ``scale`` from the parent env's action
# spec at the insertion point of this transform. ``self.parent`` is
# rebuilt with all transforms up to (but not including) ``self``, so
# its action spec is exactly the env-scale spec we need to read.
if self.initialized:
return
parent = self.parent
if parent is None:
raise RuntimeError(
"ActionScaling has not been initialized: pass explicit ``loc`` "
"and ``scale`` to the constructor, or attach this transform to "
"a TransformedEnv whose action spec is bounded so that the "
"values can be derived automatically."
)
in_key = unravel_key(self.in_keys_inv[0])
full_action_spec = parent.full_action_spec
if in_key not in full_action_spec.keys(True, True):
raise RuntimeError(
f"ActionScaling could not find key {in_key!r} in the parent "
f"environment's action spec. Available keys: "
f"{list(full_action_spec.keys(True, True))}."
)
self._init_from_spec(full_action_spec[in_key])
def _init_from_spec(self, leaf_spec: TensorSpec) -> None:
low, high = self._validate_bounded(leaf_spec)
dtype = low.dtype if low.dtype.is_floating_point else torch.get_default_dtype()
loc = ((high + low) / 2).to(dtype)
scale = ((high - low) / 2).to(dtype)
self._materialize_loc_scale(loc, scale)
def _materialize_loc_scale(self, loc: torch.Tensor, scale: torch.Tensor) -> None:
if isinstance(self.loc, nn.UninitializedBuffer):
self.loc.materialize(shape=loc.shape, dtype=loc.dtype)
self.scale.materialize(shape=scale.shape, dtype=scale.dtype)
self.loc.data.copy_(loc)
self.scale.data.copy_(scale)
@staticmethod
def _validate_bounded(action_spec: TensorSpec) -> tuple[torch.Tensor, torch.Tensor]:
space = getattr(action_spec, "space", None)
if not isinstance(space, ContinuousBox):
raise RuntimeError(
f"ActionScaling requires a bounded continuous action spec, got "
f"{type(action_spec).__name__} with space "
f"{type(space).__name__ if space is not None else None}. "
"Unbounded or discrete action specs are not supported."
)
# ``Unbounded`` specs use a ``ContinuousBox`` whose low/high are set to
# ``finfo.min`` and ``finfo.max`` respectively, so checking the spec type
# is more reliable than ``torch.isfinite``.
if isinstance(action_spec, Unbounded):
raise RuntimeError(
"ActionScaling cannot be used with an Unbounded action spec. "
"The action spec must be fully bounded for spec-based normalization."
)
low = space.low
high = space.high
# Partially unbounded: one side is finite but the other matches the
# ``finfo`` extreme used internally by ``Unbounded``.
dtype = low.dtype
if dtype.is_floating_point:
extreme_low = torch.finfo(dtype).min
extreme_high = torch.finfo(dtype).max
if (low == extreme_low).any() or (high == extreme_high).any():
raise RuntimeError(
"ActionScaling requires fully bounded actions: at least one "
"entry of the action spec is unbounded (low equals finfo.min or "
"high equals finfo.max)."
)
if not torch.isfinite(low).all() or not torch.isfinite(high).all():
raise RuntimeError(
"ActionScaling requires fully bounded actions: every entry of the "
"action spec must have a finite lower and upper bound. Got "
"non-finite values in low or high."
)
if (high <= low).any():
raise RuntimeError(
"ActionScaling requires high > low for every entry of the action "
"spec. Got entries with high <= low."
)
return low, high
@staticmethod
def _is_finitely_bounded(leaf_spec: TensorSpec) -> bool:
# ``Unbounded`` (and partially-unbounded ``Bounded``) specs encode the
# open sides with ``finfo`` extremes; mapping those through the affine
# would overflow, so they are treated as unbounded instead.
if isinstance(leaf_spec, Unbounded):
return False
low, high = leaf_spec.space.low, leaf_spec.space.high
if low.dtype.is_floating_point:
extreme_low = torch.finfo(low.dtype).min
extreme_high = torch.finfo(high.dtype).max
if (low == extreme_low).any() or (high == extreme_high).any():
return False
return bool(torch.isfinite(low).all() and torch.isfinite(high).all())
def _transform_leaf(self, leaf_spec: TensorSpec) -> TensorSpec:
dtype = (
leaf_spec.dtype
if leaf_spec.dtype.is_floating_point
else torch.get_default_dtype()
)
if self._explicit:
# Explicit loc/scale: no bounds are required from the spec. A
# bounded spec is mapped through the forward affine (monotonic,
# scale > 0); an unbounded (or partially unbounded) spec stays
# unbounded, since the affine image of an unbounded space is
# unbounded.
space = getattr(leaf_spec, "space", None)
if not isinstance(space, ContinuousBox):
raise RuntimeError(
f"ActionScaling requires a continuous action spec, got "
f"{type(leaf_spec).__name__}. Discrete action specs are "
"not supported."
)
if not self._is_finitely_bounded(leaf_spec):
return Unbounded(
shape=leaf_spec.shape,
device=leaf_spec.device,
dtype=leaf_spec.dtype,
)
loc, scale = self._loc_scale(space.low.device)
new_low = (space.low.to(dtype) - loc) / scale
new_high = (space.high.to(dtype) - loc) / scale
if not self.standard_normal:
new_low = (new_low + 1) / 2
new_high = (new_high + 1) / 2
else:
# Spec-derived loc/scale: bounds are mandatory and define the
# normalized space exactly ([-1, 1] or [0, 1]).
if not self.initialized:
self._init_from_spec(leaf_spec)
else:
self._validate_bounded(leaf_spec)
low = leaf_spec.space.low.to(dtype)
high = leaf_spec.space.high.to(dtype)
if self.standard_normal:
new_low = torch.full_like(low, -1.0)
new_high = torch.full_like(high, 1.0)
else:
new_low = torch.zeros_like(low)
new_high = torch.ones_like(high)
return Bounded(
low=new_low,
high=new_high,
shape=leaf_spec.shape,
device=leaf_spec.device,
dtype=leaf_spec.dtype,
)
def _loc_scale(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
# Only move the buffers when the device actually differs: an
# unconditional ``.to()`` inserts a copy node in the compile graph,
# whereas the device comparison is resolved at trace time (device is
# static metadata, not data), so the common same-device path stays
# copy-free and compile-friendly.
loc, scale = self.loc, self.scale
if loc.device != device:
loc = loc.to(device)
scale = scale.to(device)
return loc, scale
def _check_dim(self, action: torch.Tensor) -> None:
# Guard the silent-broadcast hazard: a per-dimension loc/scale must
# match the action's trailing dim. A scalar (or shape-[1]) loc/scale
# broadcasts freely and is left alone.
if self.loc.numel() > 1 and (
action.ndim == 0 or action.shape[-1] != self.loc.shape[-1]
):
raise ValueError(
f"action shape {tuple(action.shape)} does not match the "
f"loc/scale dimension {self.loc.shape[-1]} on its last dim."
)
def _apply_transform(self, action: torch.Tensor) -> torch.Tensor:
self._ensure_initialized()
self._check_dim(action)
loc, scale = self._loc_scale(action.device)
normalized = (action - loc) / scale
if not self.standard_normal:
normalized = (normalized + 1) / 2
return normalized
def _inv_apply_transform(self, action: torch.Tensor) -> torch.Tensor:
self._ensure_initialized()
self._check_dim(action)
loc, scale = self._loc_scale(action.device)
if not self.standard_normal:
action = action * 2 - 1
return action * scale + loc
[docs]
def normalize(self, action: torch.Tensor) -> torch.Tensor:
"""Map an env-scale action to the normalized space (the forward map)."""
return self._apply_transform(action)
[docs]
def denormalize(self, action: torch.Tensor) -> torch.Tensor:
"""Map a normalized action back to the env scale (the inverse map)."""
return self._inv_apply_transform(action)
[docs]
@classmethod
def from_stats(
cls,
*,
mean: torch.Tensor | None = None,
std: torch.Tensor | None = None,
low: torch.Tensor | None = None,
high: torch.Tensor | None = None,
eps: float = 1e-6,
**kwargs,
) -> ActionScaling:
"""Build an :class:`ActionScaling` from dataset action statistics.
Provide exactly one complete pair: ``mean`` and ``std`` (zero-mean,
unit-std normalized space) or ``low`` and ``high`` (maps the range to
``[-1, 1]``).
Keyword Args:
mean (torch.Tensor, optional): per-dimension action mean.
std (torch.Tensor, optional): per-dimension action std.
low (torch.Tensor, optional): per-dimension action minimum.
high (torch.Tensor, optional): per-dimension action maximum.
eps (float, optional): floor applied to the scale to avoid division
by zero on constant action dimensions. Defaults to ``1e-6``.
**kwargs: forwarded to the constructor (e.g. ``in_keys_inv``,
``standard_normal``).
"""
if (mean is None) != (std is None):
raise ValueError("mean and std must be provided together.")
if (low is None) != (high is None):
raise ValueError("low and high must be provided together.")
if (mean is not None) == (low is not None):
raise ValueError("Provide exactly one of (mean, std) or (low, high).")
if mean is not None:
loc = torch.as_tensor(mean, dtype=torch.get_default_dtype())
scale = torch.as_tensor(std, dtype=torch.get_default_dtype())
else:
low = torch.as_tensor(low, dtype=torch.get_default_dtype())
high = torch.as_tensor(high, dtype=torch.get_default_dtype())
loc = (low + high) / 2
scale = (high - low) / 2
if loc.shape != scale.shape:
raise ValueError(
f"loc and scale must have the same shape, got {tuple(loc.shape)} "
f"and {tuple(scale.shape)}."
)
return cls(loc=loc, scale=scale.clamp_min(eps), **kwargs)
def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase:
# The action only flows through the inv direction during env stepping;
# the ``next_tensordict`` returned by the base env does not contain it.
# Overriding ``_call`` as a no-op avoids the default loop raising a
# ``KeyError`` for the missing action key. The forward direction
# (env action -> normalized) is still wired through ``forward`` /
# ``_apply_transform`` for replay buffers and ``nn.Module`` chains.
return next_tensordict
def __repr__(self) -> str:
loc = self.loc if self.initialized else "<uninitialized>"
scale = self.scale if self.initialized else "<uninitialized>"
return (
f"{self.__class__.__name__}("
f"loc={loc}, scale={scale}, standard_normal={self.standard_normal}, "
f"in_keys_inv={self.in_keys_inv})"
)
[docs]
class FlattenAction(Transform):
"""Flatten adjacent dimensions of an action.
Mirrors :class:`~torchrl.envs.transforms.FlattenObservation`, but applies
to actions: the policy sees a flattened action space and the original
multi-dimensional shape is restored on the inv direction before the action
is passed to the base environment.
On the inv direction (policy -> env), a 1-D ``flattened`` action is
unflattened to the original ``(dim_first, ..., dim_last)`` span of the env
action. On the forward direction (env action -> flattened, used inside
replay buffers and :class:`~torch.nn.Module` chains), the adjacent dims
``[first_dim, last_dim]`` are flattened.
Args:
first_dim (int): first dimension to flatten. Must be negative unless
``allow_positive_dim`` is ``True``.
last_dim (int): last dimension to flatten (inclusive). Must be negative
unless ``allow_positive_dim`` is ``True``.
in_keys_inv (sequence of NestedKey, optional): keys read during the
``inv`` direction (policy -> env). Defaults to ``["action"]``.
Multiple keys are supported - the same flatten span is applied to
each one, which is useful for dict-structured action spaces.
out_keys_inv (sequence of NestedKey, optional): keys written during the
``inv`` direction. Defaults to ``in_keys_inv``.
in_keys (sequence of NestedKey, optional): keys read during the forward
direction (env action -> flattened). Defaults to ``in_keys_inv``.
out_keys (sequence of NestedKey, optional): keys written during the
forward direction. Defaults to ``in_keys``.
allow_positive_dim (bool, optional): if ``True``, positive dimensions
are accepted. Defaults to ``False`` so that the same transform
works regardless of the parent environment's batch size.
Keyword Args:
action_shape (sequence of int, optional): explicit pre-flatten shape
of the dimensions ``[first_dim, last_dim]``. Useful when the
transform is used outside a :class:`TransformedEnv` (e.g. inside
a replay buffer) and the original action shape cannot be derived
from a parent env. The same span is applied to every entry of
``in_keys_inv``. Defaults to ``None``, in which case the shape is
derived lazily from the parent env's action spec.
Examples:
>>> import torch
>>> from torchrl.data.tensor_specs import Bounded
>>> from torchrl.envs.transforms import FlattenAction, TransformedEnv
>>> from torchrl.testing.mocking_classes import ContinuousActionVecMockEnv
>>> base_env = ContinuousActionVecMockEnv(
... action_spec=Bounded(low=-1.0, high=1.0, shape=(3, 5))
... )
>>> env = TransformedEnv(base_env, FlattenAction(first_dim=-2, last_dim=-1))
>>> env.action_spec.shape
torch.Size([15])
"""
invertible = True
def __init__(
self,
first_dim: int = -2,
last_dim: int = -1,
in_keys_inv: Sequence[NestedKey] | None = None,
out_keys_inv: Sequence[NestedKey] | None = None,
in_keys: Sequence[NestedKey] | None = None,
out_keys: Sequence[NestedKey] | None = None,
allow_positive_dim: bool = False,
*,
action_shape: Sequence[int] | None = None,
):
if in_keys_inv is None:
in_keys_inv = ["action"]
if not isinstance(in_keys_inv, (list, tuple)):
in_keys_inv = [in_keys_inv]
if out_keys_inv is None:
out_keys_inv = copy(list(in_keys_inv))
if in_keys is None:
in_keys = copy(list(in_keys_inv))
if out_keys is None:
out_keys = copy(list(in_keys))
super().__init__(
in_keys=in_keys,
out_keys=out_keys,
in_keys_inv=in_keys_inv,
out_keys_inv=out_keys_inv,
)
if not allow_positive_dim and first_dim >= 0:
raise ValueError(
"first_dim should be smaller than 0 to accommodate for "
"envs of different batch_sizes. Set allow_positive_dim=True "
"to allow positive dimensions."
)
if not allow_positive_dim and last_dim >= 0:
raise ValueError(
"last_dim should be smaller than 0 to accommodate for "
"envs of different batch_sizes. Set allow_positive_dim=True "
"to allow positive dimensions."
)
if first_dim > last_dim:
raise ValueError(
f"first_dim ({first_dim}) must be <= last_dim ({last_dim})."
)
self._first_dim = first_dim
self._last_dim = last_dim
self.allow_positive_dim = bool(allow_positive_dim)
# Per-action-key original (pre-flatten) span, populated from the spec
# or seeded from the ``action_shape`` constructor kwarg.
self._unflatten_shapes: dict[NestedKey, tuple[int, ...]] = {}
if action_shape is not None:
action_shape = tuple(int(s) for s in action_shape)
for in_key in self.in_keys_inv:
self._unflatten_shapes[unravel_key(in_key)] = action_shape
@property
def first_dim(self) -> int:
if self._first_dim >= 0 and self.parent is not None:
return len(self.parent.batch_size) + self._first_dim
return self._first_dim
@property
def last_dim(self) -> int:
if self._last_dim >= 0 and self.parent is not None:
return len(self.parent.batch_size) + self._last_dim
return self._last_dim
@property
def _flat_merged_dim(self) -> int:
# Index of the merged dim in the post-flatten tensor. The flat tensor
# has ``(last - first)`` fewer dims than the original. For positive
# ``first_dim``, the merged dim sits at exactly ``first_dim``. For
# negative ``first_dim``, ``last_dim`` (also negative) already points
# at the merged dim in the new, shorter tensor, because
# ``last_dim - first_dim`` dims were collapsed strictly to its left.
if self._first_dim >= 0:
return self.first_dim
return self.last_dim
def _apply_transform(self, action: torch.Tensor) -> torch.Tensor:
# env-scale action -> flattened
return torch.flatten(action, self.first_dim, self.last_dim)
def _inv_apply_transform(self, action: torch.Tensor) -> torch.Tensor:
# flattened action -> env-scale (unflatten)
self._ensure_unflatten_shapes()
# ``_inv_apply_transform`` only receives a tensor, with no information
# about which ``in_keys_inv`` entry it came from. For multi-key
# transforms we cannot disambiguate, so we route those through
# ``_inv_call`` (which knows the key) and raise here. Single-key
# instances are unambiguous and remain supported.
if len(self.in_keys_inv) != 1:
raise RuntimeError(
f"FlattenAction._inv_apply_transform cannot disambiguate "
f"between {len(self.in_keys_inv)} action keys. Use "
f"``FlattenAction.inv(td)`` / ``_inv_call(td)`` instead, which "
f"know which key each tensor belongs to."
)
in_key = unravel_key(self.in_keys_inv[0])
shape = self._unflatten_shapes.get(in_key)
if shape is None:
raise RuntimeError(
f"FlattenAction has no stored unflatten shape for key "
f"{in_key!r}. Pass ``action_shape`` to the constructor or "
f"attach the transform to a TransformedEnv with a bounded "
f"action spec for this key."
)
return torch.unflatten(action, self._flat_merged_dim, shape)
def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase:
# Route each action key to its own unflatten span using the per-key
# state computed at ``transform_action_spec`` time.
if not self.in_keys_inv:
return tensordict
self._ensure_unflatten_shapes()
flat_dim = self._flat_merged_dim
for in_key, out_key in zip(self.in_keys_inv, self.out_keys_inv):
in_key_u = unravel_key(in_key)
out_key_u = unravel_key(out_key)
data = tensordict.get(out_key_u, default=None)
if data is None:
if not self.missing_tolerance:
raise KeyError(
f"'{out_key_u}' not found in tensordict {tensordict}"
)
continue
shape = self._unflatten_shapes.get(in_key_u)
if shape is None:
raise RuntimeError(
f"FlattenAction has no stored unflatten shape for key "
f"{in_key_u!r}. Pass ``action_shape`` to the constructor "
f"or attach the transform to a TransformedEnv with a "
f"bounded action spec for this key."
)
tensordict.set(in_key_u, torch.unflatten(data, flat_dim, shape))
return tensordict
def _ensure_unflatten_shapes(self) -> None:
# Lazily populate ``_unflatten_shapes`` from the parent env's action
# spec at the insertion point of this transform. ``self.parent`` is
# rebuilt with all transforms up to (but not including) ``self``, so
# its action spec is exactly the env-scale spec we need to read. If
# ``action_shape`` was provided at construction time this is a no-op.
if self._unflatten_shapes:
return
parent = self.parent
if parent is None:
return
full_action_spec = parent.full_action_spec
for in_key in self.in_keys_inv:
in_key = unravel_key(in_key)
if in_key in full_action_spec.keys(True, True):
self._unflatten_shapes[in_key] = self._span_from_spec(
full_action_spec[in_key]
)
def _span_from_spec(self, leaf_spec: TensorSpec) -> tuple[int, ...]:
ndim = len(leaf_spec.shape)
first = self._first_dim
last = self._last_dim
if first < 0:
first = ndim + first
if last < 0:
last = ndim + last
if first < 0 or last >= ndim or first > last:
raise RuntimeError(
f"FlattenAction(first_dim={self._first_dim}, last_dim={self._last_dim}) "
f"is not compatible with an action of shape {tuple(leaf_spec.shape)}."
)
return tuple(int(s) for s in leaf_spec.shape[first : last + 1])
def _flatten_leaf(self, leaf_spec: TensorSpec) -> TensorSpec:
space = getattr(leaf_spec, "space", None)
if isinstance(space, ContinuousBox):
new_low = torch.flatten(space.low, self.first_dim, self.last_dim)
new_high = torch.flatten(space.high, self.first_dim, self.last_dim)
return Bounded(
low=new_low,
high=new_high,
shape=new_low.shape,
device=leaf_spec.device,
dtype=leaf_spec.dtype,
)
shape = list(leaf_spec.shape)
ndim = len(shape)
first = self.first_dim if self.first_dim >= 0 else ndim + self.first_dim
last = self.last_dim if self.last_dim >= 0 else ndim + self.last_dim
flat = math.prod(shape[first : last + 1])
new_shape = torch.Size((*shape[:first], flat, *shape[last + 1 :]))
leaf_spec = leaf_spec.clone()
leaf_spec.shape = new_shape
return leaf_spec
def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase:
# The action does not appear in ``next_tensordict`` during env
# stepping, so we leave it untouched here. The default ``_call`` loop
# would otherwise raise ``KeyError`` for the missing action key. The
# ``forward`` path used by replay buffers is still wired through
# ``_apply_transform``.
return next_tensordict
def __repr__(self) -> str:
return (
f"{self.__class__.__name__}("
f"first_dim={int(self.first_dim)}, last_dim={int(self.last_dim)}, "
f"in_keys_inv={self.in_keys_inv}, out_keys_inv={self.out_keys_inv})"
)
[docs]
class LastAction(Transform):
"""Copies the last action into the next observation.
This is the action analogue of :class:`~torchrl.envs.transforms.CatFrames`:
the policy can condition on the action taken at the previous step (delayed
control, recurrent policies, residual action heads). On each
:meth:`~torchrl.envs.EnvBase.step` the action at time ``t`` is written
under ``out_keys`` in the ``"next"`` tensordict; on
:meth:`~torchrl.envs.EnvBase.reset` the same keys are filled with a
default value (zeros, NaN, or a user-provided fill). On a batch-unlocked
parent the default is expanded by the runtime reset batch, preserving
the action feature shape.
``out_keys`` are registered as :class:`~torchrl.data.Unbounded`
observation specs with the action's shape, dtype and device, so reset
fills (zeros on a one-hot action, NaN on a bounded action) remain
in-spec.
Args:
in_keys (NestedKey or sequence of NestedKey, optional): keys pointing
to the actions to remember. Defaults to the parent environment's
:attr:`~torchrl.envs.EnvBase.action_keys` when the transform is
attached, or ``["action"]`` otherwise.
out_keys (NestedKey or sequence of NestedKey, optional): destination
keys written into the observation. Defaults to each ``in_keys``
entry with its last component replaced by ``"last_action"``
(e.g. ``"action"`` -> ``"last_action"``,
``("agents", "action")`` -> ``("agents", "last_action")``).
Keyword Args:
default (str, number or torch.Tensor, optional): value used to fill
``out_keys`` on :meth:`~torchrl.envs.EnvBase.reset`. ``"zeros"``
(default) writes zeros matching the action spec; ``"nan"`` writes
NaNs (floating-point action specs only); a scalar is broadcast
with :meth:`~torch.Tensor.fill_`; a tensor is broadcast to the
action spec shape on the spec's device and dtype. Defaults to
``"zeros"``.
reset_key (NestedKey, optional): the reset key to be used as a
partial-reset indicator. Must be unique. If not provided, defaults
to the only reset key of the parent environment (if it has only
one) and raises an exception otherwise.
Examples:
>>> from torchrl.envs import GymEnv, TransformedEnv
>>> from torchrl.envs.transforms import LastAction
>>> env = TransformedEnv(GymEnv("Pendulum-v1"), LastAction())
>>> td = env.reset()
>>> td["last_action"]
tensor([0.])
>>> rollout = env.rollout(3)
>>> (rollout["next", "last_action"] == rollout["action"]).all()
tensor(True)
.. seealso:: :class:`~torchrl.envs.transforms.CatFrames` for stacking past
observations, :class:`~torchrl.envs.transforms.InitTracker` for
marking episode starts, and
:class:`~torchrl.trainers.algorithms.configs.LastActionConfig` for the
Hydra configuration.
"""
invertible = False
def __init__(
self,
in_keys: Sequence[NestedKey] | NestedKey | None = None,
out_keys: Sequence[NestedKey] | NestedKey | None = None,
*,
default: Literal["zeros", "nan"] | float | int | torch.Tensor = "zeros",
reset_key: NestedKey | None = None,
):
if isinstance(default, str) and default not in ("zeros", "nan"):
raise ValueError(
f"{type(self).__name__} default must be 'zeros', 'nan', a "
f"number or a tensor, got {default!r}."
)
if isinstance(default, float) and math.isnan(default):
default = "nan"
super().__init__(in_keys=in_keys, out_keys=out_keys)
if isinstance(default, torch.Tensor):
self.register_buffer("_default_value", default.clone())
self.default: Literal["zeros", "nan", "tensor"] | float | int = "tensor"
else:
self.default = default
self.reset_key = reset_key
@property
def in_keys(self) -> Sequence[NestedKey]:
in_keys = self.__dict__.get("_in_keys", None)
if in_keys is not None:
return in_keys
parent = self.parent
if parent is None:
return ["action"]
in_keys = list(parent.action_keys)
self._in_keys = in_keys
return in_keys
@in_keys.setter
def in_keys(self, value: Sequence[NestedKey] | NestedKey | None) -> None:
if value is not None:
if isinstance(value, (str, tuple)):
value = [value]
value = [unravel_key(val) for val in value]
self._in_keys = value
@property
def out_keys(self) -> Sequence[NestedKey]:
out_keys = self.__dict__.get("_out_keys", None)
if out_keys is not None:
return out_keys
derived = [_replace_last(key, "last_action") for key in self.in_keys]
if self.__dict__.get("_in_keys", None) is not None:
self._out_keys = derived
return derived
@out_keys.setter
def out_keys(self, value: Sequence[NestedKey] | NestedKey | None) -> None:
if value is not None:
if isinstance(value, (str, tuple)):
value = [value]
value = [unravel_key(val) for val in value]
self._out_keys = value
@property
def reset_key(self) -> NestedKey:
reset_key = self.__dict__.get("_reset_key", None)
if reset_key is not None:
return reset_key
parent = self.parent
if parent is None:
raise RuntimeError(FORWARD_NOT_IMPLEMENTED.format(type(self).__name__))
reset_keys = parent.reset_keys
if len(reset_keys) > 1:
raise RuntimeError(
f"Got more than one reset key in env {self.container}, cannot "
f"infer which one to use. Consider providing the reset key in "
f"the {type(self)} constructor."
)
return reset_keys[0]
@reset_key.setter
def reset_key(self, value: NestedKey | None) -> None:
if value is not None:
value = unravel_key(value)
self._reset_key = value
def _make_default(
self, in_key: NestedKey, batch_size: torch.Size | None = None
) -> torch.Tensor:
parent = self.parent
if parent is None:
raise RuntimeError(FORWARD_NOT_IMPLEMENTED.format(type(self).__name__))
try:
spec = parent.full_action_spec[in_key]
except KeyError:
raise KeyError(
f"{type(self).__name__} in_key {in_key!r} is not in the "
f"parent action spec {parent.full_action_spec}."
) from None
# Unlocked parents can reset with a larger runtime batch than the spec.
extra_batch: torch.Size | tuple[()] = ()
if batch_size is not None and not parent.batch_locked:
extra_batch = batch_size
zeros = spec.zero(extra_batch)
default = self.default
if default == "zeros":
return zeros
if default == "nan":
if not zeros.dtype.is_floating_point:
raise ValueError(
f"{type(self).__name__} default='nan' requires a "
f"floating-point action spec, got dtype={zeros.dtype} "
f"for key {in_key!r}."
)
return zeros.fill_(float("nan"))
if default == "tensor":
fill = self._default_value.to(device=zeros.device, dtype=zeros.dtype)
return zeros.copy_(fill.expand_as(zeros))
return zeros.fill_(default)
def _step(
self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
) -> TensorDictBase:
for in_key, out_key in _zip_strict(self.in_keys, self.out_keys):
action = tensordict.get(in_key, default=None)
if action is None:
if not self.missing_tolerance:
raise KeyError(
f"{self}: '{in_key}' not found in tensordict {tensordict}"
)
continue
# Policies may reuse the action buffer in-place.
next_tensordict.set(out_key, action.clone())
return next_tensordict
def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase:
return next_tensordict
def _reset(
self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
) -> TensorDictBase:
_reset = _get_reset(self.reset_key, tensordict)
for in_key, out_key in _zip_strict(self.in_keys, self.out_keys):
fill = self._make_default(in_key, tensordict_reset.batch_size)
existing = tensordict.get(out_key, default=None)
if existing is None:
tensordict_reset.set(out_key, fill)
continue
tensordict_reset.set(
out_key,
torch.where(expand_as_right(_reset, existing), fill, existing),
)
return tensordict_reset
def _reset_on_native_autoreset(
self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
) -> TensorDictBase:
return self._reset(tensordict, tensordict_reset)
[docs]
def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
raise NotImplementedError(
FORWARD_NOT_IMPLEMENTED.format(self.__class__.__name__)
)
