0. 简介

初次邂逅 WMP 这个术语时，您或许和我一样会联想到“视觉腿足领域又诞生了一个新概念”。但深入研读论文后会发现，WMP 并非新型视觉编码器，也不是简单地将 Dreamer 迁挪至机器人控制。它精准地抓住了“师生学习”范式中的真实痛点：老师与学生之间的信息天然不对等——scandots 永远无法描述 Tilt 这类需要精确距离的薄壁通道，也无法表达 Crawl 这种悬空障碍。换句话说，特权信息的表达能力存在固有上限。这一痛点不是通过“训练一个更大的学生模型”就能回避的，要么继续妥协，要么另辟蹊径。

回顾现有视觉腿足的主流方法，大致分为两类：一类是 Extreme Parkour 和 Legged Locomotion in Challenging Terrains 采用的 scandots 教师；另一类是 Robot Parkour Learning 将障碍几何作为特权信息，为每种地形单独训练一个教师。前者的瓶颈在于“scandots 表征能力有限”，后者则需要为每种地形独立训练教师，工程开销巨大。

WMP 的目标是绕开这堵墙：不再让学生模仿不完美的教师，而是让策略直接在一个统一的世界模型上生长出来。这里的关键是，世界模型将高维视觉历史压缩成具备预测能力的隐状态，策略获取的就不再是“scandots 的伪真实值”，而是“未来一段时间内会发生什么”的语义级汇总。

1.1 输入输出接口

WMP 将腿足控制建模为 POMDP 问题。机器人每一时刻的观测分为两部分：本体感知与第一视角深度图。底层状态中还包含仅在仿真中可见的特权信息，包括 scandots、足底接触力、随机化的物理参数。策略的输出是 12 维关节目标位置，再通过 PD 控制器转换为扭矩。

需要明确的是：策略实际接收的输入并非原始深度图，而是世界模型每隔若干步采样一次深度图后产出的循环隐状态。因此，策略的工作频率维持在 50Hz，而世界模型仅在 10Hz 上更新。

1.2 关键不对称：世界模型与策略的频率错位

直观上，可以把世界模型看作“慢一点的眼睛+大脑”，策略则充当“快一点的肌肉”。深度图采集本身就有约 40ms 的真机延迟，加上 RSSM 的计算开销，若让世界模型与策略同频运行，板载算力难以承受。WMP 的工程取舍是让 RSSM 每若干步执行一次，即每 100ms 更新一次，而策略在两次更新之间复用同一个隐状态。论文中对该参数从 2 扫描到 30，发现仿真中参数越小奖励越高，但在真机环境下，由于延迟不可忽略，参数为 5 时达到最优平衡点。这意味着 RSSM 的离散时间步并非简单的 0.02 秒，而是 0.1 秒。阅读代码时尤其需要注意这一点，否则容易算错 reward sum 和 batch_length。

从公式上看，四个 RSSM 组件分别对应 Recurrent、Encoder、Dynamic predictor、Decoder 四个角色，缺一不可。这种写法将世界模型拆解为“记忆 + 推断 + 预测 + 重建”四步流水线，是 Dreamer 系列在 Atari 与 DMC 上反复迭代后沉淀的稳定结构。WMP 的关键改造在于让循环步长跨越了 5 个策略时间步。

其中，GRU 实现循环模型，CNN 和 MLP 构成多模态编码器，它既预测下一时刻潜变量分布，也负责将潜变量解码回原始观测。需要厘清的是：prior 与 posterior 共享同一个解码头部，但在训练时通过 KL 项使两者相互靠近，从而将“有观测推断”与“无观测想象”绑定在同一个网络中。这种共享设计是 Dreamer 系列在 Atari 与 DMC 上反复验证有效的模式，WMP 直接将其应用于腿足任务，并通过将 num_actions 扩大 5 倍，来适配 100ms 的时间步。

2.1 为什么不直接用 ConvNet-RNN

传统的 student 使用 ConvNet-RNN：每帧深度图经过 CNN，再串联 GRU 输出 latent，最后用模仿损失与 teacher 对齐。核心问题在于，ConvNet-RNN 只是被动总结过去的视觉信号，缺乏“预测未来”的内在动力。因此，当 student 遇到未见过的地形时，latent 容易坍缩到一种“平均地形”上。RSSM 则恰恰相反，它具有显式的 prior 分布与 posterior 分布，并通过 KL 散度迫使两者靠近。这意味着，即便没有当前观测，模型也能仅凭 prior 展开未来若干步的“想象”，足以支撑机器人走过脚下未直接看到的区域。

