Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/pytorch/rl/llms.txt

Use this file to discover all available pages before exploring further.

TorchRL ships wrappers for more than a dozen external RL environment libraries. Every wrapper subclasses EnvBase and exposes the identical reset / step / rollout / check_env_specs interface — the same code that works with GymEnv works unchanged with DMControlEnv, BraxEnv, or VmasEnv. All third-party libraries are optional dependencies: the wrapper module is imported lazily at construction time and raises a clear ImportError if the underlying package is not installed.
Install only the third-party packages you need. Many of these libraries have platform-specific requirements (CUDA, JAX, IsaacLab GPU drivers). Check each project’s installation guide before proceeding.

​
GymEnv / GymWrapper

GymEnv constructs a Gymnasium (or legacy OpenAI Gym) environment directly from its registered ID. GymWrapper accepts a pre-existing gym Env object. Both expose a fully TorchRL-compatible interface including spec translation, device casting, from_pixels support, and info dict integration. 
pip install gymnasium
# or: pip install gym

from torchrl.envs import GymEnv, check_env_specs

env = GymEnv(
    "HalfCheetah-v4",
    device="cpu",
    frame_skip=4,
    from_pixels=False,
)
check_env_specs(env)

td = env.rollout(max_steps=200)
print(td["next", "reward"].mean())
For gymnasium ≥ 1.0, TorchRL currently requires gymnasium < 1.0 due to breaking API changes. See the compatibility discussion for updates.

​
DMControlEnv / DMControlWrapper

DMControlEnv wraps any task from DeepMind Control Suite. State and pixel observations, continuous and discrete action specs, and multi-environment parallelism are all supported out of the box. 
pip install dm_control

from torchrl.envs import DMControlEnv

env = DMControlEnv(
    env_name="cheetah",
    task_name="run",
    from_pixels=True,
    frame_skip=4,
    device="cpu",
)

td = env.rand_step()
print(td["next", "pixels"].shape)   # torch.Size([240, 320, 3])
print(td["next", "reward"].shape)   # torch.Size([1])

# View all available domain/task pairs
from torchrl.envs import DMControlEnv
print(DMControlEnv.available_envs)

​
BraxEnv / BraxWrapper

BraxEnv wraps Google Brax, a JAX-based differentiable physics engine. Brax uses JAX vmap for hardware-accelerated batched simulation; TorchRL converts JAX arrays to PyTorch tensors, preserving gradient flow. 
pip install brax

import torch
from torchrl.envs import BraxEnv

device = "cuda" if torch.cuda.is_available() else "cpu"
env = BraxEnv(
    "ant",
    batch_size=[256],   # 256 Brax environments via JAX vmap
    device=device,
)
env.set_seed(42)
td = env.rollout(max_steps=100)
print(td.shape)   # torch.Size([256, 100])
batch_size and num_workers are orthogonal scaling axes. BraxEnv("ant", batch_size=[128], num_workers=4) creates 4 processes, each running 128 Brax environments, for 512 total.

​
IsaacGymEnv / IsaacLabWrapper

TorchRL provides wrappers for both the legacy IsaacGym Preview (deprecated) and the actively maintained IsaacLab framework. IsaacLab envs run directly on GPU via NVIDIA Isaac Sim. 
# IsaacLab: follow https://isaac-sim.github.io/IsaacLab/installation

from torchrl.envs import IsaacGymEnv

env = IsaacGymEnv(task="Ant", num_envs=2048, device="cuda:0")
td = env.rollout(max_steps=100)
print(td.shape)   # torch.Size([2048, 100])

​
VmasEnv / VmasWrapper

VMAS (Vectorised Multi-Agent Simulator) provides differentiable multi-agent scenarios running in vectorised PyTorch. TorchRL’s wrapper exposes each agent’s observations and actions as nested TensorDict entries under each agent name. 
pip install vmas

from torchrl.envs import VmasEnv

env = VmasEnv(
    scenario="navigation",
    num_envs=32,
    n_agents=4,
    device="cpu",
)
td = env.reset()
print(td.keys())
# dict_keys(['agent_0', 'agent_1', 'agent_2', 'agent_3', 'done', ...])

td = env.rollout(max_steps=50)
print(td.shape)   # torch.Size([32, 50])
n_agents and other scenario-specific parameters (e.g. reward sparsity, collision penalties) are passed as extra keyword arguments directly to the VMAS scenario constructor, not as constructor params of VmasEnv itself. The available arguments vary by scenario — see the VMAS scenarios folder for details.

​
PettingZooEnv / PettingZooWrapper

PettingZoo covers a wide range of multi-agent environments. TorchRL’s wrapper handles both AEC (alternating) and parallel APIs, mapping each agent to a nested TensorDict key according to a MarlGroupMapType. 
pip install pettingzoo

