训练入口精读:scripts/train.py 如何把 Go2 任务真正跑起来

如果把整个仓库比作一条流水线,那么 train.py 就是总开关。
它不负责实现所有细节,但它决定了:

  • 用哪套配置训练
  • 用哪种环境
  • 用哪条算法分支
  • 训练结果保存到哪里
  • 最终会导出哪些可供评估和部署使用的产物

本文就只做一件事:顺着 scripts/train.py,讲清楚它是怎么一步步把 Go2 训练真正跑起来的。


一、train.py 到底做了什么?

train.py 压缩成伪代码,其实非常清楚:

def train_go2():
    import Cfg
    config_go2(Cfg)                     # 先装 Go2 默认配置
    override Cfg in train.py            # 再覆盖成当前训练主线的配置
    env = VelocityTrackingEasyEnv(...)  # 创建环境
    logger.log_params(...)              # 记录本次实验参数
    env = HistoryWrapper(env)           # 把单步 obs 包成 obs_history
    runner = Runner(env, device="cuda:0")   # 创建训练器
    runner.learn(...)                   # 开始 rollout + PPO 更新

从结构上看,它做的不是“定义算法”,而是“组装训练系统”。


二、配置覆盖顺序:Cfg 到底是怎么被改出来的?

理解 train.py,第一步不是看 PPO,而是看 Cfg

这个项目的配置不是 YAML 驱动,而是 Python 对象逐层覆盖。
最终训练配置来自三层叠加:

1. 基类默认配置

最底层是 legged_robot_config.py,里面定义了通用四足环境的默认值,比如:

  • observation 开关
  • reward scales
  • domain randomization
  • terrain
  • control
  • normalization

它可以理解成“全项目默认模板”。

2. Go2 专属默认配置

然后在 train.py#L19调用:

config_go2(Cfg)

这个函数在 go2_config.py 里,作用是把通用模板改成 Go2 版本,比如:

  • Go2 的 URDF
  • 默认关节角
  • Go2 的 PD 参数
  • terrain 默认行为
  • 默认 num_envs
  • 默认 observation 结构

所以这一步之后,Cfg 已经不是“通用四足配置”,而是“Go2 默认配置”。

3. train.py 再做当前实验覆盖

真正决定“这次训练学什么”的,是 train.py#L21到 train.py#L206这一大段覆盖。

因此最终配置关系可以写成:

C f g final = Override t r a i n ( config_go2 ( C f g base ) ) Cfg_{\text{final}} = \text{Override}_{train}\big(\text{config\_go2}(Cfg_{\text{base}})\big) Cfgfinal=Overridetrain(config_go2(Cfgbase))

这也是为什么读这个项目时,不能只看 go2_config.py,也不能只看 legged_robot_config.py,必须回到 train.py 看最后的覆盖结果。


三、train.py 里到底覆盖了哪些关键配置?

虽然覆盖项很多,但按功能可以分成五类。

1. 任务与课程设置

在 train.py#L21 开始,脚本先改 command curriculum 的离散化和阈值:

Cfg.commands.num_lin_vel_bins = 30
Cfg.commands.num_ang_vel_bins = 30
Cfg.curriculum_thresholds.tracking_ang_vel = 0.7
Cfg.curriculum_thresholds.tracking_lin_vel = 0.8
...
Cfg.commands.distributional_commands = True

这说明当前训练不是固定命令采样,而是带课程分布更新的 command curriculum。

2. sim2real 相关设置

在 train.py#L30 开始,脚本集中改了与 sim2real 相关的部分:

  • lag_timesteps = 6
  • control_type = "actuator_net"
  • 打开 friction / restitution / base mass / gravity / motor strength / motor offset 等随机化
  • 同时只保留 friction + restitution 进入 privileged obs

这一步很关键,因为它决定当前主线不是纯仿真 locomotion,而是带明显 sim2real 倾向的训练。

3. observation 与 privileged observation 结构

在 train.py#L78开始:

Cfg.env.num_privileged_obs = 2
Cfg.env.num_observation_history = 30
Cfg.commands.num_commands = 15
Cfg.env.num_observations = 70
Cfg.env.observe_two_prev_actions = True
Cfg.env.observe_gait_commands = True
Cfg.env.observe_clock_inputs = True

