跳转至

VIGOR: Zero-Shot Visual Generalization via Latent-Space Consistency in Model-Based Reinforcement Learning

会议: NeurIPS 2026
arXiv: 2610.02801
OpenReview: https://openreview.net/forum?id=MKDlrS5Ftr
领域: 强化学习
关键词: 基于模型的强化学习, 零样本视觉泛化, 潜空间一致性, 数据增强, TD-MPC2

一句话总结

针对基于模型的强化学习在未见视觉扰动下因误差多步递归累积而崩溃的难题,VIGOR 通过非对称弱到强数据增强、动力学级一致性约束与编码器级稳定化,实现了无需测试时微调的高鲁棒零样本视觉泛化。

研究背景与动机

基于视觉输入的强化学习(Visual RL)旨在直接从像素观测中学习连续控制策略,在机器人操作与运动控制中取得了显著成功。基于模型的强化学习(MBRL,如 TD-MPC2、Dreamer)通过在紧凑的潜空间中学习动力学模型并执行时域规划,展现出了远超无模型强化学习(MFRL)的采样效率。然而,现有视觉 MBRL 算法高度依赖训练分布与测试分布独立同分布的假设;一旦测试环境出现背景杂波、光照突变、相机视角轻微抖动等未见的视觉干扰,策略性能往往出现灾难性滑坡。

在无模型 RL 中,视觉干扰主要通过数据增强或对比表征学习缓解,且编码器的扰动通常仅影响单步的 Q 值或动作预测。但对于基于模型的 RL,视觉泛化面临独有的两级脆弱性(two-level vulnerability):第一级,测试阶段的未见视觉干扰首先使观测编码超出训练分布(OOD);第二级,这些偏移的隐变量作为初始输入被送入动力学模型,并在规划时域 \(H\) 内进行多步递归展开,导致预测误差沿时间轴迅速发散累积,最终摧毁整个 MPC 规划过程。直接把 MFRL 的强数据增强策略全量迁移到 MBRL 会破坏世界模型的预测稳定性,使智能体性能下降超过 50%。

为了在保持 MBRL 极高采样效率的同时切断两级误差传播链,本文的核心思路是约束世界模型在扰动下的潜空间轨迹一致性。核心 idea:通过非对称弱到强批量划分提供稳定锚点,以动力学级直接隐变量回归强行对齐递归轨迹,并借助编码器级表征对齐抑制交叉扰动漂移,构建扰动无关的鲁棒潜空间世界模型。

方法详解

整体框架

VIGOR 以无重构潜空间 MBRL 代表框架 TD-MPC2 为骨干网络,通过三个相互协同的组件解决两级脆弱性问题:非对称弱到强增强(AW)在批次内构造配对的弱扰动与弱到强复合扰动视图;动力学级一致性(DC)将混合隐变量递归前向推演,并将多步轨迹强制对齐至弱扰动生成的目标潜状态;编码器级稳定化(ES)则防止共享编码器在跨增强监督压力下发生表征畸变。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入观测批次 s_t"] --> B["非对称弱到强增强 AW<br/>批次等分: 弱扰动 s_t^w 与复合强扰动 s_t^ws"]
    B --> C["共享编码器 h_θ 编码<br/>混合隐变量 z_t^m = [z_t^w, z_t^ws]"]
    C --> D["编码器级稳定化 ES<br/>余弦对齐 z_t^ws 与 stop-grad(z_t^w)"]
    C --> E["动力学级一致性 DC<br/>混合隐变量递归递推 z_hat_{t+h}^m (h=1..H)"]
    E --> F["对齐目标 z_{t+h}^w<br/>强行约束递推轨迹与弱增强锚点一致"]
    D --> G["TD-MPC2 任务目标<br/>奖励预测 + Q值估计 + MPC规划执行"]
    F --> G

关键设计

1. 非对称弱到强增强(AW):兼顾训练稳定性与扰动多样性的批次划分 在 MBRL 中若全量施加如随机贴图(random-overlay)等强增强,会导致编码空间剧烈震荡,使得递归动力学模型难以收敛;若仅使用随机平移(random-shift)等弱增强,则无法覆盖真实测试环境中严苛的视觉域偏移。AW 将单一 mini-batch \(\mathcal{B}\) 随机均匀划分为互斥的弱增强子集 \(\mathcal{I}^w\) 和强增强子集 \(\mathcal{I}^{ws}\)(默认比例 \(\zeta = |\mathcal{I}^w|/|\mathcal{I}^{ws}| = 1\))。弱增强分支仅施加保留主要几何信息的随机平移 \(\tau^w\);强增强分支则施加弱到强的组合变换 \(\tau^s \circ \tau^w\)(在平移基础上叠加密集纹理/图像贴图)。两路图像在批次维度拼接后输入共享编码器 \(h_\theta\),既为后续动力学对齐提供了未经严重破坏的低噪声基准,又引入了高强度的扰动样本,保证总批大小与计算开销不变。