from torchrl.envs import PettingZooEnv

env = PettingZooEnv(
    task="simple_spread_v3",
    parallel=True,
    device="cpu",
)
td = env.reset()
td = env.rollout(max_steps=25)

​
JumanjiEnv / JumanjiWrapper

Jumanji provides JAX-based combinatorial and scheduling environments. TorchRL wraps the JAX outputs into PyTorch tensors via a jax-to-torch bridge. 
pip install jumanji

from torchrl.envs import JumanjiEnv

env = JumanjiEnv("TSP50-v1", batch_size=[64], device="cpu")
td = env.reset()
td = env.rollout(max_steps=50)
print(td.shape)   # torch.Size([64, 50])

​
OpenSpielEnv / OpenSpielWrapper

OpenSpiel covers a large catalogue of board and card games for game-theoretic RL research. Each player maps to a nested TensorDict agent group. 
pip install open_spiel

from torchrl.envs import OpenSpielEnv

env = OpenSpielEnv("connect_four", device="cpu")
td = env.reset()
td = env.rollout(max_steps=42)

​
SMACv2Env / SMACv2Wrapper

SMACv2 (StarCraft Multi-Agent Challenge v2) is a standard cooperative multi-agent benchmark running on top of StarCraft II. Requires StarCraft II and the SMAC bridge library. 
pip install smacv2
# Also install StarCraft II + SMAC maps

from torchrl.envs import SMACv2Env

env = SMACv2Env(
    map_name="10gen_terran",
    device="cpu",
)
td = env.reset()
td = env.rollout(max_steps=150)

​
SafetyGymnasiumEnv / SafetyGymnasiumWrapper

Safety-Gymnasium extends Gymnasium with cost signals for safe-RL research. The TorchRL wrapper surfaces cost alongside reward in the output TensorDict. 
pip install safety-gymnasium

from torchrl.envs.libs.safety_gymnasium import SafetyGymnasiumEnv

env = SafetyGymnasiumEnv("SafetyPointGoal1-v0", device="cpu")
td = env.rollout(max_steps=100)
print(td["next", "reward"])
print(td["next", "cost"])   # safety cost signal

​
MultiThreadedEnv / MultiThreadedEnvWrapper (EnvPool)

MultiThreadedEnv wraps EnvPool, a C++-backed multi-threaded environment that achieves extremely high throughput for Atari and MuJoCo environments. 
pip install envpool

from torchrl.envs import MultiThreadedEnv

env = MultiThreadedEnv(
    num_workers=32,
    env_name="Pong-v5",
    device="cpu",
)
td = env.rollout(max_steps=10)
print(td.shape)   # torch.Size([32, 10])

​
MeltingpotEnv / MeltingpotWrapper

Melting Pot tests social learning and cooperation in multi-agent scenarios. TorchRL maps each player to a nested TensorDict entry. 
pip install dm-meltingpot

from torchrl.envs import MeltingpotEnv

env = MeltingpotEnv("commons_harvest__open", device="cpu")
td = env.reset()
td = env.rollout(max_steps=20)

​
GenesisEnv / GenesisWrapper

Genesis is a GPU-accelerated physics simulation platform for robotics and embodied AI. The TorchRL wrapper exposes Genesis environments via the standard EnvBase interface with JAX-to-Torch tensor conversion. 
pip install genesis-world

from torchrl.envs import GenesisEnv

env = GenesisEnv("FrankaPush-v0", device="cuda:0")
td = env.rollout(max_steps=50)

​
Summary Table

ClassLibraryInstallMulti-AgentGPU-native
GymEnvGymnasium / OpenAI Gympip install gymnasium❌❌
DMControlEnvDM Controlpip install dm_control❌❌
BraxEnvBrax (JAX)pip install brax❌✅ (via JAX)
IsaacGymEnvIsaacGym (legacy)see NVIDIA docs❌✅
IsaacLabWrapperIsaacLabsee NVIDIA docs❌✅
VmasEnvVMASpip install vmas✅✅
PettingZooEnvPettingZoopip install pettingzoo✅❌
JumanjiEnvJumanji (JAX)pip install jumanji❌✅ (via JAX)
OpenSpielEnvOpenSpielpip install open_spiel✅❌
SMACv2EnvSMACv2pip install smacv2✅❌
SafetyGymnasiumEnvSafety-Gymnasiumpip install safety-gymnasium❌❌
MultiThreadedEnvEnvPoolpip install envpool❌❌
MeltingpotEnvMelting Potpip install dm-meltingpot✅❌
GenesisEnvGenesispip install genesis-world❌✅
Use from torchrl.envs import get_available_libraries to programmatically query which optional dependencies are installed in the current environment.

Build docs developers (and LLMs) love

Get started for freeTalk to us