这一组覆盖几乎直接决定了网络结构:

  • 单步 obs 是 70 维
  • 历史窗口长度是 30
  • 所以 obs_history70 × 30 = 2100
  • privileged obs 是 2 维
  • actor / critic 输入最终都是 2102 维

4. terrain 与终止条件

在 train.py#L92开始,脚本把地形和 terminal 条件调成当前 locomotion 任务适配的版本,比如:

  • mesh_type = "trimesh"
  • terrain_width = 1.0
  • terrain_length = 1.0
  • teleport_robots = False
  • 开启 terminal body height / roll-pitch 终止

这说明当前任务不是平地纯 tracking,而是有真实地形与终止约束的 locomotion。

5. reward 与 command 空间

从 train.py#L118起,脚本集中重写了 reward scales 和 command ranges。

这部分虽然细节很多,但入口层面最重要的结论是:

  • 这不是沿用 Go2 默认 reward
  • 而是显式切到了一套更接近论文 locomotion 风格的 reward 组合

四、为什么环境先建出来,再包一层 HistoryWrapper

环境创建发生在 train.py#L208:

env = VelocityTrackingEasyEnv(sim_device='cuda:0', headless=headless, cfg=Cfg)

然后才在 train.py#L214 包一层:

env = HistoryWrapper(env)

这不是装饰器式的小改动,而是当前训练链路的必要步骤。

1. 原始环境只提供单步 observation

VelocityTrackingEasyEnv 最终继承自 LeggedRobot,环境本体负责:

  • 仿真 step
  • reward
  • obs 构造
  • privileged obs 构造

但它本身并不会自动把过去 30 步 obs 串起来。

2. HistoryWrapper 把单步 obs 改成三元组字典

在 history_wrapper.py#L18:

obs, rew, done, info = self.env.step(action)
privileged_obs = info["privileged_obs"]
self.obs_history = torch.cat((self.obs_history[:, self.env.num_obs:], obs), dim=-1)
return {'obs': obs, 'privileged_obs': privileged_obs, 'obs_history': self.obs_history}, rew, done, info

所以包完以后,算法拿到的就不再只是 obs,而是:

  • obs
  • privileged_obs
  • obs_history

3. 为什么这一步不可省?

因为当前 ppo_cse 分支里的网络结构就是围绕 obs_history 设计的:

  • adaptation module 吃 obs_history
  • actor 吃 obs_history + latent
  • critic 吃 obs_history + privileged_obs

如果不包 HistoryWrapper,后面的 ActorCritic 根本拿不到它需要的输入。

可以把它写成一个简单公式:

obs_history t = [ o t − H + 1 ; o t − H + 2 ; …   ; o t ] \text{obs\_history}_t = [o_{t-H+1}; o_{t-H+2}; \dots; o_t] obs_historyt=[otH+1;otH+2;;ot]

当前这里:

  • H = 30 H = 30 H=30
  • o t ∈ R 70 o_t \in \mathbb{R}^{70} otR70
  • 所以 o b s _ h i s t o r y t ∈ R 2100 obs\_history_t \in \mathbb{R}^{2100} obs_historytR2100

这也是为什么 HistoryWrapper 在这个项目里不是“方便加一下”,而是训练主线必需的一环。


五、为什么 Runner 明确走 ppo_cse/,而不是旧版 ppo/

这一点在 import 的第一眼就已经确定了。

在 train.py#L13到 train.py#L17:

from go2_gym_learn.ppo_cse import Runner
from go2_gym_learn.ppo_cse.actor_critic import AC_Args
from go2_gym_learn.ppo_cse.ppo import PPO_Args
from go2_gym_learn.ppo_cse import RunnerArgs

也就是说,这条训练主线从入口层面就已经选定了 ppo_cse 分支。

1. 这不是运行时自动选择,而是入口脚本显式指定

仓库里确实同时保留了:

  • ppo
  • ppo_cse

train.py 并没有做条件分支,而是直接 import ppo_cse

所以“为什么走 ppo_cse”最直接的答案就是:
因为当前训练入口就是按 ppo_cse 写的。

2. 从代码意图看,ppo_cse 才是当前 RMA / adaptation 主线

