rl

TorchRL

TorchRL is a PyTorch-native toolkit for reinforcement learning, decision making, robotics, and simulation. It is not a single algorithm implementation or a narrow benchmark suite: it is a collection of composable pieces for building RL systems while keeping the code close to the PyTorch programming model. Recent work has made this especially strong for recurrent RL, MuJoCo-based control, multi-agent training, replay-buffer and collector infrastructure, and reusable loss/value-estimation components.

The library is built around three ideas:

1. Data should have names, structure, batch dimensions, and devices all the way through the training loop.
2. Environments, policies, replay buffers, objectives, and collectors should be independent modules that can be swapped without rewriting the rest of the stack.
3. Research code should scale from a local prototype to vectorized, multiprocess, distributed, compiled, recurrent, multi-agent, model-based, or offline workflows without changing the data model.

That common data model is TensorDict, a dictionary-like tensor container with PyTorch operations, device transfers, shared-memory support, memmaps, lazy views, and nn.Module wrappers.

Getting started | API reference | Tutorials | Knowledge base | Examples | SOTA implementations

What’s new in TorchRL 0.15

TorchRL 0.15 introduces powerful new capabilities for robot learning, offline RL, data interoperability, and training at scale.

• Robotics and deployment. Use MuJoCo Menagerie and MicroDuck environments, compose low-level controllers, and export TensorDict policies to ONNX. Menagerie example · Low-level control tutorial
• Model-based and offline RL. TD-MPC2, Flow Q-Learning, and Transitive RL provide components and training implementations for control and offline RL. TD-MPC2 · FQL · Transitive RL
• PyTorch and dataset interoperability. Use buffer.as_dataset() with a PyTorch DataLoader, or buffer.export(...) to exchange trajectory data with LeRobot and Hugging Face Datasets. PyTorch data integration · LeRobot tutorial
• Data collection. Use replay_write_mode="trajectory" to write complete episodes to replay, with new readiness and policy-version controls. Collectors and execution models
• Evaluation and checkpointing. Evaluate across environments, schedule evaluation during training, and resume supported recipes from checkpoints. Evaluation · Checkpointing
• Sequence models and RLHF. Build Transformer-based policies and train reward models from preference data. Transformer example · Reward-model recipe

A quick mental model

TorchRL represents an RL interaction as a TensorDict that moves through a small number of reusable components:

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

The same object can carry observations, pixels, actions, rewards, masks, recurrent states, agent groups, sampled indices, priorities, or custom task fields. The result is less glue code and fewer hidden assumptions about what each algorithm or environment returns.

Quick demo

A local rollout is just a TensorDict passed between a PyTorch module and an environment:

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

from torchrl.envs import PendulumEnv, StepCounter, TransformedEnv

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

# Policies are regular nn.Modules wrapped with explicit TensorDict keys.
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])

Nothing in this pattern is specific to Pendulum. The same keys-and-TensorDict interface is used by batched environments, multi-agent tasks, collectors, replay buffers, recurrent modules, transforms, and losses.

What TorchRL is today

TensorDict-first pipelines

RL code tends to accumulate special cases: tuples from one environment, dicts from another, separate arrays for recurrent states, masks next to data rather than inside it, and losses that silently assume a particular batch layout. TorchRL uses TensorDict to make those assumptions explicit.

TensorDict supports common tensor operations while preserving named fields:

# These operations preserve the 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"]

This is the reason TorchRL components compose: a collector can emit a TensorDict, a replay buffer can store it without losing structure, a transform can add or remove keys, and a loss can read exactly the keys it needs.

Environments and transforms

TorchRL includes native environments, wrappers for popular environment libraries, and vectorized containers for running many environments at once. The environment API exposes specs for observations, actions, rewards, and done flags, so policies and transforms can check shapes, devices, dtypes, and bounds before a training job runs for hours.

Environment support includes:

• PyTorch-native environments such as PendulumEnv, and custom MuJoCo tasks built on MujocoEnv.
• Wrappers for Gymnasium, Gym, DM Control, Brax, Jumanji, PettingZoo, VMAS, OpenSpiel, Safety-Gymnasium, Isaac Lab, and other optional libraries.
• SerialEnv, ParallelEnv, and batched wrappers for local vectorization and multiprocessing.
• Environment transforms for observation normalization, image conversion, reward transforms, action masking, action scaling (ActionScaling), action flattening (FlattenAction), auto-reset, frame stacking, state reconstruction, and more.

Transforms are first-class TorchRL modules. They can run on-device, participate in specs, and be inserted, removed, or composed without wrapping the whole environment in opaque adapter layers. Macro-control primitives expand high-level commands into sequences of low-level actions.

