reference/modules_actors

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

Actor Modules

Actor modules represent policies in RL. They map observations to actions, either deterministically or stochastically.

TensorDictModules and SafeModules

Actor(*args, **kwargs)	General class for deterministic actors in RL.
DiffusionActor(*args, **kwargs)	Diffusion-based actor for RL.
MultiStepActorWrapper(*args, **kwargs)	A wrapper around a multi-action actor.
SafeModule(*args, **kwargs)	tensordict.nn.TensorDictModule subclass that accepts a TensorSpec as argument to control the output domain.
SafeSequential(*args, **kwargs)	A safe sequence of TensorDictModules.
TanhModule(*args, **kwargs)	A Tanh module for deterministic policies with bounded action space.
RandomPolicy([action_spec, action_key])	A random policy for data collectors.

Probabilistic actors

ProbabilisticActor(*args, **kwargs)	General class for probabilistic actors in RL.
DreamerV3DiscreteActor(*args, **kwargs)	DreamerV3 one-hot categorical policy over stochastic state and belief.
DreamerV3SeededPolicy(*args, **kwargs)	Run a DreamerV3 policy with an independent, checkpointable random stream.
SafeProbabilisticModule(*args, **kwargs)	tensordict.nn.ProbabilisticTensorDictModule subclass that accepts a TensorSpec as an argument to control the output domain.
SafeProbabilisticTensorDictSequential(*args, ...)	tensordict.nn.ProbabilisticTensorDictSequential subclass that accepts a TensorSpec as argument to control the output domain.

Q-Value actors

QValueActor(*args, **kwargs)	A Q-Value actor class.
DistributionalQValueActor(*args, **kwargs)	A Distributional DQN actor class.
QValueModule(*args, **kwargs)	Q-Value TensorDictModule for Q-value policies.
DistributionalQValueModule(*args, **kwargs)	Distributional Q-Value hook for Q-value policies.

Flow policies

FlowMatchingPolicy(*args, **kwargs)	TensorDict policy that samples actions by Euler integration.
OneStepPolicy(*args, **kwargs)	TensorDict policy distilled from a flow policy.

These actors read observation and optional noise keys and write action. They sample Gaussian noise when it is absent and clip actions to [low, high] (default [-1, 1]). Bounds broadcast over actions; keys can be customized, including nested keys. Networks receive concatenated inputs:

import torch
from tensordict import TensorDict
from torchrl.modules import FlowMatchingPolicy, MLP, OneStepPolicy

flow = FlowMatchingPolicy(MLP(6, 2, num_cells=[64, 64]), action_dim=2)
student = OneStepPolicy(MLP(5, 2, num_cells=[64, 64]), action_dim=2)
td = TensorDict(
    observation=torch.randn(4, 3), noise=torch.randn(4, 2), batch_size=[4]
)
teacher_action = flow(td.clone())["action"]
student_action = student(td)["action"]

The example uses three observation and two action coordinates. Flow integration clips only its final output. Both policies can be passed directly to collectors and env.rollout. Their tensor-only models are available as policy.module; student.module(observation, noise, clamp=False) exposes raw outputs for distillation. Environment actions should use the default clipping.