MultiStepActorWrapper#

class torchrl.modules.tensordict_module.MultiStepActorWrapper(*args, **kwargs)[source]#

A wrapper around a multi-action actor.

This class enables macros to be executed in an environment. The actor action(s) entry must have an additional time dimension to be consumed. It must be placed adjacent to the last dimension of the input tensordict (i.e. at tensordict.ndim).

The action entry keys are retrieved automatically from the actor if not provided using a simple heuristic (any nested key ending with the "action" string).

An "is_init" entry must also be present in the input tensordict to track which and when the current collection should be interrupted because a “done” state has been encountered. Unlike action_keys, this key must be unique.

Parameters:

actor (TensorDictModuleBase) – An actor.
n_steps (int, optional) – the number of actions the actor outputs at once (lookahead window). Defaults to None.

Keyword Arguments:

action_keys (list of NestedKeys, optional) – the action keys from the environment. Can be retrieved from env.action_keys. Defaults to all out_keys of the actor which end with the "action" string. If chunk_keys is provided (or can be inferred from the actor’s output keys) and action_keys is omitted, the action keys are inferred from the chunk keys. For example "action_chunk" and ("vla_action", "chunk") both map to "action".
chunk_keys (list of NestedKeys, optional) – the keys written by the wrapped actor that hold action chunks. Defaults to VLA-style chunk outputs (("vla_action", "chunk") first, then keys ending in "_chunk") when present, and to action_keys otherwise. When a chunk key differs from the corresponding environment action key, the chunk key itself is used as the cache. A separate *_orig cache key is only introduced when the chunk key and the action key are the same.
init_key (NestedKey, optional) – the key of the entry indicating when the environment has gone through a reset. Defaults to "is_init" which is the out_key from the InitTracker transform.
keep_dim (bool, optional) – whether to keep the time dimension of the macro during indexing. Defaults to False.
replan_interval (int, optional) – re-query the wrapped actor after this many actions have been consumed from the cache (receding-horizon execution; the actor call is skipped in between, which is the point of action chunking for expensive policies such as VLAs). Must be in [1, n_steps]; replan_interval=1 re-plans at every step (closed loop). Defaults to None, i.e. the whole cache is consumed before re-querying (open loop). With n_steps=None the bound is enforced at execution time against the actual chunk length instead.

Examples

>>> import torch.nn
>>> from torchrl.modules.tensordict_module.actors import MultiStepActorWrapper, Actor
>>> from torchrl.envs import CatFrames, GymEnv, TransformedEnv, SerialEnv, InitTracker, Compose
>>> from tensordict.nn import TensorDictSequential as Seq, TensorDictModule as Mod
>>>
>>> time_steps = 6
>>> n_obs = 4
>>> n_action = 2
>>> batch = 5
>>>
>>> # Transforms a CatFrames in a stack of frames
>>> def reshape_cat(data: torch.Tensor):
...     return data.unflatten(-1, (time_steps, n_obs))
>>> # an actor that reads `time_steps` frames and outputs one action per frame
>>> # (actions are conditioned on the observation of `time_steps` in the past)
>>> actor_base = Seq(
...     Mod(reshape_cat, in_keys=["obs_cat"], out_keys=["obs_cat_reshape"]),
...     Mod(torch.nn.Linear(n_obs, n_action), in_keys=["obs_cat_reshape"], out_keys=["action"])
... )
>>> # Wrap the actor to dispatch the actions
>>> actor = MultiStepActorWrapper(actor_base, n_steps=time_steps)
>>>
>>> env = TransformedEnv(
...     SerialEnv(batch, lambda: GymEnv("CartPole-v1")),
...     Compose(
...         InitTracker(),
...         CatFrames(N=time_steps, in_keys=["observation"], out_keys=["obs_cat"], dim=-1)
...     )
... )
>>>
>>> print(env.rollout(100, policy=actor, break_when_any_done=False))
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([5, 100, 2]), device=cpu, dtype=torch.float32, is_shared=False),
        action_orig: Tensor(shape=torch.Size([5, 100, 6, 2]), device=cpu, dtype=torch.float32, is_shared=False),
        counter: Tensor(shape=torch.Size([5, 100, 1]), device=cpu, dtype=torch.int32, is_shared=False),
        done: Tensor(shape=torch.Size([5, 100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        is_init: Tensor(shape=torch.Size([5, 100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([5, 100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                is_init: Tensor(shape=torch.Size([5, 100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                obs_cat: Tensor(shape=torch.Size([5, 100, 24]), device=cpu, dtype=torch.float32, is_shared=False),
                observation: Tensor(shape=torch.Size([5, 100, 4]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([5, 100, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([5, 100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([5, 100, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([5, 100]),
            device=cpu,
            is_shared=False),
        obs_cat: Tensor(shape=torch.Size([5, 100, 24]), device=cpu, dtype=torch.float32, is_shared=False),
        observation: Tensor(shape=torch.Size([5, 100, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([5, 100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([5, 100, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([5, 100]),
    device=cpu,
    is_shared=False)

MultiStepActorWrapper#

Docs

Tutorials

Resources