from torchrl.envs import Compose, DoubleToFloat, ObservationNorm, TransformedEnv
from torchrl.envs.libs.gym import GymEnv

base_env = GymEnv("HalfCheetah-v4", device="cuda:0")
env = TransformedEnv(
    base_env,
    Compose(
        ObservationNorm(in_keys=["observation"]),
        DoubleToFloat(),
    ),
)

Collectors and execution models

Collectors are the bridge between policies and environments. A collector owns the execution loop, batches trajectories, handles devices, and can update policy weights while environments keep running.

TorchRL includes single-process, async, multiprocess, and distributed collectors. This lets the same policy and loss code be used across small smoke tests, GPU-heavy simulation, CPU environment farms, or asynchronous evaluation setups.

from torchrl.collectors import Collector

collector = Collector(
    create_env_fn=env,
    policy=policy,
    frames_per_batch=1024,
    total_frames=1_000_000,
)

for data in collector:
    # data is a TensorDict with time, environment, and key structure preserved.
    train_step(data)

For larger jobs, the collector family adds async execution, multiple worker processes, weight updaters, evaluator loops, profiling hooks, and fake-data helpers for testing downstream code without stepping an expensive environment.

TorchRL 0.14 made it easier to scale data collection across local processes and distributed workers through a shared Collector API. For example, if make_env creates an environment, use multiple local processes with:

multi_collector = Collector(
    make_env, policy, num_collectors=4, sync=True, frames_per_batch=1024
)

Here sync=True waits for all four workers. This call returns a MultiSyncCollector, the implementation that manages synchronous worker processes.

The same Collector API supports several execution modes:

• Single process: Collector(make_env, policy, ...) collects in the current process.
• Synchronous workers: num_collectors=4, sync=True collects across four local workers and waits for each when forming a batch.
• Asynchronous workers: num_collectors=4, sync=False receives data as workers finish.
• Distributed workers: backend="ray", "rpc", or "distributed" selects a distributed execution backend.

For asynchronous environment stepping and batched policy inference, AsyncBatchedCollector combines an AsyncEnvPool, which runs environments independently, with an inference server that groups their policy requests:

from torchrl.collectors import AsyncBatchedCollector

collector = AsyncBatchedCollector(
    [make_env] * 4, policy=policy, frames_per_batch=1024, env_backend="threading"
)

The async collector can also write to replay in the background. In 0.15, replay_write_mode="trajectory" writes complete episodes to replay, buffer.sample(wait=True) waits for sampleable data on supported buffers, and multi_collector.worker_policy_versions() reports each worker’s policy version when track_policy_version=True.

See collector APIs, distributed collectors, async environments, and collectors and replay.

Replay buffers and offline data

TorchRL replay buffers are modular: storage, sampler, writer, collate function, transforms, prefetching, priority updates, and device movement are separate pieces. That makes it possible to use the same interface for simple in-memory replay, memmap-backed storage, prioritized replay, CUDA-aware sampling, offline datasets, or custom storage layouts. HERReplayBuffer relabels goals for goal-conditioned RL. For compact observations, NextObservationDelta and NextStateReconstructor support reconstruction when sampling from replay.

from torchrl.data import LazyMemmapStorage, TensorDictPrioritizedReplayBuffer

buffer = TensorDictPrioritizedReplayBuffer(
    storage=LazyMemmapStorage(1_000_000),
    alpha=0.7,
    beta=0.5,
    batch_size=256,
    prefetch=2,
)

buffer.extend(collector_batch)
sample = buffer.sample()

Replay buffers understand TensorDict structure, so they can store trajectories, nested agent data, recurrent states, HER relabeling metadata, or offline datasets without flattening everything into parallel Python containers.

Modules, distributions, and policies

TorchRL modules are ordinary PyTorch modules with explicit input and output keys. The library provides actors, critics, actor-critic operators, recurrent modules, distribution wrappers, exploration modules, world models, decision transformers, robot-learning models, and helper utilities for inferring specs from environments.

A stochastic actor can be assembled from familiar PyTorch layers:

from tensordict.nn import TensorDictModule
from tensordict.nn.distributions import NormalParamExtractor
from torch import nn
from torchrl.modules import ProbabilisticActor, TanhNormal

params = TensorDictModule(
    nn.Sequential(
        nn.LazyLinear(256),
        nn.Tanh(),
        nn.Linear(256, 2),
        NormalParamExtractor(),
    ),
    in_keys=["observation"],
    out_keys=["loc", "scale"],
)