在 ppo_cse/init.py#L43里:

class RunnerArgs(PrefixProto, cli=False):
    algorithm_class_name = 'RMA'

而旧版 ppo/init.py#L47 里是:

algorithm_class_name = 'PPO'

这至少说明从命名和组织上,ppo_cse 更接近当前这套:

  • asymmetric actor-critic
  • adaptation module
  • history-based student policy

的训练主线。

3. 入口脚本的配置也和 ppo_cse 设计完全匹配

比如:

  • num_privileged_obs = 2
  • num_observation_history = 30
  • HistoryWrapper
  • sim2real 式 domain rand
  • adaptation supervision

这些都不是最传统的单帧 PPO 配置,而明显是在为 ppo_cse 这一套网络和训练逻辑服务。

因此更完整的说法是:
不是 Runner 在运行时“决定”选 ppo_cse,而是 train.py 从一开始就把当前实验定义成了 ppo_cse 这条分支上的训练任务。

历史长度设为 30,是一个经验上的折中:窗口太短,时间信息不足,难以稳定辨识 friction/restitution ;窗口太长,输入维度过大、训练更慢、优化更难。


六、Runner 在这里到底接管了什么?

当 train.py#L216执行:

runner = Runner(env, device=f"cuda:{gpu_id}")
runner.learn(num_learning_iterations=100000, init_at_random_ep_len=True, eval_freq=100)

控制权就从入口脚本转移到 ppo_cse/init.py。

Runner 在初始化时主要做了三件事:

1. 构建 ActorCritic

见 ppo_cse/init.py#L70:

actor_critic = ActorCritic(
    self.env.num_obs,
    self.env.num_privileged_obs,
    self.env.num_obs_history,
    self.env.num_actions,
).to(self.device)

这里网络输入维度直接来自环境,不是手写常数。

2. 构建 PPO 算法对象

见 ppo_cse/init.py#L93:

self.alg = PPO(actor_critic, device=self.device)

3. 初始化 rollout storage

见 ppo_cse/init.py#L96:

self.alg.init_storage(
    self.env.num_train_envs,
    self.num_steps_per_env,
    [self.env.num_obs],
    [self.env.num_privileged_obs],
    [self.env.num_obs_history],
    [self.env.num_actions]
)

因此 train.py 自己并不直接写任何 PPO 循环。
它只负责组装;真正的 rollout、GAE、PPO update、adaptation update 都由 Runner.learn() 接管。


七、logger.log_params(...) 为什么会生成 parameters.pkl

这是入口脚本里一个很容易被忽略、但对复现实验和部署都很关键的点。

在 train.py#L210:

logger.log_params(
    AC_Args=vars(AC_Args),
    PPO_Args=vars(PPO_Args),
    RunnerArgs=vars(RunnerArgs),
    Cfg=vars(Cfg)
)

vars() 是 Python 内置函数。

它通常返回对象的 dict,也就是“这个对象当前有哪些属性”。

比如:

vars(AC_Args)

就会把 AC_Args 类里的参数变成一个字典,大概像这样:

{
    'init_noise_std': 1.0,
    'actor_hidden_dims': [512, 256, 128],
    'critic_hidden_dims': [512, 256, 128],
    ...
}

它记录了四类信息:

  • AC_Args
  • PPO_Args
  • RunnerArgs
  • Cfg

也就是说,这一行把“模型结构参数 + PPO 超参数 + runner 参数 + 环境配置”一起打包记录下来。

1. 为什么它会落到 runs/ 目录里?

因为在 train.py#L225 到 train.py#L227,主函数先做了:

logger.configure(
    logger.utcnow(f'gait-conditioned-agility/%Y-%m-%d/{stem}/%H%M%S.%f'),
    root=Path(f"{MINI_GYM_ROOT_DIR}/runs").resolve(),
)

这里:

  • root 指向 runs/
  • stem 是脚本名 train
  • 时间戳决定实验子目录

所以一次训练会被记录到类似:

runs/gait-conditioned-agility/2026-01-15/train/011500.160115/

这样的目录下。

2. 为什么说它最终会生成 parameters.pkl

