预测潜在世界模型的闭环性能：LunarLander中非马尔可夫奖励下MPC和基于模型的强化学习的离线检查点选择

# 预测潜在世界模型的闭环性能：面向月球着陆器中非马尔可夫奖励的离线检查点选择用于MPC和基于模型的强化学习

来源：https://arxiv.org/html/2607.01736

###### 摘要

我们研究如何仅凭验证时诊断指标预测学到的潜在世界模型的下游闭环性能。从世界模型训练运行中选择合适的检查点非常困难：验证损失和多步预测RMSE在闭环性能崩溃后仍会持续改善。我们提出了一套源自最优控制理论的结构性验证时诊断指标，并将其应用于具有成形奖励的Gymnasium LunarLander-v3环境。我们训练了一个RSSM\[5 (https://arxiv.org/html/2607.01736#bib.bib4),4 (https://arxiv.org/html/2607.01736#bib.bib5)\]世界模型，并将每个检查点的CEM-MPC回报视为闭环质量的真实标准。通过评估40个指标与该真实标准的关系，我们发现最强的单一预测指标是奖励可观测性分数（ROF），它衡量奖励预测器对可观测子空间的依赖程度。我们将ROF与三个结构正则化器组合成一个单一的离线检查点选择分数——复合奖励可观测性分数（CROF）。CROF选择的世界模型训练出一个基于模型的A2C策略，该策略在公平比较中比无模型的A2C基线高出约24.5个回报点，同时使用约65倍更少的真实环境交互；同样的世界模型也驱动了一个强大的零样本CEM-MPC策略。代码与数据：https://github.com/nsmoly/LunarLander_RSSM。

## 1 引言

基于模型的强化学习承诺通过构建环境的预测模型——即世界模型——并在想象中用于规划或策略优化，从而减少学习有效策略所需的真实环境交互次数。Dreamer或循环状态空间模型（RSSM）架构\[5 (https://arxiv.org/html/2607.01736#bib.bib4),4 (https://arxiv.org/html/2607.01736#bib.bib5)\]已成为世界模型的标准骨干，既能通过模型预测控制（MPC）\[5 (https://arxiv.org/html/2607.01736#bib.bib4)\]实现在线规划，也能通过想象展开\[4 (https://arxiv.org/html/2607.01736#bib.bib5),6 (https://arxiv.org/html/2607.01736#bib.bib6)\]使用强化学习（RL）进行策略学习。

图1：LunarLander-v3：四个离散动作（无操作、左/主/右推进器），8维观测，从随机初始状态出发的软着陆目标。

基于模型的强化学习中的一个基本挑战是目标不匹配\[8 (https://arxiv.org/html/2607.01736#bib.bib8)\]：世界模型的训练目标（预测精度）不一定与下游目标（闭环性能）对齐。一个具有出色一步或多步预测精度的模型，如果其学到的潜在动态缺乏这些用途所需的结构特性，仍可能不适合规划或策略优化。

在本工作中，我们专注于LunarLander-v3 Gymnasium环境。它在先前工作中并未成为详细的RSSM/Dreamer式分析的对象，同时结合了两个有用的特性：8维观测和4个离散动作，两者都足够小，允许对学到的潜在进行直接结构分析；以及——这正是本文的动机所在——一个由于基于势能的成形奖励加上由模拟器内部标志触发的终止事件而无法仅从当前观测和动作完全恢复的奖励（第4.8节）。因此，LunarLander作为更广泛类别环境（具有非马尔可夫或成形奖励）的可处理代理，在这些环境中，标准的一步预测指标被证明是糟糕的检查点选择标准。我们开发了借鉴经典控制理论（可控性、可观测性、Kalman\[7 (https://arxiv.org/html/2607.01736#bib.bib1)\]）的结构诊断指标，并通过实验表明，它们能够识别出那些比标准模型选择标准产生更好下游策略的世界模型检查点，无论是用于MPC还是在想象中训练的基于模型的演员-评论家策略。

我们在本工作中的贡献如下：

1. 一套40个结构性的验证时指标，用于世界模型评估，借鉴经典控制理论：基于雅可比的可控性/可观测性分析（固定和时变线性化）以及三个奖励/观测子空间对齐分数。
2. 复合奖励可观测性分数（CROF），一个单一的离线检查点选择分数，结合了奖励/观测子空间对齐分数与三个结构分数（可控性秩、可观测性秩、开环观测误差）。
3. 在LunarLander-v3上的实验证据表明，CROF能够跟踪零样本CEM-MPC回报和下游A2C训练质量，从而能够离线选择更好的世界模型检查点用于闭环使用，而标准选择标准（验证损失、多步RMSE、敏感性）则选择了具有较差闭环性能的过拟合检查点。
4. 一个数据效率比较，表明CROF选择的世界模型训练出的基于模型的A2C策略，在公平比较中比无模型的A2C基线使用约65倍更少的真实环境交互。

## 2 相关工作

我们的工作借鉴了三个领域：用于控制的世界模型、基于模型的强化学习中的目标不匹配问题，以及动态系统的经典控制理论分析。

### 2.1 世界模型与基于模型的强化学习

RSSM系列——PlaNet\[5 (https://arxiv.org/html/2607.01736#bib.bib4)\], Dreamer\[4 (https://arxiv.org/html/2607.01736#bib.bib5)\], 和 DreamerV3\[6 (https://arxiv.org/html/2607.01736#bib.bib6)\]——学习用于规划和基于想象的演员-评论家训练的潜在动态；PETS\[2 (https://arxiv.org/html/2607.01736#bib.bib7)\]将概率集成动态与MPC相结合，实现了强大的样本效率。所有这些方法都通过最大化预测似然进行训练，并通过下游任务性能进行评估，而不分析学到的动态中决定规划质量的结构特性。

### 2.2 目标不匹配

Lambert等人\[8 (https://arxiv.org/html/2607.01736#bib.bib8)\]指出了目标不匹配问题：为一阶预测精度优化动态模型并不能保证良好的下游闭环性能，并提出了重新加权训练目标作为初步缓解措施。我们的工作提出了一个补充性的解释，说明这种不匹配为何发生：预测精度（即使多步）捕捉的是统计拟合，而规划质量取决于奖励相关、可控和可观测子空间之间的结构对齐——这是一种训练损失无法衡量的几何属性。

### 2.3 任务相关表示

几项工作学习在训练时保留任务相关信息的表示：DeepMDP\[3 (https://arxiv.org/html/2607.01736#bib.bib9)\]通过基于双模拟的奖励和转移结构损失，SOLAR\[14 (https://arxiv.org/html/2607.01736#bib.bib11)\]通过鼓励局部线性动态用于从图像进行LQR规划，以及Agarwal等人\[1 (https://arxiv.org/html/2607.01736#bib.bib12)\]通过行为相似性的对比嵌入。这些方法修改了训练目标。我们的工作是补充性的：我们提供事后验证时诊断指标，用于量化所学结构是否支持闭环规划，而与训练过程无关。

### 2.4 学习系统的控制理论分析

可控性和可观测性源于Kalman\[7 (https://arxiv.org/html/2607.01736#bib.bib1)\]，他展示了相应矩阵的秩决定了线性系统能否被完全控制和观测；Moore\[9 (https://arxiv.org/html/2607.01736#bib.bib2)\]通过平衡实现理论扩展了这一概念，识别出同时可控且可观测的维度。Sussillo和Barak\[13 (https://arxiv.org/html/2607.01736#bib.bib10)\]将局部雅可比和不动点分析应用于神经科学中训练过的RNN，揭示了高维循环网络中的可解释低维动态。我们将这些工具改编到RSSM设置中，并将它们组合成标量奖励/观测对齐分数（RCF、OCF、ROF），用作无需环境访问即可预测闭环性能的事后验证时诊断指标。

## 3 方法

### 3.1 世界模型架构

我们使用一个受RSSM启发的世界模型\[4 (https://arxiv.org/html/2607.01736#bib.bib5),6 (https://arxiv.org/html/2607.01736#bib.bib6)\]，针对LunarLander-v3进行了调整，其观测为\(p=8\)维（6个连续物理维度加上2个二进制腿接触指示器），动作空间为\(K=4\)个离散动作。RSSM核心维护一个确定性隐藏状态\(\bm{h}_t \in \mathbb{R}^{d_h}\)（\(d_h=256\)，单层GRU）和一个随机潜在变量\(\bm{z}_t \in \mathbb{R}^{d_z}\)（\(d_z=16\)），给出完整的潜在状态\(\bm{s}_t = [\bm{h}_t; \bm{z}_t] \in \mathbb{R}^{272}\)。动作\(a_t \in \{0,1,2,3\}\)被独热编码为\(\bm{u}_t \in \mathbb{R}^{4}\)。RSSM的转移（先验）是：

\[
\bm{s}_{t+1} = f_\theta(\bm{s}_t, \bm{u}_t) = \bigl[ \operatorname{GRU}(\bm{h}_t, [\bm{z}_t; \bm{u}_t]); \; \bm{\mu}_{\text{prior}}(\bm{h}_{t+1}) \bigr],
\tag{1}
\]

后验编码器是\(q_\theta(\bm{z}_t \mid \bm{h}_t, \bm{o}_t)\)，先验和后验都是2层MLP（LayerNorm + SiLU），预测对角高斯分布的均值和log-std；观测编码器是一个2层MLP，将\(\bm{o}_t \in \mathbb{R}^8\)映射为特征向量，并与\(\bm{h}_t\)拼接后输入后验。

##### 具有异构头的结构化解码器。

解码器并非将所有8个观测维度等同处理，而是使用一个共享的2层MLP骨干（LayerNorm + SiLU，输入\(d_h + d_z = 272\)）连接四个头：一个MSE回归的物理头（\(\mathbb{R}^6\)，连续状态）、一个BCE接触头（\(\mathbb{R}^2\) logits，推理时使用sigmoid）、一个BCE done头（\(\mathbb{R}^1\)），以及一个独立的3层奖励MLP（\(\mathbb{R}^{272} \to 256 \to 256 \to 1\)），它直接操作于\([\bm{h}; \bm{z}]\)而非共享骨干，因此奖励预测器能够发展自己的表示，而不会与重构争夺骨干容量。观测解码器为\(\hat{\bm{o}}_t = g_\theta(\bm{s}_t)\)（物理值拼接接触logits的sigmoid）；奖励预测器为\(\hat{r}_t = \rho_\theta(\bm{s}_t)\)。

##### 训练损失。

总损失是一个加权和：

\[
\mathcal{L} = w_r \mathcal{L}_{\text{recon}} + w_\rho \mathcal{L}_{\text{rew}} + w_d \mathcal{L}_{\text{done}} + w_{\text{KL}} \mathcal{L}_{\text{KL}},
\tag{2}
\]

权重为\(w_r = 1.0\), \(w_\rho = 1.2\), \(w_d = 0.5\), \(w_{\text{KL}} = 1.0\)（KL自由位下限\(\beta_{\text{KL}} = 0.5\)）。\(\mathcal{L}_{\text{recon}}\)结合了物理的MSE和接触的BCE；\(\mathcal{L}_{\text{rew}}\)是平方奖励误差；\(\mathcal{L}_{\text{done}}\)是done标志的BCE；\(\mathcal{L}_{\text{KL}} = \text{KL}(q_\theta \| p_\theta)\)是后验-先验KL散度。

### 3.2 通过交叉熵方法（CEM）的MPC

由于LunarLander-v3具有离散动作空间（4个动作：无操作、左推进器、主发动机、右推进器），标准的基于梯度的规划不适用。我们使用适用于离散动作的交叉熵方法（CEM）\[11 (https://arxiv.org/html/2607.01736#bib.bib3)\]：在每个真实时间步，CEM迭代地细化一个规划视界\(H=25\)步上每一步的分类分布。每次迭代从当前分布中采样\(N=384\)个候选动作序列，在世界模型先验中开环地展开每个序列，同时累积折扣奖励\(\sum_{k=1}^H \gamma^{k-1} \hat{r}_{t+k}\)（\(\gamma=0.97\)），保留前\(E=48\)个精英，并以平滑因子\(\alpha=0.7\)将每步的logits向精英动作频率更新。经过\(I=4\)次这样的迭代后，执行最佳序列的第一个动作，并在下一个真实时间步重复该过程。规划视界\(H=25\)与用于指标评估的展开视界相匹配，因此离线指标恰好衡量了CEM运行的机制。

### 3.3 通过潜在想象的演员-评论家策略

遵循Dreamer\[4 (https://arxiv.org/html/2607.01736#bib.bib5)\]，我们在世界模型的潜在想象中完全训练一个演员-评论家策略（A2C）。演员和评论家是紧凑的3层MLP（隐藏维度256），读取潜在状态\([\bm{h}; \bm{z}]\)；演员输出一个4个动作上的分类分布，评论家预测标量价值估计。每个训练步骤从回放数据集中采样一个真实的观测-动作序列，将RSSM后验在真实数据上预热5步，然后通过使用演员采样的动作（无观测校正）展开先验来想象接下来的15步。价值目标是\(\lambda\)-回报\[12 (https://arxiv.org/html/2607.01736#bib.bib14)\]（\(\gamma=0.99\), \(\lambda=0.95\)），从世界模型的奖励头计算得到；演员通过带优势裁剪的REINFORCE更新，评论家通过MSE更新，熵在训练期间从0.5衰减至0.15。我们以学习率\(3 \times 10^{-5}\)和批次大小256训练了1000个epoch。关键是，这个阶段使用了零额外的真实环境交互：所有学习都在学到的世界模型上发生，使用的正是用于训练它的相同离线数据集。

### 3.4 世界模型质量指标

我们在四个类别中定义指标；完整的按指标列表及其与MPC回报的经验相关性在表1（参见第4.5节）中汇总。

#### 3.4.1 验证损失分解

五个标准训练诊断指标：val_loss（总加权验证损失，如(2)式）、其组成部分val_kl（后验-先验KL）和val_recon（物理+接触重构）、以及一步后验RMSE：post_obs_rmse和post_rew_rmse，它们衡量带有编码器校正的闭环预测质量。

#### 3.4.2 多步开环展开

这些指标衡量CEM MPC在规划期间经历的情况：无编码器校正的多步预测质量。在后验预热5步后，RSSM先验使用真实动作展开25步，平方误差在有效（非填充、未终止前）位置掩码，并对平方根进行降阶。我们记录九个量：观察RMSE在展开的第一步（ol_obs_start），整个视界的平均值（ol_obs_avg），每个展开的最大值（ol_obs_max）