actor = ProbabilisticActor(
    params,
    in_keys=["loc", "scale"],
    out_keys=["action"],
    distribution_class=TanhNormal,
    distribution_kwargs={"low": -1.0, "high": 1.0},
    return_log_prob=True,
)

The explicit key contract makes it clear what data a module consumes and produces, and it allows losses, collectors, and transforms to be reconfigured without editing the model itself.

Objectives, returns, and trainers

TorchRL objectives are loss modules that read TensorDict keys, compute losses, and expose configurable key mappings. They cover policy-gradient methods, actor-critic algorithms, Q-learning, offline RL, imitation learning, model-based RL, and multi-agent RL.

Examples include PPO, SAC, DQN, TD3, REDQ, IQL, CQL, Decision Transformer, Dreamer, CrossQ, GAIL, behavior cloning, ACT, MAPPO, IPPO, and QMIX/VDN. Value-estimator utilities provide GAE, TD(lambda), V-trace, lambda returns, multi-agent advantages, and vectorized return computation.

from torchrl.objectives import ClipPPOLoss
from torchrl.objectives.value import GAE

loss = ClipPPOLoss(actor_network=actor, critic_network=critic)
advantage = GAE(value_network=critic, gamma=0.99, lmbda=0.95)

data = advantage(data)
losses = loss(data)
loss_value = losses["loss_objective"] + losses["loss_critic"] + losses["loss_entropy"]

For higher-level workflows, TorchRL also provides trainer utilities and Hydra configuration dataclasses that assemble environments, networks, collectors, losses, optimizers, loggers, hooks, and schedules into reproducible recipes.

Multi-agent, model-based, and imitation learning

Multi-agent data is represented as TensorDict structure rather than a separate parallel convention. Agent observations, actions, rewards, masks, and shared state can live under nested keys such as ("agents", "observation"), while losses and modules declare which keys they use.

TorchRL supports multi-agent environments and algorithms through VMAS, PettingZoo, Melting Pot, SMACv2, OpenSpiel, multi-agent trainers, and dedicated objectives including MAPPO and IPPO, as well as MultiAgentGAE, ValueNorm, PopArtValueNorm, RunningValueNorm, and cross-agent critic utilities.

The same component style also covers model-based and imitation-learning work: Dreamer/DreamerV3 objectives and RSSM modules, Decision Transformer components, BCLoss for behavior cloning, and ACTModel with ACTLoss for action chunking all share the same TensorDict and key-dispatch conventions as the online RL algorithms.

Additional specialized workflows

TorchRL also includes support for specialized workflows, including LLM post-training experiments. The LLM stack provides conversation containers, Hugging Face/vLLM/SGLang integration points, GRPO and SFT objectives, async collectors, weight-update helpers, and tool-use transforms. Entry points include the LLM reference and the GRPO implementation.

Performance and PyTorch integration

TorchRL is designed to stay close to PyTorch execution. Components are TensorDict-aware, vectorized where possible, and increasingly friendly to torch.compile, CUDA, shared memory, memmaps, and distributed execution.

Performance-sensitive areas include:

• vectorized return and advantage computation;
• recurrent GRU/LSTM reset handling with scan and Triton backends;
• compact sequence layouts for recurrent value estimation, shifted value estimation, and chunked value-network forwards;
• async collectors and policy weight synchronization;
• prioritized replay and CUDA-aware replay-buffer paths;
• memmap-backed data movement for large offline or distributed jobs.

Where to start

If you want to…	Start with…
Learn the basic environment and TensorDict loop	Getting started and the quick demo above
Train a classic continuous-control agent	PPO, SAC, or TD3 implementations
Build custom environment preprocessing	Environment transforms
Scale data collection	Collectors and distributed collectors
Store large or prioritized data	Replay buffers
Work with recurrent policies	Recurrent modules and state lifecycle docs
Train multi-agent systems	Multi-agent objectives and multi-agent examples
Explore MuJoCo macro policies	Macro primitives and MuJoCo tutorials
Try language-model post-training experiments	LLM reference and GRPO

Installation

TorchRL targets Python 3.10+, PyTorch 2.11+, and TensorDict 0.15.x (0.15.2 or later). PyTorch 2.15 requires Python 3.11+, so the Python 3.10 wheels are built against PyTorch 2.14.

Install the stable release:

pip install torchrl

This standard PyPI wheel is the right default for most users, including CPU prioritized replay buffers and workloads that do not use prioritized replay. Starting with TorchRL 0.13, Linux CUDA wheels are also published for users who want the CUDA-based prioritized replay-buffer implementations. Install the CUDA wheel from the PyTorch wheel index that matches your PyTorch CUDA runtime (replace cu132 with the CUDA build you use):

