Skip to content

README_zh.md

Latest commit

 

History

History
201 lines (144 loc) · 7.59 KB

README_zh.md

File metadata and controls

201 lines (144 loc) · 7.59 KB
Raw

ROCm_Robotics_RL_Lab

一个面向 AMD GPU 平台的机器人仿真强化学习实验仓库，当前经过完整训练与验证的基线是 robosuite Panda + Lift。

本项目完整的开发、RL 训练和测试验证，均基于 AMD Ryzen AI Max+ 395 笔记本平台完成，充分利用了其集成的 Radeon 8060S 显卡进行 OpenGL 仿真渲染与 ROCm / PyTorch AI 计算。这也是一个非常适合机器人 RL 实验的便携开发平台。

English version: see README.md

📁 项目结构

ROCm_Robotics_RL_Lab/
├── docs/                        # 文章与发布素材
├── environments/
│   ├── gym_wrapper.py           # robosuite -> Gymnasium 适配器
│   └── pick_cube_place_cup.py   # 自定义抓取放置环境原型
├── scripts/
│   ├── quickstart.py            # 快速开始示例
│   ├── train_sac.py             # Panda Lift 的 SAC 训练脚本
│   ├── train_ppo.py             # Panda Lift 的 PPO 训练脚本
│   ├── evaluate.py              # 评估与视频录制脚本
│   ├── seed_sweep.py            # 多 seed 运行脚本
│   └── param_sweep.py           # 参数 sweep 脚本
├── model_loading.py             # SB3 模型加载辅助工具
├── requirements.txt
├── README.md
└── README_zh.md

🚀 快速开始

1. 环境配置 (AMD GPU / ROCm)

# 创建虚拟环境
uv venv .venv --python 3.12
source .venv/bin/activate

# 安装 PyTorch ROCm 版本（AMD GPU）
uv pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.1

验证 GPU 是否识别：

python -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))"

安装依赖：

uv pip install -r requirements.txt
uv pip install stable-baselines3[extra]

# robosuite 一次性初始化
python .venv/lib/python3.12/site-packages/robosuite/scripts/setup_macros.py

2. 运行快速示例

cd ROCm_Robotics_RL_Lab
python scripts/quickstart.py

这会在 robosuite 的 Lift 任务上训练一个小型 Panda SAC 策略。

3. 完整训练

# SAC 训练
python scripts/train_sac.py --total-timesteps 500000 --n-envs 4

# PPO 训练
python scripts/train_ppo.py --total-timesteps 1000000 --n-envs 8

# 实时观察训练动作（单环境）
python scripts/train_sac.py --total-timesteps 500000 --n-envs 1 --render

# 如需显示 Stable-Baselines3 进度条，可显式启用
python scripts/train_sac.py --total-timesteps 500000 --n-envs 4 --progress-bar

当前 train_sac.py 已针对并行环境做了更稳健的默认设置：

  • 会把 --save-freq / --eval-freq 按 n_envs 自动换算成 SB3 回调频率，避免并行环境下评估和保存过稀
  • 默认超参已经向“稳定成功率”而不是“快速刷高稠密奖励”倾斜：learning_rate=1e-4、batch_size=512、learning_starts=20000、tau=0.002、gradient_steps=1
  • 这组默认值会降低 critic / actor 抖动，减少过早利用塑形奖励的倾向，更适合当前 Lift 任务的稀疏成功目标
  • 训练环境会在成功时结束回合，并把时间上限作为 truncated 处理，减少把超时误当成真正终止带来的价值学习偏差
  • 基于第一轮参数 sweep，目前默认的成功导向配置是给成功回合额外 +100 的终止奖励（--success-bonus），并不额外惩罚超时失败（--timeout-penalty 0）；这组参数在多 seed 平均表现最好
  • best_success 现在默认使用 20 个评估回合和 20% 的成功率门槛（--n-eval-episodes 20 --min-best-success-rate 0.20），在可靠性和训练速度之间做平衡
  • best_success/best_metrics.json 会记录该 checkpoint 的 timestep、success_rate 和 mean_reward，训练结束时终端也会打印同样的信息