2.2 RSSM 在 WMP 仓库里的真实形态

WMP 直接复用了 dreamerv3-torch 的 RSSM，但将 num_actions 改为 num_actions × update_interval，使得 dynamic 在 100ms 时间步内接收 5 个动作。这一改造看似只是参数缩放，但实际上将模型的“动作时间分辨率”从 50Hz 压低到 10Hz，与策略的运行频率显式解耦。下面这段 WorldModel 的训练步骤，需要注意它并不依赖 imagined rollout 做 RL，而是仅用真实仿真数据来训练世界模型，这与 Dreamer 的原始范式有本质差异：

# dreamer/models.py
def _train(self, data):
    data = self.preprocess(data)
    with tools.RequiresGrad(self):
        with torch.cuda.amp.autocast(self._use_amp):
            embed = self.encoder(data)
            post, prior = self.dynamics.observe(
                embed, data["action"], data["is_first"]
            )
            kl_loss, kl_value, dyn_loss, rep_loss = self.dynamics.kl_loss(
                post, prior, self._config.kl_free,
                self._config.dyn_scale, self._config.rep_scale,
            )
            preds = {}
            for name, head in self.heads.items():
                grad_head = name in self._config.grad_heads
                feat = self.dynamics.get_feat(post)
                feat = feat if grad_head else feat.detach()
                preds[name] = head(feat)
            losses = {n: -p.log_prob(data[n]) for n, p in preds.items()}
            scaled = {k: v * self._scales.get(k, 1.0) for k, v in losses.items()}
            model_loss = sum(scaled.values()) + kl_loss
        metrics = self._model_opt(torch.mean(model_loss), self.parameters())
    return {k: v.detach() for k, v in post.items()}, _, metrics

这段代码完成了三项任务：首先用 MultiEncoder 将本体感知与图像编码为 embed；然后让 RSSM.observe 在整个轨迹上一次性计算出 posterior post 与 prior prior；最后在 decoder 和 reward 两个头部上计算 NLL，加上 KL 项即获得模型损失。这里的关键是，grad_heads 默认只包含 decoder 和 reward，这意味着 cont head 在 WMP 中被禁用——四足任务不需要 termination 概率，因为 episode 的结束由仿真器判定，无需模型预测。

难点提示：KL_free 是怎么回事：训练 RSSM 时容易走向两个极端，要么 prior 与 posterior 对齐过度导致 posterior 坍缩，要么放任不管导致开环预测发散。kl_free=1.0 为 KL 项设置了一个“免罚区”，即每个时间步若 KL 散度小于 1 nat 则不再施加梯度。这相当于告诉模型“分布间距接近即可，不必完美对齐”。直觉上，就像驾校教练只关注你是否压线，而不在乎方向盘晃动的幅度。

3.1 为什么策略要拿确定性隐状态而不是随机隐状态

WMP 在论文中反复强调，策略获取的是 deterministic 部分而非 stochastic 部分。这里值得进一步说明：随机隐状态带有噪声采样，每次 forward 都会抖动，让策略去对齐它就相当于“对着随机数学开车”；而确定性隐状态是 GRU 的 deterministic state，包含了到目前为止所有压缩的视觉与本体历史，是真正稳定可靠的“世界摘要”。论文中使用 t-SNE 将六种地形下的确定性隐状态投射到二维平面，可以看到 Slope、Stair、Gap、Climb、Crawl、Tilt 之间有清晰的边界，仅 Slope 与 Climb 轻微重叠，这与“Climb 可视为 90 度的 Slope”这一直观感受相吻合。

3.2 ActorCriticWMP 的真实拼装方式

以下代码是 WMP 仓库中 ActorCriticWMP 的 act 路径。可以看到，策略的输入由三部分拼接而成：本体历史 latent、3 维 command、世界模型 latent。这种“三合一”设计让策略既能利用 RSSM 的视觉语义，又能保留 RMA 风格的 explicit context，是将两条路线优点融合到同一个 actor 中的关键写法。需要注意的是，actor 与 critic 的 wm_feature_encoder 参数相互独立，这对 PPO 的收敛至关重要：

