Research map / Reinforcement learning / Visual and representation RL
TorchRL: A data-driven decision-making library for PyTorch
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Presents TorchRL, a modular PyTorch reinforcement learning and control library built around the TensorDict data carrier, with reproduction, distributed, offline, multi-agent and efficiency benchmarks.
Problem
PyTorch lacks a native, comprehensive library for decision-making and control that is both well-integrated and flexible enough for large teams and diverse RL tasks, and existing libraries trade off modularity against ease of use.
Approach
Introduces TensorDict, a dictionary-like, batch-size-aware PyTorch data carrier used as the common interface between independent components (environments, collectors, replay buffers, modules, losses), so that components can be swapped and reused without coupling.
Key results
- Reproduction results over 5 seeds on MuJoCo (e.g., SAC HalfCheetah 11077±323, TD3 Hopper 2809±618) and Atari (e.g., PPO Pong 20.52±0.58).
- Vectorized VMAS collection reaches about 538,930 fps with 20,000 envs versus about 1,581 fps for RLlib.
- GAE computation: TorchRL-vec 1.33 ms versus CleanRL 1.43 ms, Tianshou 5.15 ms, SB3 14.1 ms, RLlib 9.38 ms.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use TorchRL instead of RLlib for vectorized data collection | Simple spread in VMAS · Simple spread, 30000 vectorized environments | frames per second: 422003 vs 1336 | Helped |
| Use TorchRL-vec instead of CleanRL for GAE computation | GAE implementations · 1000 trajectories of 1000 time steps | speed (ms): 1.43 vs 1.33 | Helped |
| Use TorchRL async collector instead of TorchRL sync collector | TorchRL data collection · Breakout-v5 | frames per second: 10394 -> 19401 | Helped |
| Use TorchRL async collector instead of TorchRL sync collector | TorchRL data collection · HalfCheetah-v4 | frames per second: 23527 -> 32894 | Helped |
| Use TorchRL async collector instead of TorchRL sync collector | TorchRL data collection · Pendulum-v1 | frames per second: 24530 -> 31584 | Helped |
| Use TensorDict rather than PyTree for scalar-leaf increment in a deeply nested structure | TensorDict · 100-level nested dictionary | runtime (ms): 0.916 ms vs 1.0555 ms | Hurt |
| Use TensorDict rather than PyTree for CUDA casting | TensorDict · 100-level nested dictionary | runtime (ms): 1.181 ms vs 1.302 ms | Hurt |
| Use TensorDict rather than PyTree for split | TensorDict · Batch-size operation on deeply nested structure | runtime (ms): 1.015 ms against 4.254 ms | Hurt |
| Use TensorDict functional calls instead of functorch functional calls | Transformer functional calls · Narrow, deep Transformer | runtime (ms): 8.423 ms against 10.26071 ms | Hurt |
Surprising findings
- TensorDict overhead is higher than PyTree for some operations, e.g., split takes 4.254 ms versus 1.015 ms for PyTree, and value increment is about 115% slower.
- Data collection with RLlib on Breakout-v5 reaches only 97 fps versus 19401 fps for the TorchRL async collector.
Limitations
- TensorDict PyTree-style operations (split, chunk, unbind) carry significant overhead compared with PyTree equivalents, which the authors plan to address.
- Functional calls in TensorDict are in a first version and are slower than functorch in the reported test (10.26 ms vs 8.42 ms).
Experimental setup
Datasets: D4RL OpenML ImageNet
Benchmarks: MuJoCo (HalfCheetah, Hopper, Walker2D, Ant) Atari (Pong, Freeway, Boxing, Breakout) D4RL medium-replay (HalfCheetah, Hopper) VMAS simple spread / multi-robot tasks
Metrics: mean final reward std over seeds frames per second runtime ms data collection speed
Baselines: RLlib Tianshou CleanRL Stable-baselines3 Gymnasium
Builds on: PyTorch OpenAI Gym tensordict torch.vmap
Closest papers by meaning
- PWM: Policy Learning with Multi-Task World Models (ICLR 2025 · similarity 0.87)
- Gradient Boosting Reinforcement Learning (ICML 2025 · similarity 0.86)
- Efficient Reinforcement Learning by Discovering Neural Pathways (NeurIPS 2024 · similarity 0.85)
- Diffusion Guidance Is a Controllable Policy Improvement Operator (NeurIPS 2026 · similarity 0.85)
- XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning (ICLR 2025 · similarity 0.85)
- On Generalization Across Environments In Multi-Objective Reinforcement Learning (ICLR 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.