pytorch/rl

TorchRL – 原生 PyTorch 强化学习工具包

是什么 – TorchRL 是一个基于 PyTorch 构建的库，为强化学习（RL）研究和生产提供 可组合 的构建模块。它 不 提供单一算法；相反，它提供统一的数据容器（TensorDict）和一组可互换的模块（环境、策略、收集器、回放缓冲区、损失对象、训练器等），可在保持与原生 PyTorch 编程模型一致的前提下自由组合。

核心理念

理念	TorchRL 的实现方式
显式、命名的数据	所有张量都在 TensorDict 内部传递，该容器在整个训练循环中携带字段名、批量维度和设备位置信息。
模块化堆栈	环境、策略、收集器、回放缓冲区和损失模块是独立、可替换的组件。
从原型到生产可扩展	相同的 API 可用于单进程滚动、多进程异步收集器、分布式训练、编译（torch.compile）或 CUDA 加速管道，无需代码更改。

主要组件

• TensorDict – 支持完整 PyTorch 操作、设备传输、共享内存和内存映射的字典式容器。
• 环境与变换 – 原生 PendulumEnv、MuJoCo 任务，以及 Gymnasium、DM-Control、Brax、PettingZoo 等的包装器；变换（观测归一化、动作缩放、帧堆叠、HER 等）是一等模块。
• 收集器 – 同步、异步、多进程和分布式收集器，可批量轨迹、将数据移动到正确设备，并可实时更新策略权重。
• 回放缓冲区 – 模块化存储（内存内、懒加载内存映射、CUDA 友好）、优先采样、HER、离线数据集处理。
• 模块 / 策略 – 带有显式输入/输出键契约的常规 nn.Module；包括随机动作器、评论家、循环网络、分布包装器、世界模型组件。
• 目标函数 – PPO、SAC、DQN、TD3、REDQ、IQL、CQL、决策变换器、DreamerV3、MAPPO/IPPO、QMIX/VDN、行为克隆等损失模块，全部读写命名键。
• 训练器与 Hydra 配置 – 高级工具，将环境、收集器、损失、优化器和日志连接为可复现的配方。

最近亮点（v0.13）

• 使用 Triton 支持的 GRU/LSTM 重置实现更快的循环路径。
• 新增 MuJoCo 环境、卫星示例和宏观控制原语。
• 扩展多智能体支持：MAPPO、IPPO、MultiAgentGAE、价值归一化工具。
• 异步优先回放缓冲区写入、紧凑观测存储、可选 CUDA 内核用于回放。
• 新增变换（ActionScaling、FlattenAction、NextObservationDelta 等）和价值估计器改进。

典型工作流（快速演示）

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

# 带有简单变换堆栈的环境
env = TransformedEnv(PendulumEnv(), StepCounter(max_steps=200))

# 以常规 nn.Module 表示的策略，带有显式 TensorDict 键
policy = TensorDictModule(
    nn.Sequential(nn.LazyLinear(64), nn.Tanh(), nn.Linear(64, 1), nn.Tanh()),
    in_keys=["observation"],
    out_keys=["action"],
)

# 一步滚动 – 结果是一个包含观测、动作、奖励等的 TensorDict
rollout = env.rollout(max_steps=32, policy=policy)
print(rollout.batch_size)          # torch.Size([32])
print(rollout["next", "reward"].shape)  # torch.Size([32])

相同的 TensorDict 可以直接输入收集器、存储到回放缓冲区，并由损失模块消费，无需任何格式转换。

适合谁使用？

• 研究人员：构建新 RL 算法（需要灵活的数据流、组件轻松替换、与 PyTorch 的 autograd/compile 紧密集成）。
• 工程师：将 RL 管道扩展到多个 CPU/GPU 或分布式集群（异步收集器、CUDA 友好回放缓冲区、内存映射存储）。
• 机器人/仿真团队：需要原生 MuJoCo 或自定义环境，并支持设备内变换。
• 多智能体或基于模型的 RL 项目（提供 VMAS、PettingZoo 包装器、DreamerV3、决策变换器组件）。
• LLM 后训练 实验（TorchRL 包含一个小型 LLM 堆栈，用于 GRPO/SFT 风格微调）。

安装

# 稳定版本（仅 CPU 回放缓冲区）
pip install torchrl

# 启用 CUDA 的回放缓冲区（将 cu118 替换为你的 CUDA 版本）
pip install "torchrl==0.13.0+cu118" \
    --extra-index-url https://download.pytorch.org/whl/cu118

# 可选扩展（额外环境/工具）
pip install "torchrl[utils]"          # Hydra、日志等
pip install "torchrl[gym_continuous]" # Gymnasium 连续控制
pip install "torchrl[atari]"          # Atari 支持
pip install "torchrl[marl]"           # 多智能体库（PettingZoo、VMAS、…）

文档与资源

• 入门与教程： https://pytorch.org/rl/stable/index.html#getting-started
• API 参考： https://pytorch.org/rl/stable/reference/index.html
• SOTA 实现： sota-implementations/ 文件夹（PPO、SAC、TD3 等）
• 基准测试与 CI 状态： README 中的徽章链接到实时仪表板。
• 社区： Discord 服务器、Twitter 和 GitHub 问题用于支持。

---

总结： TorchRL 是一个功能齐全、以 PyTorch 为先的 RL 工程框架，让您在保持代码库整洁和 PyTorch 原生的前提下，对广泛的 RL、多智能体、基于模型以及甚至 LLM 微调工作流进行原型设计、扩展和实验。