# rsl_rl/modules/actor_critic_wmp.py
def act(self, observations, history, wm_feature, **kwargs):
    latent_vector = self.history_encoder(history)
    command = observations[:, self.privileged_dim + 6:self.privileged_dim + 9]
    wm_latent_vector = self.wm_feature_encoder(wm_feature)
    concat_observations = torch.concat(
        (latent_vector, command, wm_latent_vector), dim=-1,
    )
    self.update_distribution(concat_observations)
    return self.distribution.sample()

def evaluate(self, critic_observations, wm_feature, **kwargs):
    wm_latent_vector = self.critic_wm_feature_encoder(wm_feature)
    concat_observations = torch.concat(
        (critic_observations, wm_latent_vector), dim=-1,
    )
    return self.critic(concat_observations)

这段代码的工程价值在于，它将三个不同来源的信息分别压缩再拼接：history_encoder 是一个 MLP，接收 5 步本体感知历史并输出 latent，对应 RMA/HIM 那条线的 explicit context；wm_feature_encoder 是另一个 MLP，将 RSSM 输出的 1536 维 deterministic state 压缩到 32 维，供 actor 使用；critic_wm_feature_encoder 是 critic 专用的同结构 MLP，但参数独立。这三路 latent 再加上 3 维 command，构成 actor 的最终输入。

3.3 不对称 Actor-Critic 的细节

注意 evaluate 中，critic 接收的是 critic_observations，即 privileged_obs，包含 scandots 和 base linear vel。这种 asymmetric AC 方案沿袭自 Pinto et al., 2017。这里的关键是，论文强调，即便有了 scandots，critic 在 Tilt/Crawl 这类 scandots 表达不完整的地形中，仍需要依靠世界模型 latent 才能给出准确的 value 估计，否则 baseline 会偏移，PPO 的 advantage 也随之偏离。换句话说，世界模型 latent 不仅是 actor 的眼睛，也是 critic 的“补盲工具”。

4. 训练时的监督模块：仿真中同时学习世界模型与策略

4.1 单阶段训练的工程设计

WMP 与 student-teacher 最大的区别在于只有一个训练阶段。仓库中的 WMPRunner.learn 在每个迭代中同时执行三件事：用 PPO+AMP 训练 actor-critic、用监督回归训练 depth predictor、用 ELBO 训练世界模型。这种“同一个 episode 内，三模型并行训练”的写法看上去复杂，但工程上非常稳定。因为 RSSM 不参与 RL 的反传（策略的梯度通过 stop-gradient 截断），所以三个优化器之间不会相互干扰。

为什么单阶段比两阶段更强：student-teacher 两阶段训练的根本痛点是“上限被 teacher 限制”，因为 student 的模仿误差会累积。WMP 将世界模型嵌入到 PPO 的感知通路中，使世界模型的训练目标（重建过去观测、预测未来）与策略的训练目标（最大化 return）正交，但又共享同一份仿真数据。这意味着双方不会争抢数据，还能并行收敛。简单来说，将感知与决策放在同一个训练循环中，比串行的两阶段流水线更接近动物的学习方式。

# rsl_rl/runners/wmp_runner.py
for it in range(self.current_learning_iteration, tot_iter):
    with torch.inference_mode():
        for i in range(self.num_steps_per_env):
            if (self.env.global_counter % self.wm_update_interval == 0):
                wm_embed = self._world_model.encoder(wm_obs)
                wm_latent, _ = self._world_model.dynamics.obs_step(
                    wm_latent, wm_action, wm_embed, wm_obs["is_first"],
                )
                wm_feature = self._world_model.dynamics.get_deter_feat(wm_latent)
                wm_is_first[:] = 0
            history = self.trajectory_history.flatten(1).to(self.device)
            actions = self.alg.act(obs, critic_obs, amp_obs, history,
                                   wm_feature.to(self.env.device))
            obs, privileged_obs, rewards, dones, infos, *_ = self.env.step(actions)
        self.alg.compute_returns(critic_obs, wm_feature.to(self.env.device))
    mean_value_loss, mean_surrogate_loss, *_ = self.alg.update()
    if (sum_wm_dataset_size > self.wm_config.train_start_steps):
        if (it % self.depth_predictor_cfg["training_interval"] == 0):
            depth_mse_loss = self.train_depth_predictor()
        wm_metrics = self.train_world_model()

4.2 Depth Predictor 的角色