2. 动力学级一致性(DC):递归预测轨迹向弱增强锚点对齐 针对未见扰动在多步递归推演中发散放大的核心痛点,DC 强迫动力学模型学会忽略视觉干扰带来的表征漂移,直接输出任务内在的真实物理状态转移。在时域长度为 \(H\) 的轨迹段起始时刻 \(t\),动力学模型以 AW 拼接得到的混合表征 \(z_t^m\) 作为初始状态 \(z_t^m = \hat{z}_t^m\)。在后续各步 \(h=1,\dots,H\),动力学模型递归迭代推演状态 \(\hat{z}_{t+h}^m = d_\theta(\hat{z}_{t+h-1}^m, a_{t+h-1})\)。DC 损失定义为推演潜状态与目标后继状态弱增强编码之间的 \(\ell_2\) 回归误差: $\(\mathcal{L}_{\text{DC}}(\theta) = \sum_{h=1}^H \left\| d_\theta(\hat{z}_{t+h-1}^m, a_{t+h-1}) - \operatorname{sg}(z_{t+h}^w) \right\|_2^2\)$ 其中 \(\operatorname{sg}(\cdot)\) 表示阻断梯度算子,\(z_{t+h}^w = h_\theta(\tau^w(s_{t+h}, \upsilon^w))\) 为对应的弱增强编码。通过将单步与长时域多步预测的目标均锚定在弱增强轨迹上,动力学模型被强力正则化,确保受扰动状态在展开过程中始终被拉回任务关键的流形上。

3. 编码器级稳定化(ES):防止跨增强监督导致共享表征畸变 在 TD-MPC2 架构中,唯一的图像编码器 \(h_\theta\) 必须同时为动力学模型、Q 值网络以及在线采样规划器(MPC)提供状态输入。由于 DC 损失在动力学层面对强弱增强隐变量施加了单向的跨增强回归梯度,这极易导致编码器发生表征漂移(encoder drift),损害静态特征判别力。ES 引入对称特征规范化对齐损失,在每个时间步约束同一帧画面在弱到强增强视角下的编码与弱增强视角下的编码保持余弦相似度接近: $\(\mathcal{L}_{\text{ES}}(\theta) = \left\| \frac{z_t^{ws}}{\|z_t^{ws}\|_2} - \frac{\operatorname{sg}(z_t^w)}{\|\operatorname{sg}(z_t^w)\|_2} \right\|_2^2\)$ 该设计阻止了强增强梯度逆向破坏特征流形,构成了阻断表征漂移的第一道防线,使 DC 能够在干净稳定的潜特征基础上专注于动力学转移一致性的拟合。

损失函数 / 训练策略

VIGOR 的总训练目标联合了 TD-MPC2 原生损失与新增的一致性约束。对于从重放缓冲区采样的轨迹段 \(\Gamma = (s_i, a_i, r_i, s_{i+1})_{i=t}^{t+H}\),总体损失为: $\(\mathcal{L}_{\text{total}}(\theta) = \mathbb{E}\left[ \sum_{i=t}^{t+H} \lambda^{i-t} \left( \mathcal{L}_{\text{rew}}(\theta) + \mathcal{L}_Q(\theta) + \mathcal{L}_{\text{DC}}(\theta) \right) + \mathcal{L}_{\text{ES}}(\theta) \right]\)$ 其中 \(\lambda \in (0, 1]\) 为时间衰减因子。由于新增模块全部通过辅助损失发挥作用,在策略推理和在线 MPC 规划阶段不引入任何额外结构或时间开销。

实验关键数据

主实验

在遵循 RL-ViGen 协议评测的 DeepMind Control Suite(6 个连续控制任务)和 Robosuite(4 个机器人操作任务)上,涵盖 7 种未见视觉扰动类型(如背景视频、颜色突变、光照、纹理、相机扰动等)与不同难度。下表汇总了各代表性基线在全扰动场景下的平均得分(Mean)与四分位均值(IQM):

测试基准 评估指标 DrQ-v2 (MFRL) SVEA (MFRL) SRM (MFRL) TD-MPC2 (MBRL) DreamerV3 (MBRL) Dr. G (MBRL) VIGOR (本文)
DMC Mean 626.1 726.2 752.0 599.1 635.3 623.7 777.6
DMC IQM 739.5 815.5 834.9 672.2 738.3 704.8 856.9
Robosuite Mean 41.0 82.8 91.8 40.2 64.3 54.0 131.8
Robosuite IQM 15.7 38.7 50.5 1.7 30.1 3.6 58.7

消融实验

消融实验深入验证了各个关键模块在零样本泛化中的不可替代性:

