pytorch/rl

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

TorchRL – PyTorchネイティブ強化学習ツールキット

何であるか – TorchRLはPyTorchの上に構築されたライブラリで、強化学習(RL)の研究およびプロダクションに向けた組み合わせ可能な構成要素を提供します。これは単一のアルゴリズムを提供するものではありません。代わりに、統一されたデータコンテナ(TensorDict)と、交換可能なモジュール(環境、ポリシー、コレクタ、リプレイバッファ、損失オブジェクト、トレーナーなど)を提供し、ネイティブなPyTorchプログラミングモデルに忠実なまま、自由に組み合わせることができます。

コアとなるアイデア

アイデア TorchRLの実装方法
明示的で名前付きのデータ すべてのテンソルは、フィールド名、バッチ次元、デバイス配置を保持するTensorDict内を移動します。トレーニングループ全体で一貫性を持ちます。
モジュールスタック 環境、ポリシー、コレクタ、リプレイバッファ、損失モジュールは独立しており、交換可能で、自由に組み合わせられます。
プロトタイプからプロダクションまでスケーラブル 同期プロセス、非同期マルチプロセスコレクタ、分散トレーニング、torch.compileまたはCUDA加速パイプラインなど、コード変更なしで同じAPIが利用可能。

主要なコンポーネント

  • TensorDict – PyTorchの操作、デバイス間転送、共有メモリ、メモリマップを完全にサポートする辞書型コンテナ。
  • 環境と変換 – ネイティブなPendulumEnv、MuJoCoタスク、Gymnasium、DM-Control、Brax、PettingZooなどへのラッパー。変換(観測正規化、アクションスケーリング、フレームスタック、HERなど)は一等モジュールとして扱われます。
  • コレクタ – 同期、非同期、マルチプロセス、分散コレクタ。軌道をバッチ化し、データを適切なデバイスに移動し、リアルタイムでポリシー重みを更新可能。
  • リプレイバッファ – モジュール式ストレージ(メモリ内、遅延メモリマップ、CUDA対応)、優先度サンプリング、HER、オフラインデータセット処理。
  • モジュール / ポリシー – 明示的な入出力キー契約を持つ通常のnn.Module。確率的アクター、クライティック、再帰ネット、分布ラッパー、ワールドモデルコンポーネントを含む。
  • 目的関数 – PPO、SAC、DQN、TD3、REDQ、IQL、CQL、Decision Transformers、DreamerV3、MAPPO/IPPO、QMIX/VDN、行動クラッシングなど、すべて名前付きキーを読み書きする損失モジュール。
  • トレーナーとHydra設定 – 環境、コレクタ、損失、最適化、ログを再現可能なレシピに接続する高レベルユーティリティ。

最近のハイライト(v0.13)

  • TritonバックエンドによるGRU/LSTMリセットの高速化。
  • 新しいMuJoCo環境、衛星例、マクロ制御プリミティブ。
  • 多エージェントサポートの拡張:MAPPO、IPPO、MultiAgentGAE、価値正規化ユーティリティ。
  • 非同期優先リプレイバッファ書き込み、コンパクトな観測ストレージ、オプションのCUDAカーネルによるリプレイ。
  • 追加の変換(ActionScalingFlattenActionNextObservationDeltaなど)および価値推定器の改善。

一般的なワークフロー(クイックデモ)

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))

# 明示的なTensorDictキーを持つ通常のnn.Moduleとして表現されたポリシー
policy = TensorDictModule(
    nn.Sequential(nn.LazyLinear(64), nn.Tanh(), nn.Linear(64, 1), nn.Tanh()),
    in_keys=["observation"],
    out_keys=["action"],
)

# 1ステップのロールアウト – 結果は観測、アクション、報酬などを含む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との密接な統合が必要)。
  • エンジニア:多数のCPU/GPUや分散クラスタにRLパイプラインをスケーリング(非同期コレクタ、CUDA対応リプレイバッファ、メモリマップストレージ)。
  • ロボット工学/シミュレーションチーム:ネイティブなMuJoCoまたはカスタム環境と、デバイス上で実行される変換を必要とする。
  • 多エージェントまたはモデルベースRLプロジェクト(VMAS、PettingZooラッパー、DreamerV3、Decision-Transformerコンポーネントを提供)。
  • LLMの後処理実験(TorchRLはGRPO/SFTスタイルのファインチューニング用に小さなLLMスタックを含む)。

インストール

# ステーブルリリース(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、…)

ドキュメントとリソース


結論: TorchRLは、コードベースをクリーンかつPyTorchネイティブに保ちながら、幅広いRL、多エージェント、モデルベース、さらにはLLMファインチューニングワークフローのプロトタイピング、スケーリング、実験を可能にする、フル機能のPyTorch第一のRLエンジニアリングフレームワークです。

関連

  • プロジェクト
  • プロジェクト
  • プロジェクト
  • Dispatch
  • プロジェクト