TorchRL: Modular Reinforcement Learning for PyTorch

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/pytorch/rl/llms.txt

Use this file to discover all available pages before exploring further.

​

The TensorDict Data Model

TensorDict
  -> policy module writes actions and log-probs
  -> environment reads actions and writes next observations, rewards, done flags
  -> collector batches trajectories from one or many workers
  -> replay buffer stores, samples, prioritizes, and transforms data
  -> loss module reads named keys and writes differentiable losses
  -> optimizer updates ordinary PyTorch parameters

# These operations preserve structure and operate on every compatible value.
batch = torch.stack(list_of_tensordicts, dim=0)
batch = batch.reshape(-1)
batch = batch.to("cuda")
mini_batch = batch[:128]

# Nested keys make multi-agent, recurrent, and next-state data explicit.
reward = batch["next", "reward"]
agent_obs = batch["agents", "observation"]
hidden = batch["recurrent_state", "h"]

​

Library Structure

​

Key Features

• TensorDict-first pipelines — every component reads and writes named tensor fields; no silent shape assumptions or hidden data conventions
• Broad environment coverage — PyTorch-native environments plus wrappers for Gymnasium, DM Control, Brax, Jumanji, PettingZoo, VMAS, OpenSpiel, Safety-Gymnasium, Isaac Lab, MuJoCo, and more
• Composable transforms — first-class transform modules that run on-device, participate in specs, and can be inserted or removed without wrapping the entire environment
• Flexible collectors — Collector, MultiSyncCollector, MultiAsyncCollector, and distributed variants for everything from smoke tests to GPU-heavy simulation farms
• Modular replay buffers — separate storage, sampler, writer, and transform pieces; supports prioritized replay, HER, memmap-backed storage, and offline datasets
• Rich policy modules — ProbabilisticActor, TanhNormal, ValueOperator, recurrent LSTM/GRU wrappers, and exploration modules with explicit in/out key contracts
• Complete objectives library — PPO, SAC, DQN, TD3, REDQ, IQL, CQL, Dreamer/DreamerV3, Decision Transformer, behavior cloning, and multi-agent objectives
• Multi-agent support — MAPPO, IPPO, MultiAgentGAE, value normalization, VMAS and PettingZoo environments, and nested TensorDict conventions for agent groups
• LLM post-training — GRPO and SFT objectives, conversation containers, Hugging Face/vLLM/SGLang integration, async collectors, and weight-update helpers
• SOTA implementations — readable, research-ready implementations of PPO, SAC, DQN, TD3, REDQ, IMPALA, CrossQ, GAIL, Dreamer, Decision Transformer, GRPO, and multi-agent algorithms
• PyTorch-native performance — torch.compile compatibility, CUDA-aware replay, Triton/scan recurrent backends, vectorized return computation, and distributed execution

import torch
from tensordict.nn import TensorDictModule
from torch import nn

from torchrl.envs import PendulumEnv, StepCounter, TransformedEnv

# A PyTorch-native environment with a transform stack.
env = TransformedEnv(PendulumEnv(), StepCounter(max_steps=200))

# Policies are regular nn.Modules with explicit TensorDict key contracts.
policy = TensorDictModule(
    nn.Sequential(
        nn.LazyLinear(64),
        nn.Tanh(),
        nn.Linear(64, 1),
        nn.Tanh(),
    ),
    in_keys=["observation"],
    out_keys=["action"],
)

rollout = env.rollout(max_steps=32, policy=policy)
assert rollout.batch_size == torch.Size([32])
assert rollout["next", "reward"].shape[:1] == torch.Size([32])

​

Where to Start