pytorch/rl

A modular, primitive-first, python-first PyTorch library for Reinforcement Learning.

TorchRL – PyTorch‑native Reinforcement‑Learning Toolkit

What it is – TorchRL is a library built on top of PyTorch that provides composable building blocks for reinforcement‑learning (RL) research and production. It does not ship a single algorithm; instead it offers a unified data container (TensorDict) and a set of interchangeable modules (environments, policies, collectors, replay buffers, loss objects, trainers, etc.) that can be mixed‑and‑matched while staying close to the native PyTorch programming model.

Core ideas

Idea How TorchRL implements it
Explicit, named data All tensors travel inside a TensorDict that carries field names, batch dimensions and device placement throughout the training loop.
Modular stack Environments, policies, collectors, replay buffers and loss modules are independent, swap‑able components.
Scalable from prototype to production The same API works for a single‑process rollout, multi‑process async collectors, distributed training, compiled (torch.compile) or CUDA‑accelerated pipelines without code changes.

Main components

  • TensorDict – dictionary‑like container with full PyTorch ops, device transfers, shared‑memory and mem‑map support.
  • Environments & Transforms – native PendulumEnv, MuJoCo tasks, and wrappers for Gymnasium, DM‑Control, Brax, PettingZoo, etc.; transforms (obs‑norm, action‑scaling, frame‑stacking, HER, etc.) are first‑class modules.
  • Collectors – synchronous, async, multiprocess and distributed collectors that batch trajectories, move data to the right device and can update policy weights on‑the‑fly.
  • Replay Buffers – modular storage (in‑memory, lazy‑memmap, CUDA‑aware), prioritized sampling, HER, offline‑dataset handling.
  • Modules / Policies – regular nn.Modules wrapped with explicit input‑/output‑key contracts; includes stochastic actors, critics, recurrent nets, distribution wrappers, world‑model components.
  • Objectives – loss modules for PPO, SAC, DQN, TD3, REDQ, IQL, CQL, Decision Transformers, DreamerV3, MAPPO/IPPO, QMIX/VDN, behavior‑cloning, etc., all reading/writing named keys.
  • Trainers & Hydra configs – high‑level utilities that wire environments, collectors, losses, optimizers and logging into reproducible recipes.

Recent highlights (v0.13)

  • Faster recurrent paths with Triton‑backed GRU/LSTM resets.
  • New MuJoCo environments, satellite examples, and macro‑control primitives.
  • Expanded multi‑agent support: MAPPO, IPPO, MultiAgentGAE, value‑norm utilities.
  • Async prioritized replay‑buffer writes, compact observation storage, optional CUDA kernels for replay.
  • Additional transforms (ActionScaling, FlattenAction, NextObservationDelta, etc.) and value‑estimator improvements.

Typical workflow (quick demo)

import torch
from tensordict.nn import TensorDictModule
from torch import nn
from torchrl.envs import PendulumEnv, StepCounter, TransformedEnv

# Environment with a simple transform stack
env = TransformedEnv(PendulumEnv(), StepCounter(max_steps=200))

# Policy expressed as a regular nn.Module 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"],
)

# One‑step rollout – the result is a TensorDict containing observations, actions, rewards, etc.
rollout = env.rollout(max_steps=32, policy=policy)
print(rollout.batch_size)          # torch.Size([32])
print(rollout["next", "reward"].shape)  # torch.Size([32])

The same TensorDict can be fed to a collector, stored in a replay buffer, and consumed by a loss module without any format conversion.

Who should use it?

  • Researchers building new RL algorithms (need flexible data flow, easy swapping of components, and tight integration with PyTorch’s autograd/compile).
  • Engineers scaling RL pipelines to many CPUs/GPUs or distributed clusters (async collectors, CUDA‑aware replay buffers, mem‑map storage).
  • Robotics / simulation teams that want native MuMuCo or custom environments with on‑device transforms.
  • Multi‑agent or model‑based RL projects (provides VMAS, PettingZoo wrappers, DreamerV3, Decision‑Transformer components).
  • LLM post‑training experiments (TorchRL includes a small LLM stack for GRPO/SFT style fine‑tuning).

Installation

# Stable release (CPU‑only replay buffers)
pip install torchrl

# CUDA‑enabled replay buffers (replace cu118 with your CUDA version)
pip install "torchrl==0.13.0+cu118" \
    --extra-index-url https://download.pytorch.org/whl/cu118

# Optional extras for extra environments / utilities
pip install "torchrl[utils]"          # Hydra, logging, etc.
pip install "torchrl[gym_continuous]" # Gymnasium continuous‑control
pip install "torchrl[atari]"          # Atari support
pip install "torchrl[marl]"           # Multi‑agent libs (PettingZoo, VMAS, …)

Documentation & resources

Bottom line: TorchRL is a full‑featured, PyTorch‑first RL engineering framework that lets you prototype, scale, and experiment with a wide range of RL, multi‑agent, model‑based, and even LLM‑fine‑tuning workflows while keeping the codebase clean and PyTorch‑native.

Related

  • Project
    LucasAlegre/morl-baselines

    MORL‑Baselines is a PyTorch library of reliable multi‑objective reinforcement‑learning algorithms (single‑ and multi‑policy, SER/ESR). It follows the MO‑Gymnasium API, provides utilities, automatic W&B logging, and integrates with the Open RL Benchmark for reproducible experiments.

  • Project
    Denys88/rl_games

    RL Games is a PyTorch‑based, high‑performance reinforcement‑learning library. It offers fast GPU‑accelerated training (PPO, SAC, multi‑agent, self‑play), integrates with EnvPool and Triton kernels, supports many simulators (MuJoCo, Isaac Gym, MJLab, Atari, StarCraft II), provides Colab notebooks and ONNX export, and includes utilities for experiment tracking and population‑based hyper‑parameter search.

  • Project
    yoshitomo-matsubara/torchdistill

    torchdistill is a PyTorch‑based library that lets you run knowledge‑distillation experiments (teacher‑student training) by writing only a YAML configuration. It provides many modern KD losses, a forward‑hook manager for extracting intermediate features, and example scripts for vision and NLP tasks. Install via `pip install torchdistill` and cite the associated papers when used in research.

  • Dispatch
    Keep the Tokens Flowing: Lessons from 16 Open-Source RL Libraries

    Hugging Face surveyed 16 open-source RL libraries and found that async RL training separates inference and training onto different GPU pools, uses a rollout buffer, and pushes weights asynchronously, with Ray dominating orchestration and NCCL broadcast the common weight‑sync method.

  • Project
    meta-pytorch/torchforge

    A PyTorch-native agentic RL library that separates infrastructure concerns from model concerns to simplify scalable reinforcement learning experimentation.