pytorch/rl

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

What it solves

TorchRL provides a modular, PyTorch-native toolkit for building reinforcement learning (RL) systems. It eliminates the need for excessive "glue code" and rigid data structures by using a common data model that allows environments, policies, replay buffers, and loss functions to be swapped and composed without rewriting the entire training stack.

How it works

The library is centered around TensorDict, a dictionary-like container that carries named tensors (observations, actions, rewards, etc.) through the training loop. This allows data to maintain its structure, batch dimensions, and device location across different components:

  • Policies: PyTorch modules that read from and write to specific keys in a TensorDict.
  • Environments: Modules that handle the interaction loop, supporting native PyTorch environments and wrappers for libraries like Gymnasium and MuJoCo.
  • Collectors: Bridges that manage the execution loop, batching trajectories from one or many workers (single-process, async, or distributed).
  • Replay Buffers: Modular storage systems for sampling and transforming data, supporting prioritized replay and memmap-backed storage.
  • Loss Modules: Objectives that compute differentiable losses (e.g., PPO, SAC, DQN) based on named keys in the TensorDict.

Who it’s for

TorchRL is designed for researchers and developers working on reinforcement learning, robotics, decision-making, and simulation, as well as those performing LLM post-training experiments (such as GRPO).

Highlights

  • Composable Architecture: Independent modules for environments, policies, and buffers that can be swapped easily.
  • Broad Algorithm Support: Implementations for PPO, SAC, TD3, DQN, DreamerV3, and multi-agent algorithms like MAPPO and IPPO.
  • High Performance: Optimized for torch.compile, CUDA, and distributed execution, with specialized Triton backends for recurrent RL.
  • Extensive Environment Integration: Native support for MuJoCo and wrappers for Gymnasium, PettingZoo, and Isaac Lab.
  • LLM Integration: Tools for LLM post-training, including SFT and GRPO objectives and vLLM/SGLang integration.

Related

  • Project
    meta-pytorch/torchforge

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

  • Project
    tensorforce/tensorforce

    A modular TensorFlow-based deep reinforcement learning framework for research and practice, providing flexible components to build and deploy RL agents.

  • Project
    yihaosun1124/OfflineRL-Kit

    A PyTorch library for offline reinforcement learning that provides a standardized framework and implementations of state-of-the-art model-free and model-based algorithms.

  • Project
    TorchDSP/torchsig

    An open-source signal processing machine learning toolkit based on PyTorch that simplifies the creation, augmentation, and transformation of RF signals for ML research.

  • Project
    ai4co/rl4co

    A unified reinforcement learning framework and benchmark for combinatorial optimization, designed to decouple algorithm research from engineering implementation.