MuJoCo + PPO 二阶倒立摆学习框架

这是一个为学习 PPO 而写的最小强化学习项目。环境是 MuJoCo 中的小车二阶倒立摆，评估时从末端竖直朝下的悬挂姿态开始，目标是 swing-up 后稳定在倒立平衡点；训练默认混合使用悬挂、随机角度和少量倒立初始状态，避免 PPO 在纯悬挂 swing-up 任务上太早卡住。算法是从零实现的 PPO，代码尽量按数据流拆开，方便你从采样一直看到损失函数。

本机已经有 /home/sustech/humanoid_rl/unitree_rl_mjlab 和 conda 环境 /home/sustech/miniconda3/envs/unitree_rl_mjlab。这个项目复用其中已经安装好的 mujoco、torch 等依赖，但不把二阶倒立摆强行塞进 Unitree 机器人任务框架里，这样更适合学习。

目录

assets/double_inverted_pendulum.xml      MuJoCo MJCF 模型
double_pendulum_ppo/envs/                环境与同步向量环境
double_pendulum_ppo/ppo/                 Actor-Critic、rollout buffer、PPO 更新
configs/double_pendulum_ppo.yaml         默认训练配置
scripts/train.py                         训练入口
scripts/evaluate.py                      评估/渲染入口
scripts/plot_metrics.py                  绘制训练曲线
tests/                                  轻量 smoke tests

运行环境

推荐直接使用已有的 Unitree/Mjlab 环境：

conda activate unitree_rl_mjlab
cd /home/sustech/RL_proj

如果你要在新环境里安装：

python -m pip install -r requirements.txt

CUDA / PyTorch 检查

当前 unitree_rl_mjlab 环境已经安装 CUDA 11.8 构建的 PyTorch：

torch==2.7.1+cu118
torchvision==0.22.1+cu118
torchaudio==2.7.1+cu118

检查 GPU 是否真的可用：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/check_cuda.py

如果 nvidia-smi 不能正常显示 GPU，说明系统 NVIDIA driver 没有对当前会话可用，重装 PyTorch 也不能启用 GPU。确认驱动正常后，如需重装 CUDA 11.8 版 PyTorch，可使用：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python -m pip install \
  --index-url https://download.pytorch.org/whl/cu118 \
  --upgrade --force-reinstall \
  torch==2.7.1+cu118 torchvision==0.22.1+cu118 torchaudio==2.7.1+cu118

训练

快速 smoke run：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/train.py \
  --total-timesteps 4096 \
  --num-envs 4 \
  --num-steps 64 \
  --run-name smoke

正式一点的训练：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/train.py \
  --config configs/double_pendulum_ppo.yaml

训练会同时写入 metrics.csv 和 TensorBoard event 文件。启动曲线面板：

tensorboard --logdir runs --port 6006

然后在浏览器打开 http://localhost:6006。关键曲线建议先看 task/max_pose_alignment_score、task/upright_streak_steps、task/tip_speed、 task/swing_reversal_event、task/tip_speed_over_limit、reward/upright_hold_reward。 如果只想写 CSV，可以加 --no-tensorboard。

默认配置使用 vector_env: subproc，会用多进程并行跑多个 MuJoCo 环境。env.num_envs 可以按 CPU 核数调大，例如 16、32、64；GPU 主要负责 PPO 网络前向和更新，MuJoCo step 主要吃 CPU。

默认课程采用 teacher + BC + PPO finetune：先用启发式 teacher 生成从 hanging 起摆的动作数据，行为克隆预训练 actor，再用 PPO 继续优化。这样避免 PPO 从完全静止 hanging 纯随机探索时长期找不到正确起摆相位。配置已打开 include_phase_observation，观测里包含起摆节拍相位，方便无记忆 MLP 模仿 teacher 的换向节奏。

课程后期仍会用到按高度分桶覆盖的 qpos/qvel assisted reset 状态库。它不是主线，只负责补充中高位接住经验。

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/generate_assisted_resets.py \
  --mode mixed \
  --max-states 50000 \
  --height-bins 0.18:0.35,0.35:0.50,0.50:0.65,0.65:0.80,0.80:0.95 \
  --out assets/assisted_reset_states_balanced.npz

生成 teacher 数据：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/generate_teacher_data.py \
  --episodes 500 \
  --steps 700 \
  --max-samples 200000 \
  --min-episode-height 0.15 \
  --reset-kick-prob 0.2 \
  --reset-angular-velocity-scale 1.0 \
  --action-limit 120 \
  --out assets/teacher_swingup_data.npz

行为克隆预训练：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/pretrain_bc.py \
  --data assets/teacher_swingup_data.npz \
  --epochs 20 \
  --batch-size 4096 \
  --device cuda \
  --out runs/bc_teacher_swingup.pt

检查 BC 起摆能力：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/evaluate.py \
  --checkpoint runs/bc_teacher_swingup.pt \
  --episodes 5 \
  --deterministic \
  --no-sleep \
  --device cuda

正常情况下 hmax 应该能到 0.25 左右。然后用 BC checkpoint 初始化 PPO：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/train.py \
  --config configs/double_pendulum_ppo.yaml \
  --init-checkpoint runs/bc_teacher_swingup.pt \
  --device cuda \
  --total-timesteps 30000000 \
  --num-envs 128 \
  --run-name bc_ppo_30m

