DOCSEARCHbyAlgolia
Submit Docs

TorchRL Documentation

repository·main·Indexed 25 days ago

https://github.com/pytorch/rl

A modular, primitive-first, PyTorch-native toolkit for reinforcement learning, decision making, robotics, and simulation. TorchRL uses TensorDict as its core data model to ensure composability and scalability across local and distributed workflows. It provides independent modules for environments, policies, replay buffers, objectives, and collectors, supporting a wide range of libraries including Gymnasium, Gym, and MuJoCo.

Tokens
103.4K
Snippets
198
Records
422
Agent score
86%

What's inside TorchRL

  1. Overview of TorchRL

    main

    TorchRL is a PyTorch-native toolkit designed for reinforcement learning, decision making, robotics, and simulation. It provides a collection of composable modules that follow the PyTorch programming model, emphasizing a TensorDict-first approach.

    Key design principles include:

    • Structured Data: Data maintains names, structure, batch dimensions, and device locality throughout the training loop using TensorDict.
    • Modularity: Environments, policies, replay buffers, objectives, and collectors are independent, swappable modules.
    • Scalability: Research code can scale from local prototypes to distributed, vectorized, or multi-agent workflows without changing the underlying data model.

  4. Understand the TorchRL Configuration System

    main
    TorchRL uses a configuration system built on top of Hydra to manage reinforcement learning experiments. It utilizes structured, dataclass-based configurations that allow for composition, overriding, and extension. This system enables reproducible experiments by using YAML files and command-line overrides to define environments, models, collectors, and training parameters.

  5. Use torchrl.modules for RL neural networks

    main
    The torchrl.modules package provides a collection of RL-specific neural network modules built on top of tensordict.nn.TensorDictModule. These modules are designed to work with tensordict data structures and support features like spec-based construction, probabilistic policies, exploration strategies, value networks, and safe modules for action constraints.

  6. Use Monte Carlo Tree Search (MCTS) components from torchrl.modules.mcts

    main
    The torchrl.modules.mcts package provides components for Monte Carlo Tree Search (MCTS), specifically focusing on score computation modules. These modules are used to balance exploration and exploitation during tree search algorithms by calculating scores for different nodes or actions.

  8. Use the torchrl.envs package for environment management

    main

    The torchrl.envs package provides a unified API to handle environments from various backends (Gym, DMControl, Brax, Jumanji, etc.). It uses tensordict.TensorDict as the foundation for data organization, allowing for arbitrary inputs, outputs, and nested or batched data structures.

    Key capabilities include:

    • Unified API: Consistent interface across different backends.
    • Vectorization: Built-in support for parallel and batched environments.
    • Transforms: A system for preprocessing observations and actions.
    • Multi-agent: Native support for multi-agent RL.
    • Flexible backends: Easy integration with popular RL frameworks.

  10. Understand TorchRL Environment (Env) and Transforms

    main

    In TorchRL, an Env implements the torchrl.envs.EnvBase API (including reset, step, and specs) and uses a tensordict-based input/output contract.

    Instead of using Gym wrappers, TorchRL uses Transforms (torchrl.envs.transforms.Transform). Transforms are used to modify input/output specs, reset data, step data, or inverse action data, typically by composing them into a torchrl.envs.transforms.TransformedEnv.

  11. Understand TorchRL service deployment profiles

    main

    TorchRL service examples demonstrate different ways to place inference, logging, and replay buffer components. The training loop remains backend-neutral by consuming domain-compatible clients regardless of the placement.

    | Profile | Inference | Logger | Replay buffer | Training loop |
| --- | --- | --- | --- | --- |
| Single process | Background thread | Direct | Direct | Driver |
| Multiprocess | Spawned process | Spawned process | Direct | Driver |
| Ray | Ray actor | Ray actor | Ray actor | Driver |

  12. Compare ParallelEnv, SerialEnv, and MultiThreadedEnv

    main

    TorchRL provides different ways to handle multiple environments depending on your needs:

    1. ParallelEnv: Uses process-based parallelism. It is highly flexible and can wrap any arbitrary TorchRL environment, but has higher overhead due to spawning Python processes.
    2. SerialEnv: Executes environments serially in a single process. This is primarily used for testing and debugging the behavior of batched environments without the complexity of subprocesses.
    3. MultiThreadedEnv: Uses the EnvPool library for multithreaded execution. It offers higher performance but is less flexible: you can only use environments implemented in EnvPool (e.g., certain Atari or Classic Control environments).