lucidrains/q-transformer

Q‑Transformer – 使用自回歸 Q 函數的離線強化學習

是什麼 – 2023年 Google DeepMind 論文《Q‑Transformer》的 PyTorch 實作。將 Q 函數視為自回歸模型，預測離散動作序列，從而實現機器人任務的可擴展離線強化學習。

核心元件

• QRoboticTransformer – 基於 MaxViT 的視覺 Transformer 骨幹，接收影片幀與可選的文字指示，輸出離散動作空間的 Q 值。
• QLearner – 在代理產生的回放記憶資料集上執行 Q 學習的訓練迴圈。
• Agent – 在使用者提供的環境（BaseEnvironment 的子類）中執行模型以收集軌跡。
• ReplayMemoryDataset – 用於儲存（狀態, 動作, 奖勵, 下一狀態, 完成）元組的簡單資料集包裝器。

如何取得

pip install q-transformer

典型工作流程

1. 定義環境 – 實作 BaseEnvironment 或使用提供的 MockEnvironment，該環境預期輸入影片張量與指示嵌入。
2. 建立模型 – 使用視覺超參數（例如小型 MaxViT）實例化 QRoboticTransformer，並指定離散動作數量與分箱數。
3. 收集資料 – Agent(model, environment, ...)() 執行多個 episode，將轉移儲存在回放緩衝區中，並寫入磁碟。
4. 訓練 – QLearner(model, dataset=ReplayMemoryDataset(), ...)() 執行 Q 學習（支援 dueling head、n 步回報、梯度累積等）。
5. 推論動作 – 訓練完成後，呼叫 model.get_optimal_actions(video, instructions) 以取得新影片-指示對的最佳離散動作序列。

獨特之處

• 實作 自回歸 Q 函數概念，使模型能高效處理多維離散動作空間。
• 包含可選的 dueling 架構、n 步回報與保守正則化占位符。
• 專為機器人設計：模型可處理多攝影機影片（(3, 6, 224, 224)）與語言指示。
• 倉儲追蹤詳細的 TODO 清單，反映持續進行的研究方向（例如 beam-search 解碼、對過去動作的交叉注意力、Gumbel 探索）。

適合誰使用

• 探索操作或具身代理離線強化學習的研究人員。
• 需要從影片與語言中實現離散化機器人控制即插即用基線的實務者。
• 任何希望重現或擴展 Q‑Transformer 論文的人（該實作已被 Kotb 等，2024 年重現）。

引用 – README 提供原始 Q‑Transformer 論文與相關研究（FlashAttention、持續學習正則化等）的 BibTeX 條目。

---

以上所有細節均直接取自倉儲的 README。