pytorch/rl

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

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 微調工作流程進行原型設計、擴展與實驗。