reference/generated/torchrl.modules.LowLevelController

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

LowLevelController

class torchrl.modules.LowLevelController(*args, **kwargs)

Deploy a TensorDict policy on independent controller instances.

The adapter and policy run on a private, flattened TensorDict. Only the low-level action and the namespaced next state are written back. Policy weights are shared; recurrent and adapter state remain independent for every environment and every member of the selected group.

Parameters:

• policy (TensorDictModuleBase) – pretrained policy with explicit TensorDict inputs, outputs, and primers for persistent state.

• decision_spec (Composite) – unbatched spec for one controller’s high-level decisions. All leaves are held fixed during a closed-loop macro step.

Keyword Arguments:

• adapter (TensorDictModuleBase or Transform, optional) – module applied before the policy, using keys relative to the group. Defaults to None (pass inputs directly to the policy).

• group_key (NestedKey, optional) – group whose batch dimensions enumerate controller instances. Defaults to None (the root TensorDict).

• state_key (NestedKey, optional) – state namespace within the group. Defaults to “controller”. The non-private namespace is retained by collectors and replay buffers.

• policy_action_key (NestedKey, optional) – action output of the wrapped policy. Defaults to “action”.

• action_key (NestedKey, optional) – destination action within the group. Defaults to “action”.

• reset_key (NestedKey, optional) – additional per-instance reset signal in the group’s next observation, such as “fallen”. Defaults to None (ordinary environment resets only).

The module does not disable gradients or change the policy’s training mode. ClosedLoopMultiAction controls inference when the controller is deployed in an environment.

Examples

>>> import torch
>>> from tensordict.nn import TensorDictModule
>>> from torchrl.data import Bounded, Composite
>>> from torchrl.envs.transforms import ClosedLoopMultiAction
>>> 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=2)
>>> td = env.reset().set("command", torch.ones(1))
>>> env.step(td)["next", "observation"]
tensor([2], dtype=torch.int32)
>>> env.close()

See also

ClosedLoopMultiAction repeatedly executes the controller against fresh observations; MicroDuckSkillController specializes the controller for a task-conditioned MicroDuck policy; and LowLevelControllerConfig exposes this class through Hydra configuration.

forward(tensordict: TensorDictBase) → TensorDictBase

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

make_tensordict_primer() → Transform

Return the group’s state initialization and partial-reset transforms.

Returns:

primers that infer group shapes from the parent env and retain the original policy and adapter initialization values.

Return type:

Transform