虽然 ml_logger 的内部实现不在这个仓库里,但从本仓库的运行产物和下游使用方式可以非常明确地反推出这一点:

  • runs/.../parameters.pkl 实际存在
  • play.py#L39 会读取它
  • deploy_policy.py#L22 也会读取它

因此在这个项目里,logger.log_params(...) 的工程意义就是:
把当前实验的完整配置快照落盘,供后续评估脚本和部署脚本重新还原训练环境。

3. 为什么这一步这么重要?

因为部署时并不只需要网络权重,还需要知道:

  • command 空间怎么定义
  • observation 开关怎么开
  • action scale 是多少
  • gait 参数范围是什么
  • reward / terrain / domain rand 训练时是什么样

这些都不在 .jit 里,而是在 parameters.pkl 里。

所以从“能不能复现训练”和“能不能正确部署”两个角度看,logger.log_params(...) 都不是附属操作,而是主链条的一部分。


八、为什么 train.pylog_params,再包 HistoryWrapper

这是一个细节,但很值得讲。

代码顺序是:

env = VelocityTrackingEasyEnv(...)
logger.log_params(...)
env = HistoryWrapper(env)
runner = Runner(...)

这样写的好处是:

  • 先记录原始实验配置
  • 再做算法侧需要的包装
  • 保证落盘的是“真实实验参数”,而不是某个 wrapper 的运行时派生状态

而且 HistoryWrapper 的核心行为来自 Cfg.env.num_observation_history,这个参数本身已经被 logger.log_params 记录进去了,所以不会丢失历史窗口信息。


九、train.py 最后到底把一场训练启动成了什么样子?

把入口脚本、环境、wrapper、runner 合在一起,可以得到当前训练主线的完整逻辑:

1. 配置层

先从:

  • 基类默认配置
  • Go2 默认配置

出发,再由 train.py 覆盖成当前实验的最终 Cfg

2. 环境层

创建 VelocityTrackingEasyEnv,也就是基于 LeggedRobot 的 Go2 速度跟踪环境。

3. 输入组织层

HistoryWrapper 把单步 observation 组织成:

  • obs
  • privileged_obs
  • obs_history

其中:

obs_history ∈ R 30 × 70 = R 2100 \text{obs\_history} \in \mathbb{R}^{30 \times 70} = \mathbb{R}^{2100} obs_historyR30×70=R2100

4. 算法层

构建 ppo_cse.Runner,内部再构建:

  • ActorCritic
  • PPO
  • RolloutStorage

5. 日志与产物层

训练前先用 logger.configure(...) 确定 run 目录,再用 logger.log_params(...) 记录实验快照,最终形成:

  • parameters.pkl
  • checkpoint .pt
  • 供部署使用的 .jit

十、总结

scripts/train.py 的价值不在于它实现了多少复杂算法,而在于它把这个项目真正“接通了”。它通过 config_go2(Cfg) 先装载 Go2 默认配置,再在入口脚本内集中覆盖 observation、command、reward、terrain、domain randomization 与 sim2real 相关设置,最终得到当前实验的完整 Cfg。随后,它创建 VelocityTrackingEasyEnv,再通过 HistoryWrapper 把原本只提供单步 observation 的环境包装成同时返回 obsprivileged_obsobs_history 的接口,这一步直接决定了后续 adaptation module、actor 和 critic 的输入形式。算法侧,train.py 并没有走旧版 ppo/,而是显式导入 go2_gym_learn.ppo_cseRunnerAC_ArgsPPO_ArgsRunnerArgs,因此当前训练主线从入口层面就被定义成了一条带 obs_history、privileged learning 和 adaptation supervision 的 RMA 风格分支。最后,入口脚本在真正开始训练前调用 logger.configure(...)logger.log_params(...),把本次实验的算法参数与环境配置快照记录到 runs/ 目录下;结合仓库已有的运行产物以及评估、部署脚本对 parameters.pkl 的读取逻辑,可以确认这一步正是后续复现实验和实机部署所依赖的配置落盘环节。换句话说,train.py 不是算法细节的堆砌,而是整个 Go2 训练系统的总装入口:它把配置、环境、历史观测、算法分支、日志与产物管理一次性接成了一条可训练、可评估、可部署的完整链路。

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