lucidrains/q-transformer

Implementation of Q-Transformer, Scalable Offline Reinforcement Learning via Autoregressive Q-Functions, out of Google Deepmind

Q‑Transformer – 自動回帰Q関数を用いたオフライン強化学習

何であるか – 2023年Google DeepMindの論文『Q‑Transformer』のPyTorch実装。Q関数を自動回帰モデルとして扱い、離散的アクションのシーケンスを予測することで、ロボットタスクにおけるスケーラブルなオフライン強化学習を可能にする。

主要なコンポーネント

QRoboticTransformer – MaxViTに基づくビジョン・トランスフォーマーのバックボーン。動画フレームとオプションのテキスト指示を受け取り、離散化されたアクション空間のQ値を出力する。
QLearner – エージェントによって生成されたリプレイメモリデータセット上でQ学習を実行するトレーニングループ。
Agent – ユーザーが提供する環境（BaseEnvironmentのサブクラス）でモデルを実行し、軌道を収集する。
ReplayMemoryDataset – (状態, アクション, 報酬, 次の状態, 終了)タプルを格納するシンプルなデータセットラッパー。

入手方法

pip install q-transformer

一般的なワークフロー

環境の定義 – BaseEnvironmentを実装するか、動画テンソルと指示埋め込みを想定する提供された MockEnvironment を使用する。
モデルの作成 – 視覚的なハイパーパラメータ（例：小さなMaxViT）を指定し、離散的アクション数とビン数を設定して QRoboticTransformer をインスタンス化する。
データ収集 – Agent(model, environment, ...)() がエピソードを実行し、遷移をリプレイバッファに保存し、ディスクに書き込む。
学習 – QLearner(model, dataset=ReplayMemoryDataset(), ...)() がQ学習を実行（ダウエルヘッド、nステップリターン、勾配蓄積などをサポート）。
アクションの推論 – 学習後、model.get_optimal_actions(video, instructions) を呼び出して、新しい動画・指示ペアに対する最適な離散的アクションシーケンスを取得する。

特徴

自動回帰Q関数のアイデアを実装しており、多次元離散的アクション空間を効率的に扱える。
オプションのダウエルアーキテクチャ、nステップリターン、保守的な正則化プレースホルダーを含む。
ロボット工学向けに設計：マルチカメラ動画（(3, 6, 224, 224)）と言語指示を処理する。
リポジトリは詳細なTODOリストを追跡しており、継続的な研究方向（例：ビームサーチデコード、過去のアクションへのクロスアテンション、Gumbelベースの探索）を反映している。

誰が使うべきか

引用 – READMEには、オリジナルのQ‑Transformer論文および関連研究（FlashAttention、継続的学習正則化など）のBibTeXエントリが提供されている。

上記のすべての詳細は、リポジトリのREADMEから直接引用されています。