Diffusion Pusht Keypoints

🚀 擴散策略 / PushT（關鍵點）模型卡片

本模型為擴散策略（Diffusion Policy），依據論文Diffusion Policy: Visuomotor Policy Learning via Action Diffusion，針對 gym-pusht 中的 PushT 環境進行訓練，僅使用關鍵點觀測。

注意：原研究通過圖像修復（inpainting）進行僅關鍵點的條件訓練。在此，我們將觀測與智能體位置一起編碼，並將編碼結果作為去噪 U-Net 的全局條件。

🚀 快速開始

使用 python lerobot/scripts/eval.py -p lerobot/diffusion_pusht 對模型進行 50 個回合的評估，輸出結果將保存至 outputs/eval 目錄。

更多信息，請參考 LeRobot 庫（特別是評估腳本）。

📦 安裝指南

暫未提及安裝相關內容，跳過該章節。

💻 使用示例

基礎用法

python lerobot/scripts/eval.py -p lerobot/diffusion_pusht

🔧 技術細節

訓練詳情

本模型使用 LeRobot@cc2f6e7 進行訓練。

模型使用 [LeRobot 的訓練腳本](TODO link with commit hash one the PR is merged) 和 pusht_keypoints 數據集進行訓練，訓練命令如下：

python lerobot/scripts/train.py \
  hydra.job.name=diffusion_pusht_keypoints \
  hydra.run.dir=outputs/train/2024-07-03/13-52-44_diffusion_pusht_keypoints \
  env=pusht_keypoints \
  policy=diffusion_pusht_keypoints \
  training.save_checkpoint=true \
  training.offline_steps=200000 \
  training.save_freq=20000 \
  training.eval_freq=10000 \
  training.log_freq=50 \
  training.num_workers=4 \
  eval.n_episodes=50 \
  eval.batch_size=50 \
  wandb.enable=true \
  wandb.disable_artifact=true \
  device=cuda \
  use_amp=true

訓練曲線可在 https://wandb.ai/alexander-soare/lerobot/runs/5z9d8q9q/overview 查看。

在 Nvida RTX H100 上訓練該模型大約需要 5 小時。

評估詳情

模型在 gym-pusht 的 PushT 環境中進行評估。每個回合有兩個評估指標：

• 與目標的最大重疊度（在上述圖表中顯示為 eval/avg_max_reward），範圍為 [0, 1]。
• 成功率：最大重疊度是否至少達到 95%。

以下是 500 個回合評估的指標：

每次滾動評估的結果可在 eval_info.json 中查看。

📄 許可證

本模型採用 Apache-2.0 許可證。