DepthPredictor 是 WMP 工程链路中非常关键但容易被忽视的一环。仿真中 4096 个并行 A1 实例，若每个都运行深度相机，IsaacGym 的 GPU 会直接崩溃。核心问题在于，WMP 给出了一个折中方案——仅让 1024 个 env 真正运行深度相机（camera_num_envs=1024），其余 env 则利用 forward_height_map 加上 ConvTranspose 解码器近似生成深度图，并通过监督学习使两路对齐。以下代码是 DepthPredictor.forward：

# rsl_rl/modules/depth_predictor.py
def forward(self, forward_heightmap, prop):
    x = torch.concat((forward_heightmap, prop), dim=-1)
    x = self.encoder(x)
    x = x.reshape(
        [-1, self.h_list[0], self.w_list[0],
         self._embed_size // (self.h_list[0] * self.w_list[0])]
    )
    x = x.permute(0, 3, 1, 2)
    x = self.layers(x)
    mean = x.permute(0, 2, 3, 1)
    if self._cnn_sigmoid:
        mean = F.sigmoid(mean)
    return mean

这段代码所做的是“将 525 维 forward heightmap + 33 维 prop 解码为 64×64×1 深度图”，监督信号来自真实相机渲染的深度图与该预测结果的 MSE。进一步来看，这是一个聪明的工程取舍：仿真中 forward heightmap 几乎是免费的（直接从地形网格采样），而相机渲染成本高昂。因此，让大部分 env 使用廉价的 forward heightmap 信号来训练 RSSM，关键 env 用真实深度图监督 depth predictor，最终所有 env 在 RSSM 阶段获得的都是“视觉化”的输入。

直觉理解：可以把 depth predictor 想象成一个“假摄像头”。它本身没有 RGB-D 的物理传感能力，但因为知道地形高度场和机器人姿态，它可以伪造出几何上正确的深度图。RSSM 并不关心深度图是真是假，只要分布一致即可。

5. 推理时的执行模块：板载推理的延迟约束

5.1 板载流水线

真机上的 WMP 部署在 Unitree A1 板载的 Jetson NX 上，没有外接计算盒。这意味着，整个 RSSM + actor 都需要在 NX 的 GPU 上运行。策略的输入数据流大致如下：Intel D435i 以 60Hz 输出 424×240 深度图 → 经过空间和时间滤波 → 中心裁剪并下采样至 64×64 → 送入 RSSM 编码器。整条链路包含了 100ms 的延迟。这一延迟在仿真训练时被显式建模：env 每隔 5 个 timestep 才将深度图送入 RSSM，并刻意附加 100ms 的滞后。

5.2 Sim-to-real 的关键超参

仓库 a1_amp_config.py 中 class depth 的几行配置直接写入了真机部署的契约，每一行都对应一个不可或缺的 sim-to-real 假设。若跨平台部署时不重新审视这些数字，整个 WMP 的延迟与视场假设都会失效，这也是开源用户最容易踩的坑之一。需要注意的是，配置中的每一项都是仿真与真机协商的结果，并非随意填写的默认值：

# legged_gym/envs/a1/a1_amp_config.py
class depth:
    use_camera = True
    camera_num_envs = 1024
    update_interval = 5  # 5 works without retraining, 8 worse
    original = (64, 64)
    resized = (64, 64)
    horizontal_fov = 58
    near_clip = 0
    far_clip = 2
    dis_noise = 0.0
    scale = 1
    invert = True

这里有几个值得关注的数值：horizontal_fov=58 与 D435i 的默认视场一致；far_clip=2 表示超过 2m 的深度会被截断，因为对于 A1 这种 0.25m 高的小型四足机器人来说，2m 之外的远景对 1 秒内的步态决策意义不大；update_interval=5 直接对应 RSSM 的频率，注释中特意标注了“5 不需要重训，8 就明显变差”，这是论文 ablation 的工程结论。

6. 训练目标：三套损失并存

6.1 世界模型与 PPO 的损失装配

WMP 的总损失可以拆分为三块叠加：世界模型 ELBO、策略 PPO+AMP、深度预测 MSE。三块损失分别由三个独立的优化器更新各自的参数，互不交叉。这里的关键是，RSSM 的梯度仅在 ELBO 内部流动，不会被 PPO 的 surrogate 扰动；反过来，PPO 对 wm_feature 也施加 stop-gradient，避免 RL 信号污染世界模型。三组公式如下：

第一块是 RSSM 的 ELBO，第二块是策略侧的 PPO + AMP 风格奖励（外加 vel-predict 辅助损失），第三块是 depth predictor 的 MSE。三者通过共享数据缓冲区交换信息，但梯度互不相通。这里的关键是，PPO 的 advantage 通过 compute_returns(critic_obs, wm_feature) 获取 critic 的 value，而 critic 同样接入了世界模型 latent，于是世界模型间接影响了所有三个损失，但梯度只在 ELBO 中反传。这种“数据共享、梯度隔离”的写法，是 WMP 单阶段训练能够稳定收敛的根本工程保障。

6.2 关键超参表

# dreamer/configs.yaml + a1_amp_config.py 节选
dyn_deter:        512
dyn_stoch:        32
dyn_discrete:     32
units:            512
kl_free:          1.0
dyn_scale:        0.5
rep_scale:        0.1
batch_size:       16
batch_length:     64
train_steps_per_iter: 10
train_start_steps:    10000
model_lr:         1e-4
update_interval:  5     # WMP / 100ms
camera_num_envs:  1024  # depth-truth envs
amp_reward_coef:  0.01  # 0.5 * 0.02
entropy_coef:     0.01
vel_predict_coef: 1.0

需要明确的是，几个有反差的取舍：batch_length=64 对应 6.4 秒的训练片段，论文 Figure 4 显示 6.4 秒是性能拐点，再短记忆不足，再长 RSSM 反传不稳定；amp_reward_coef=0.5×0.02=0.01 是 AMP 风格奖励的实际权重，小到不会盖过 tracking reward，但又足以让步态自然；train_start_steps=10000 让 RSSM 在策略开始训练前先离线吸收一波数据，避免冷启动阶段 policy 获得无意义的 latent。

7. 可选模式：t-SNE 可视化与开环预测

7.1 循环状态的可视化

论文 Section V-B 的 t-SNE 实验，将六种地形下的确定性隐状态投射到二维平面，可以看到不同地形的 cluster 几乎完美分离。这意味着，确定性隐状态不仅压缩了视觉历史，还隐式编码了“当前是在爬楼还是在过缝”这类语义信息。值得注意的是，策略不需要任何额外的地形分类头，仅凭确定性隐状态就能在 actor MLP 中学会“不同地形不同步态”的策略分支。这与 [Robot Parkour Learning] 必须为每种地形单独训练教师形成了鲜明对比。

解读：t-SNE cluster 是如何自然分离的：RSSM 训练时唯一的目标是“重建观测 + 预测未来”，没有任何显式的地形分类损失。但由于不同地形的视觉模式与本体响应序列差异显著（例如爬楼时的关节角度模式与穿缝时完全不同），世界模型必须将它们编码到 latent 空间的不同区域，才能有效压缩信息。换句话说，地形分离是 ELBO 的副产物，而非设计目标。这种“无监督学到结构”的现象，正是世界模型路线最优雅之处。

WMP 论文的另一项亮眼实验，是将仿真训练的 RSSM 拿到真机上做开环 rollout：给定真实的第一帧观测和真实的动作序列，让 RSSM 想象后续若干秒的深度图。结果显示，在 Crawl 任务中，模型预测的悬空横梁形状与真实形状并不完全一致（仿真中未见过这种形状），但“机器人能穿过的缝隙”的位置和角度高度吻合。换句话说，RSSM 并未学习“物体识别”，而是学到了“可通行性”这一关键属性，这才是 sim-to-real 平滑迁移的根本原因。

从公式上看，即便没有 reward head 主导训练（reward.loss_scale=0），decoder 重建出的深度图本身已是足够丰富的监督信号。这种“不依靠 reward 也能学习”的特性是 Dreamer 系列的标志，WMP 将其发挥得更加彻底——既然 RSSM 不参与策略梯度反传，那么 reward head 在 WMP 中实际上只是一个观测变量，保留它只是为了与原始 dreamerv3-torch 保持兼容。这意味着，如果有人想将 WMP 移植到没有 reward 信号的纯模仿场景中，直接移除 reward head 对训练效果几乎没有影响。

8. 总结

WMP 是 Dreamer 路线在视觉腿足领域的真正落地，而非 student-teacher 路线的横向变体。它的核心贡献并非“用世界模型替代 student”，而是“让世界模型同时承担视觉感知与未来预测两项职责”，从而绕过了 scandots 的表达力天花板，也规避了 ConvNet-RNN 缺乏预测压力的缺陷。