纠错,疑问,交流: 请进入讨论区或请点击进入页面，扫码加入微信群或Q群进行交流

获取最新文章: 扫一扫加入“创客智造”公众号

欢迎加入我们的openduckmini交流群，微信扫描右侧二维码立即进群交流

训练系统

理解如何实现训练模型，包含训练管线，runner架构，PPO训练参数，回调机制，数据流，训练命令和技巧等

概述

Open Duck Playground 的训练系统基于 Brax 的 PPO（Proximal Policy Optimization，近端策略优化） 算法实现。训练系统提供了从环境配置到策略学习的完整管线，包括训练循环管理、检查点保存、日志记录和模型导出。

训练管线流程

┌───────────────────────────────────────────────────────────┐
│                    训练管线总览                              │
└───────────────────────────────────────────────────────────┘

初始化 Runner
    │
    ▼
加载环境（Joystick / Standing）
    │
    ▼
配置 PPO 参数（从 locomotion_params 获取）
    │
    ▼
配置网络工厂（policy_hidden_layer_sizes 等）
    │
    ▼
启动训练循环
    │
    ├──→ 环境交互
    │      │
    │      ├── 域随机化（每步修改物理参数）
    │      │
    │      ├── 策略推理（actor 网络）
    │      │
    │      └── 环境执行 step()
    │
    ├──→ 策略优化
    │      │
    │      ├── GAE 优势估计
    │      ├── PPO clip 更新
    │      └── 价值网络更新
    │
    ├──→ 定期回调
    │      │
    │      ├── progress_fn: 记录训练指标
    │      └── policy_params_fn: 保存检查点 + 导出 ONNX
    │
    └──→ 达到 num_timesteps → 完成

Runner 架构

训练系统通过 BaseRunner 抽象基类和 OpenDuckMiniV2Runner 具体类实现双层架构。

BaseRunner（`playground/common/runner.py`）

负责训练流程的通用部分：

组件	说明
输出目录管理	创建检查点和日志目录结构
TensorBoard 日志	使用 `SummaryWriter` 记录训练指标
检查点保存	使用 Orbax 保存 JAX 参数检查点
ONNX 导出	每次保存检查点时触发 ONNX 导出
PPO 训练配置	从 `locomotion_params` 加载配置
JAX 编译缓存	配置 GPU 编译缓存以加速后续训练

OpenDuckMiniV2Runner（`playground/open_duck_mini_v2/runner.py`）

负责机器人特定配置：

class OpenDuckMiniV2Runner(BaseRunner):
    def __init__(self, args):
        # 选择环境
        available_envs = {
            "joystick": (joystick, joystick.Joystick),
            "standing": (standing, standing.Standing),
        }
        self.env = env_class(task=args.task)
        self.eval_env = env_class(task=args.task)
        self.randomizer = randomize.domain_randomize
        self.action_size = self.env.action_size
        self.obs_size = self.env.observation_size["state"][0]

PPO 训练参数

训练参数从 Brax 的 locomotion_params 加载（基于 BerkeleyHumanoidJoystickFlatTerrain 配置），通过 brax_ppo_config() 获取。关键参数：

参数	说明
`num_timesteps`	总训练步数（默认 150M，可配置到 300M）
`network_factory.policy_hidden_layer_sizes`	策略网络隐层尺寸（默认 [256, 256]）
`network_factory.value_hidden_layer_sizes`	价值网络隐层尺寸
`num_envs`	并行环境数量
`batch_size`	批次大小
`unroll_length`	轨迹展开长度
`learning_rate`	学习率
`clip_epsilon`	PPO clip 参数
`entropy_cost`	熵正则化系数
`gae_lambda`	GAE lambda 参数
`reward_scaling`	奖励缩放

回调机制

progress_fn（进度回调）

每轮评估完成后调用，记录所有训练指标：

def progress_callback(self, num_steps, metrics):
    for metric_name, metric_value in metrics.items():
        self.writer.add_scalar(metric_name, metric_value, num_steps)
    
    print(f"STEP: {num_steps} reward: {metrics['eval/episode_reward']}")

记录的指标包括：

eval/episode_reward：评估回合奖励
eval/episode_reward_std：评估奖励标准差
loss/policy_loss：策略损失
loss/value_loss：价值损失
metrics/...：环境特定指标

policy_params_fn（策略参数回调）

定期保存检查点和导出模型：

def policy_params_fn(self, current_step, make_policy, params):
    # 1. 使用 Orbax 保存完整参数检查点
    orbax_checkpointer.save(path, params)
    
    # 2. 导出 ONNX 模型
    onnx_path = f"{output_dir}/{datetime}_{current_step}.onnx"
    export_onnx(params, self.action_size, self.ppo_params, self.obs_size, onnx_path)

数据流

┌────────────┐    ┌──────────────┐    ┌────────────┐
│ 并行环境     │    │ 域随机化      │    │ 策略网络     │
│ (num_envs) │───→│ domain_      │───→│ actor/critic│
│            │    │ randomize()  │    │            │
└────────────┘    └──────────────┘    └──────┬─────┘
                                             │
                                             ▼
              ┌─────────────────────────────────────┐
              │        经验回放缓冲区                  │
              │    (transitions: obs, action,        │
              │     reward, done, next_obs)          │
              └──────────────┬──────────────────────┘
                             │
                             ▼
              ┌─────────────────────────────────────┐
              │         PPO 更新                      │
              │  1. GAE 优势估计                       │
              │  2. 策略 clip 更新                     │
              │  3. 价值网络回归更新                    │
              └─────────────────────────────────────┘

训练命令

# 标准训练（摇杆任务，平坦地形 + 回差建模）
uv run playground/open_duck_mini_v2/runner.py \
    --task flat_terrain_backlash \
    --num_timesteps 300000000

# 站立任务训练
uv run playground/open_duck_mini_v2/runner.py \
    --env standing \
    --task flat_terrain \
    --num_timesteps 150000000

# 从检查点恢复训练
uv run playground/open_duck_mini_v2/runner.py \
    --task rough_terrain_backlash \
    --restore_checkpoint_path checkpoints/2025_01_15_123456_1000000 \
    --num_timesteps 300000000

训练技巧与最佳实践

从简单到复杂：先在平坦地形 + 无回差的简单配置上训练，再过渡到复杂场景
逐步增加难度：初始降低推力幅度，训练稳定后逐渐增加
使用回差建模：_backlash 任务能显著提高策略在真实机器人上的表现
合理配置时间步数：
- 简单站立：150M 步
- 摇杆控制：300M 步
- 回差建模：需要更多步数
观察 TensorBoard：关注奖励收敛趋势和成本项的变化

后续步骤

训练完成后：

找到最新导出的 .onnx 模型文件
使用 mujoco_infer.py 进行实时推理验证
在不同地形上测试策略的泛化能力