如果只想跑短一点的验证版：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/train.py \
  --config configs/double_pendulum_ppo.yaml \
  --init-checkpoint runs/bc_teacher_swingup.pt \
  --device cuda \
  --total-timesteps 12000000 \
  --num-envs 64 \
  --run-name bc_ppo_12m

训练时重点观察 l1max / max_link1_score、amax / max_alignment_score、hmax / max_height_score、upright_reward 和 success_reward。hmax 只表示末端 tip 高度，可能被“第一节还挂着、第二节折起来”的假高状态骗过；l1max 和 amax 才能说明第一节是否真的被甩上去、两节是否同时接近倒立。如果 hmax 上升但 l1max≈0，说明策略仍然只是在甩第二节。

当前启发式 teacher 只能提供低/中高度起摆示范，不等价于“到顶端并接住”的专家。可以先用开环搜索脚本诊断 teacher 能否从严格 hanging 到达高位：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/search_open_loop_teacher.py \
  --iterations 12 \
  --population 192 \
  --elites 24 \
  --steps 700 \
  --segment-steps 35 \
  --action-limit 120 \
  --include-phase-observation \
  --out assets/open_loop_top_teacher.npz

也可以用采样 MPC 生成一条可重放的开环轨迹：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/generate_mpc_teacher.py \
  --steps 700 \
  --horizon-segments 24 \
  --segment-steps 5 \
  --iterations 4 \
  --population 192 \
  --elites 24 \
  --action-limit 120 \
  --include-phase-observation \
  --out assets/mpc_top_teacher.npz

这两个脚本保存的 .npz 都包含 observations、actions、qpos/qvel、link1_score、height_score、alignment_score。如果搜索结果仍然是 hmax≈0.3~0.4 且 link1≈0、align≈0，说明 teacher 还没有把第一节摆杆甩上去，不能作为完整到顶端示范；这时应优先训练高位 assisted_reset 的 catch 策略，再逐步把初始高度下放，而不是继续让 PPO 从底部盲撞顶端奖励。

如果不想用 assisted reset，把配置里的 assisted_reset_prob 都设为 0.0。

如果想用单进程方便调试，把配置改成：

vector_env: sync
env:
  num_envs: 8

训练结果会保存到 runs/<run_name>_<time>/，里面有：

• config.yaml：本次训练配置快照
• metrics.csv：每次 PPO 更新后的指标
• checkpoints/model_*.pt：周期性 checkpoint
• last.pt：最后一个 checkpoint

评估与可视化

无渲染评估：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/evaluate.py \
  --checkpoint runs/<run_name>/last.pt \
  --episodes 5 \
  --deterministic

打开 MuJoCo viewer：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/evaluate.py \
  --checkpoint runs/<run_name>/last.pt \
  --episodes 3 \
  --deterministic \
  --render

诊断“第一脚破对称”问题时，可以只在开头给一个固定动作，再交还给策略：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/evaluate.py \
  --checkpoint runs/<run_name>/last.pt \
  --episodes 3 \
  --deterministic \
  --render \
  --warmup-steps 20 \
  --warmup-action 1.0

如果加 warmup 后能继续摆，但不加 warmup 不动，说明主要问题是 deterministic policy 在完全 hanging 对称点输出接近 0；如果加 warmup 后也接不住，说明后续 pump/catch 还没有学好。

绘制曲线：

/home/sustech/miniconda3/envs/unitree_rl_mjlab/bin/python scripts/plot_metrics.py \
  runs/<run_name>/metrics.csv

PPO 代码阅读路线

建议按这个顺序看：

1. double_pendulum_ppo/envs/double_inverted_pendulum.py：观测、奖励、终止条件。
2. double_pendulum_ppo/ppo/agent.py：高斯策略和价值函数。
3. double_pendulum_ppo/ppo/storage.py：rollout buffer 与 GAE。
4. double_pendulum_ppo/ppo/ppo.py：ratio、clip surrogate、value loss、entropy bonus。
5. scripts/train.py：把环境采样和 PPO 更新串起来。

可调参数建议

• 想看探索影响：调 agent.init_log_std 和 ppo.entropy_coef。
• 想看 PPO clip 影响：调 ppo.clip_coef，同时观察 approx_kl 和 clip_frac。
• 学习不稳定：减小 ppo.learning_rate 或增加 ppo.num_steps。
• 想只学倒立平衡而不是 swing-up：把 env.initial_pose 改成 upright，并把 env.angle_limit 改成例如 1.0471975512。
• 想强制从纯悬挂状态训练：把 env.initial_pose 改成 hanging，但 PPO 会更容易卡在局部最优。

与 unitree_rl_mjlab 的关系

unitree_rl_mjlab 是面向 Unitree 机器人和大规模任务注册的框架；这里保留它的 MuJoCo/PyTorch 运行环境优势，但算法和环境完全展开在本项目内。学清楚这个项目后，再读 unitree_rl_mjlab/src/tasks/... 和它的训练脚本会轻松很多。