pip install "torchrl==0.15.0+cu132" --extra-index-url https://download.pytorch.org/whl/cu132

The CUDA wheel is optional: if you do not need CUDA prioritized replay buffers, or if your prioritized replay buffers run on CPU, keep using pip install torchrl.

Install common optional dependencies:

pip install "torchrl[utils]"              # Hydra, logging, and development utilities
pip install "torchrl[gym_continuous]"     # Gymnasium continuous-control environments
pip install "torchrl[atari]"              # Atari support
pip install "torchrl[offline-data]"       # Offline datasets and data helpers
pip install "torchrl[marl]"               # Multi-agent environment libraries
pip install "torchrl[llm-vllm]"           # LLM API with vLLM backend on Linux
pip install "torchrl[llm-sglang]"         # LLM API with SGLang backend on Linux

Some optional libraries are platform- or Python-version-specific. If you are building a reproducible environment, install PyTorch first from the appropriate PyTorch installation selector, then install TorchRL and the optional extras you need.

Install the nightly builds when working against nightly PyTorch:

pip install --pre tensordict-nightly torchrl-nightly

For local development, keep the TorchRL and TensorDict checkouts on compatible branches and avoid re-resolving an already selected PyTorch build:

git clone https://github.com/pytorch/tensordict
git clone https://github.com/pytorch/rl
uv pip install --no-deps -e tensordict
uv pip install --no-deps -e rl

The C++ extension paths used by prioritized replay buffers require a compatible PyTorch version. If you see undefined-symbol errors, consult the versioning issues guide.

Documentation and learning resources

• Stable documentation
• API reference
• Tutorials
• Knowledge base
• Benchmarks
• Nightly CI status

Introductory material:

• TorchRL paper
• TalkRL podcast episode
• TorchRL intro at PyTorch Day 2022
• PyTorch 2.0 Q&A: TorchRL

Examples, tutorials, and implementations

TorchRL ships examples for small features and complete training recipes:

• SOTA implementations for PPO, SAC, DQN, TD3, REDQ, Decision Transformer, Dreamer, CrossQ, GAIL, IMPALA, multi-agent algorithms, GRPO, and more.
• Examples for distributed collectors, replay buffers, RLHF, MuJoCo satellite control, and other focused workflows.
• Tutorials for environment design, transforms, collectors, losses, recurrent state handling, MuJoCo macros, and end-to-end training.

The implementations are meant to be readable starting points, not black-box benchmarks. They show how TorchRL components fit together and can be copied into research code when a full trainer abstraction is not the right fit.

Ecosystem and publications

TorchRL is domain-agnostic and is used across robotics, control, simulation, drug discovery, multi-agent RL, combinatorial optimization, and research infrastructure. Selected projects and papers include:

• ACEGEN: Reinforcement learning of generative chemical agents for drug discovery.
• BenchMARL: Benchmarking multi-agent reinforcement learning.
• BricksRL: A platform for democratizing robotics and reinforcement learning research and education with LEGO.
• OmniDrones: An efficient and flexible platform for reinforcement learning in drone control.
• RL4CO: Reinforcement learning for combinatorial optimization.
• Robohive: A unified framework for robot learning.

Citation

If you use TorchRL, please cite:

@misc{bou2023torchrl,
      title={TorchRL: A data-driven decision-making library for PyTorch},
      author={Albert Bou and Matteo Bettini and Sebastian Dittert and Vikash Kumar and Shagun Sodhani and Xiaomeng Yang and Gianni De Fabritiis and Vincent Moens},
      year={2023},
      eprint={2306.00577},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

Asking questions

If you find a bug, please open an issue in this repository. For broader RL in PyTorch questions, use the PyTorch reinforcement learning forum.

Contributing

Contributions are welcome. See CONTRIBUTING.md for the full contribution guide and the call for contributions for open areas where help is especially useful.

To work on an open issue, comment /assign on the issue. This will assign the issue to you.

Trusted collaborators

TorchRL is maintained by Vincent Moens (vmoens), Achintya Paningapalli (theap06), and Ben Sprenger (bsprenger). The following trusted collaborators are also available for review guidance and contributor advice:

• coder-jayp

If you would like to be added to this list, reach out to Vincent Moens (vmoens).

For local development, install pre-commit hooks with:

pre-commit install

Status and license

TorchRL is released as a PyTorch beta feature. Breaking changes can happen, but TorchRL aims to introduce them with deprecation warnings over multiple release cycles.

TorchRL is licensed under the MIT License. See LICENSE for details.