如果你希望把阈值调得更严格，可以显式设置，例如：

python scripts/train_sac.py --total-timesteps 500000 --n-envs 4 --n-eval-episodes 20 --min-best-success-rate 0.20

如果你发现策略依旧偏向“拿塑形奖励但不完成任务”，可以进一步加大成功导向：

python scripts/train_sac.py --total-timesteps 500000 --n-envs 4 --n-eval-episodes 20 --min-best-success-rate 0.20 --success-bonus 100 --timeout-penalty 0

如果你想显式写出当前推荐的稳定训练配置，可以直接使用：

python scripts/train_sac.py --total-timesteps 500000 --n-envs 4 --learning-rate 1e-4 --batch-size 512 --learning-starts 20000 --gradient-steps 1 --tau 0.002 --n-eval-episodes 20 --min-best-success-rate 0.20 --success-bonus 100 --timeout-penalty 0

如果你怀疑成功率受随机种子影响很大，可以直接批量跑多个 seed：

python scripts/seed_sweep.py --seeds 42 123 456 --total-timesteps 500000 --n-envs 4

这个脚本会：

  • 依次调用现有的 scripts/train_sac.py
  • 优先评估每个 run 的 best_success checkpoint；如果没有，则回退到最终模型
  • 在 models/seed_sweeps/<timestamp>/summary.json 中汇总每个 seed 的训练 / 评估结果

如果你想先确认将执行哪些命令而不真正开跑，可以使用：

python scripts/seed_sweep.py --seeds 42 123 456 --dry-run

4. 评估模型

python scripts/evaluate.py --model models/sac_lift_final.zip --algo sac --n-episodes 10
python scripts/evaluate.py --model models/best/best_model.zip --algo sac --no-render
python scripts/evaluate.py --model <path> --algo ppo --record-video --video-dir videos/

评估脚本现在会按整回合统计 success：只要回合中任意一步完成任务，就记为成功，而不是只检查最后一帧状态。

🎯 当前项目重点

当前仓库中已经跑通并完成验证的主线是：

  • robosuite Lift 任务
  • Panda 机械臂
  • Stable-Baselines3 的 SAC / PPO
  • AMD GPU + ROCm + PyTorch 训练
  • 基于 OpenGL 的渲染、评估与视频录制

pick_cube_place_cup.py 仍作为后续自定义任务原型保留在仓库中，但当前文档化和验证完成的基线是 Panda Lift。

📊 训练参数建议

SAC 参数

参数 推荐值 说明
learning_rate 1e-4 学习率
buffer_size 1,000,000 经验回放缓冲区
batch_size 512 批大小
learning_starts 20,000 开始更新前的随机采样步数
gradient_steps 1 更保守的更新频率
gamma 0.99 折扣因子
tau 0.002 软更新系数

PPO 参数

参数 默认值 说明
learning_rate 3e-4 学习率
n_steps 2048 每次更新步数
batch_size 64 批大小
n_epochs 10 训练轮数
clip_range 0.2 裁剪范围

📈 监控训练

使用 TensorBoard 查看训练曲线：

tensorboard --logdir logs/

🔗 相关资源

📝 注意事项

  1. 当前仓库验证完成的主线是 robosuite Panda + Lift。
  2. 若要实时观察训练过程，请使用 --render，并将 --n-envs 设为 1。
  3. 若在无图形界面的环境中评估模型，建议使用 --no-render 以避免 GLFW / DISPLAY 报错。
  4. 训练脚本默认关闭 SB3 的富文本进度条，以避免部分环境在退出时打印 tqdm / rich 的清理 traceback；如需显示，可添加 --progress-bar。

📜 许可证

MIT License