配置 DMC Mean (Easy / Hard / Avg) Robosuite Mean (Easy / Hard / Extreme / Avg) 相对性能 (DMC / Robo) 说明
VIGOR (完整模型) 863.9 / 781.1 / 822.5 332.0 / 236.7 / 146.0 / 238.3 100% / 100% 默认 AW (ζ=1) + DC + ES
w/o DC & ES 730.7 / 522.3 / 626.5 74.2 / 83.2 / 80.2 / 79.2 76% / 33% 移除动力学与编码器双级一致性
w/o DC 816.4 / 652.1 / 734.3 128.9 / 108.5 / 95.9 / 111.1 89% / 46% 仅保留编码器对齐,缺乏动力学递推约束
w/o ES 826.7 / 713.8 / 770.2 311.8 / 193.5 / 127.6 / 211.0 93% / 88% 缺乏编码器稳定,强干扰下出现表征漂移
weak-only (全弱增强) 773.9 / 606.6 / 690.2 134.2 / 112.6 / 93.4 / 113.4 83% / 47% 扰动多样性不足,难以应对极端域偏移
strong-only (全强增强) 423.8 / 374.9 / 399.4 187.1 / 137.3 / 121.3 / 148.6 48% / 62% 缺乏低噪声锚点,破坏 MBRL 动力学学习
weak-to-strong-only 550.6 / 481.3 / 516.0 297.6 / 138.1 / 130.1 / 188.6 62% / 79% 不做批次划分,持续强扰动损害表征

关键发现

  • 动力学级约束在严苛视觉偏移下占主导地位:在扰动更剧烈、要求精确接触的 Robosuite 操作任务上,去掉 DC 使得性能暴跌至完整模型的 46%,而去掉 ES 仅下降至 88%。这证明编码器表征即便相似,若缺乏多步动力学一致性,残余误差在规划展开中仍会被级联放大。
  • 全强增强在 MBRL 中产生灾难性反噬:在 MFRL 中行之有效的全量强数据增强在 DMC 上仅取得完整 VIGOR 48% 的性能,证实了无差别破坏批次几何会彻底扰乱世界模型的递推建模能力;AW 的非对称批次划分是稳定世界模型的必要条件。
  • 扰动不敏感性(Augmentation-agnostic):当将默认的 random-overlay 替换为机制截然不同的 random-conv(卷积纹理干扰)或 SRM(频域谱掩码)时,VIGOR 依然维持了卓越的零样本泛化表现,证明性能增益来自于隐空间一致性约束的归纳偏置,而非对某种特定数据增强手段的拟合。

亮点与洞察

  • 切中 MBRL 视觉泛化的核心机制痛点:精准定位并形式化了基于模型 RL 特有的「编码器初始 OOD 扰动 + 动力学递归时域展开累积」的两级失效模式,给出了针对性的系统解决方案。
  • 优雅的非对称批次与跨步递归对齐设计:巧妙设计 AW 划分,利用单次前向混合编码与以弱增强为 Target 的直接隐变量回归,既免去了高开销的对比学习负样本维护,又实现了长时域推演轨迹与真实流形的牢固针定。
  • 兼顾极高采样效率与零样本鲁棒性:彻底打破了「MFRL 泛化强但低效、MBRL 高效但脆弱」的权衡困境,在保留 TD-MPC2 原生快速收敛速度的前提下,在复杂操作扰动基准上超越既有 SOTA 达 43.6%。

局限与展望

  • 局限于外观表象扰动,未解决几何大形变:VIGOR 目前主要针对背景、色彩、纹理和光照等外观级视觉偏移;在附录中评估未见相机大角度旋转时性能依然显著下降,未来需要引入具有几何等变性或 3D 感知的世界模型表征。
  • 尚未拓展至物理动力学与任务级泛化:当前假设系统底层物理动力学(质量、摩擦力、重力)不变,仅抵御观测端扰动;未来结合上下文自适应机制应对环境物理参数跳变是重要演进方向。
  • 缺乏实体机械臂真机闭环验证:目前所有评测均在 MuJoCo 仿真环境下完成,真机部署中的光影非线性变化与真实传感器噪声尚待进一步实验检验。

相关工作与启发

  • vs TD-MPC2: 原生 TD-MPC2 缺乏任何视觉一致性机制,在未见视觉扰动下得分极低(DMC 599.1,Robosuite 40.2);VIGOR 在完全继承其高效无重构 MPC 架构的同时,赋予其顶尖的抗视觉扰动能力。
  • vs Dr. G (Dreamer 扩展): Dr. G 基于重构型世界模型 Dreamer 引入双重对比学习和逆动力学辅助损失;VIGOR 基于无重构 TD-MPC2,无需构建负样本对或重建图像,通过 AW+DC+ES 的直接潜变量回归机制在泛化与效率上均大幅优于 Dr. G。
  • vs SVEA / SRM (MFRL 强基线): SVEA 和 SRM 依赖单步 Q 函数正则或频域增强,仅能防护单步决策漂移;在面对多步时空展开的控制任务时,VIGOR 展现出显著更强的规划抗噪能力。

评分

  • 新颖性: ⭐⭐⭐⭐ [针对 MBRL 两级误差传播机制量身定制了非对称增强与递归动力学一致性架构]
  • 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 DMC 与 Robosuite 10 项任务、7 类视觉扰动、细致消融与 UMAP 隐空间流形可视化]
  • 写作质量: ⭐⭐⭐⭐⭐ [逻辑严密,两级脆弱性动机鲜明,方法与实验论证高度契合]
  • 价值: ⭐⭐⭐⭐⭐ [极大推进了视觉 MBRL 向现实世界复杂未见视觉场景落地的可靠性]