ClosedLoopMultiAction#

class torchrl.envs.transforms.ClosedLoopMultiAction(controller: TensorDictModuleBase, *, steps: int, decision_spec: Composite | None = None, reward_aggregation: Literal['last', 'stack', 'sum', 'mean'] = 'sum', exploration_type: InteractionType = InteractionType.DETERMINISTIC, no_grad: bool = True, dim: int = 1, stack_observations: bool = False)[source]#

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.

Parameters:

controller (TensorDictModuleBase) – low-level policy, typically LowLevelController.

Keyword Arguments:

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 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()

ClosedLoopMultiAction#

Docs

Tutorials

Resources