解耦视觉表征与动态学习，构建精准视觉世界模型Decoupling visual representation learning from dynamics learning to build an accurate visual world model
Younggyo Seo · Danijar Hafner · Hao Liu · Fangchen Liu · Stephen James · Kimin Lee · Pieter Abbeel | KAIST · UC Berkeley · Google Research · University of Toronto
MWM（Masked World Models）是一种基于模型的视觉强化学习框架，核心思想是将视觉表征学习与动态学习解耦：先用带卷积特征掩码的自编码器（ViT + 辅助奖励预测）学习高质量表征，再在该表征空间上训练隐空间动态模型。这一设计使世界模型能够捕捉小物体等精细视觉细节，在 Meta-world 50 个任务上将成功率从基线的 67.9% 提升至 81.7%。MWM (Masked World Models) is a model-based visual reinforcement learning framework whose core idea is to decouple visual representation learning from dynamics learning: an autoencoder with convolutional feature masking (ViT + auxiliary reward prediction) first learns high-quality representations, and a latent dynamics model is then trained on top of that representation space. This design lets the world model capture fine visual details such as small objects, raising the success rate on the 50 Meta-world tasks from the baseline 67.9% to 81.7%.
视觉模型强化学习（visual model-based RL）有潜力从视觉观测中实现高样本效率的机器人学习，但现有方法通常将视觉表征学习与动态建模端到端联合优化，难以准确建模机器人与小物体之间的交互。Visual model-based RL has the potential to deliver sample-efficient robot learning from visual observations, but existing methods usually optimize visual representation learning and dynamics modeling jointly end-to-end, which makes it hard to accurately model the interaction between robots and small objects.
"the current approaches typically train a single model end-to-end for learning both visual representations and dynamics, making it difficult to accurately model the interaction between robots and small objects."
图 1（论文 Figure 1）：MWM 框架总览。左侧：反复更新带卷积特征掩码与奖励预测的自编码器；右侧：在自编码器隐空间中学习隐动态模型（RSSM 变体）。动态学习期间自编码器参数冻结，实现真正的解耦更新。Figure 1 (paper Figure 1):Overview of the MWM framework. Left: the autoencoder with convolutional feature masking and reward prediction is updated repeatedly; right: a latent dynamics model (an RSSM variant) is learned in the latent space of the autoencoder. The autoencoder parameters stay frozen during dynamics learning, giving a genuinely decoupled update.
81.7%Meta-world 50 任务成功率（MWM）Success rate on the 50 Meta-world tasks (MWM)
端到端优化在表征质量与动态准确性之间存在内在权衡：世界模型需同时负责"看清画面"和"预测未来"，两者目标相互干扰。另一方面，类 MAE（masked autoencoder）的像素块掩码方式虽计算高效，但难以捕捉块内精细细节（如小目标位置），限制了其在视觉控制中的应用。MWM 正是为此而生。End-to-end optimization carries an inherent trade-off between representation quality and dynamics accuracy: the world model has to see the image clearly and predict the future at the same time, and the two objectives interfere with each other. On the other hand, MAE-style (masked autoencoder) pixel-patch masking is computationally efficient but struggles to capture fine details inside a patch, such as the location of a small object, which limits its use in visual control. MWM was created exactly for this.
02方法 MethodMethod
MWM 将整个学习循环分为三个交替步骤：(i) 用卷积特征掩码 + 辅助奖励预测训练自编码器；(ii) 在自编码器冻结表征空间上训练隐动态模型；(iii) 与环境交互收集新样本。MWM splits the whole learning loop into three alternating steps: (i) train the autoencoder with convolutional feature masking + auxiliary reward prediction; (ii) train the latent dynamics model on the frozen representation space of the autoencoder; (iii) interact with the environment to collect new samples.
与 MAE 的像素块掩码不同，MWM 先通过 convolution stem（3 层卷积 + 线性投影）将图像 ot 转换为卷积特征序列 hct，再以掩码比例 m = 75% 随机遮掩这些特征向量，送入 ViT encoder（4 层）和 ViT decoder（3 层）重建原始像素。
由于早期卷积层会在空间上混合低级细节，模型可以从相邻非遮掩特征中恢复块内所有细节——既保留了 MAE 的稳定性与计算效率，又弥补了其对精细细节捕捉不足的缺陷。
Unlike the pixel-patch masking of MAE, MWM first turns the image ot into a convolutional feature sequence hct through a convolution stem (3 convolutional layers + a linear projection), then randomly masks these feature vectors with a masking ratio of m = 75% and feeds them to a ViT encoder (4 layers) and a ViT decoder (3 layers) that reconstruct the raw pixels. Because the early convolutional layers mix low-level details spatially, the model can recover every detail inside a patch from neighboring unmasked features, which keeps the stability and computational efficiency of MAE while remedying its weakness at capturing fine details.
辅助奖励预测（Reward Prediction）：为将任务相关信息编码进表征，自编码器同时预测奖励 r̂t。具体做法是在 ViT decoder 输入中追加一个可学习 mask token，对应输出接线性头预测奖励。实验证明这一设计对最终性能至关重要（去掉后性能显著下降）。
Auxiliary reward prediction (Reward Prediction):To encode task-relevant information into the representation, the autoencoder also predicts the reward r̂t. Concretely, a learnable mask token is appended to the ViT decoder input and its output is fed to a linear head that predicts the reward. Experiments show that this design is essential to final performance, which drops markedly once it is removed.
隐空间动态模型（Latent Dynamics Learning）Latent Dynamics Model (Latent Dynamics Learning)
视觉表征学习完成后，利用冻结的自编码器表征 zc,0t（无掩码编码结果）训练 RSSM 变体。与原始 DreamerV2 的区别在于：Representation Model 以 zc,0t 而非原始像素为输入，图像解码器替换为"视觉表征解码器"（重建表征而非像素）。这使动态模型可以专注于建模状态转移，而不必同时学习高维像素重建。
Once visual representation learning is finished, an RSSM variant is trained on the frozen autoencoder representation zc,0t (the mask-free encoding). The difference from the original DreamerV2 is that the Representation Model takes zc,0t rather than raw pixels as input, and the image decoder is replaced by a visual-representation decoder that reconstructs representations instead of pixels. This lets the dynamics model concentrate on modeling state transitions without also having to learn high-dimensional pixel reconstruction.
整体优化目标 Lmwm(φ, θ) 包含两项：
• 视觉表征学习：重建损失（MSE）+ 奖励预测损失
• 动态学习：表征重建损失 + 奖励损失 + KL 散度（RSSM）
The overall objective Lmwm(φ, θ) has two terms: • Visual representation learning: reconstruction loss (MSE) + reward prediction loss • Dynamics learning: representation reconstruction loss + reward loss + KL divergence (RSSM)
图 3（论文 Figure 3）：Meta-world 六个难度较高的视觉操纵任务（需要机器人臂与小物体交互）上的学习曲线。MWM（蓝）在样本效率和最终性能上均显著优于 DreamerV2（黄）。特别地，Pick Place 任务中 MWM 达到 >80% 成功率，而 DreamerV2 几乎无法解决该任务。Figure 3 (paper Figure 3):Learning curves on six harder Meta-world visual manipulation tasks that require the robot arm to interact with small objects. MWM (blue) is clearly better than DreamerV2 (yellow) in both sample efficiency and final performance. In particular, on Pick Place MWM reaches a >80% success rate, while DreamerV2 can hardly solve the task.
03实验 ExperimentsExperiments
在三个基准上评估 MWM：Meta-world（50 任务）、RLBench（2 任务）、DeepMind Control Suite（视觉运动控制）。视觉输入统一为 64×64×3 RGB 图像，基线为 DreamerV2。MWM is evaluated on three benchmarks: Meta-world (50 tasks), RLBench (2 tasks) and the DeepMind Control Suite (visual locomotion control). Visual inputs are uniformly 64×64×3 RGB images and the baseline is DreamerV2.
图 4 & 5（论文 Figure 4, 5）：上：Meta-world 50 任务汇总成功率（步数归一化）；RLBench Reach Target（中）和 Push Button（右）学习曲线。下：DeepMind Control Suite 三个视觉任务的 episode return。实线为均值，阴影为 bootstrap 置信区间。Figures 4 & 5 (paper Figure 4, 5):Top: aggregate success rate over the 50 Meta-world tasks (normalized by steps); learning curves for RLBench Reach Target (middle) and Push Button (right). Bottom: episode return on three visual tasks of the DeepMind Control Suite. Solid lines are means, shaded areas are bootstrap confidence intervals.
消融研究 AblationsAblations
图 6（下图）系统分析了三个关键设计选择：Figure 6 (below) systematically analyzes three key design choices:
卷积特征掩码 vs. 像素块掩码（MAE）：卷积特征掩码显著优于像素块掩码，验证了捕捉块内精细细节对视觉控制的重要性。Convolutional feature masking vs. pixel-patch masking (MAE):Convolutional feature masking is clearly better than pixel-patch masking, confirming how important capturing fine intra-patch details is for visual control.
掩码比例 m：m = 75% 优于 m ∈ {0.25, 0.5}，说明强正则化可防止模型从相邻卷积特征中"走捷径"。但 m = 0.9（过强）会损害性能。Masking ratio m:m = 75% is better than m ∈ {0.25, 0.5}, showing that strong regularization keeps the model from taking shortcuts through neighboring convolutional features. But m = 0.9, which is too strong, hurts performance.
奖励预测：去掉辅助奖励预测目标后性能大幅下降，说明重建目标单独不足以学习任务相关信息。Reward prediction:Removing the auxiliary reward prediction objective causes a large performance drop, showing that the reconstruction objective alone is not enough to learn task-relevant information.
图 6 & 7（论文 Figure 6, 7）：上：三组消融曲线（卷积特征掩码、掩码比例、奖励预测）在 Meta-world 三个任务上的表现。下：RLBench Reach Target 任务中，真实帧、自编码器重建（Recon）与 MWM/DreamerV2 隐动态预测帧的对比。MWM 能准确预测红色目标块位置，DreamerV2 则失败。Figures 6 & 7 (paper Figure 6, 7):Top: three groups of ablation curves (convolutional feature masking, masking ratio, reward prediction) on three Meta-world tasks. Bottom: on RLBench Reach Target, a comparison of ground-truth frames, autoencoder reconstructions (Recon) and the latent-dynamics predictions of MWM / DreamerV2. MWM predicts the position of the red target block accurately, whereas DreamerV2 fails.
定性分析（Figure 7）揭示了解耦设计的工作机制：自编码器忠实重建所有细节（包括蓝、橙色干扰块），而动态模型的预测则选择性地只追踪任务相关组件（红色目标块），忽略无关细节。The qualitative analysis (Figure 7) reveals how the decoupled design works: the autoencoder faithfully reconstructs every detail, including the blue and orange distractor blocks, while the predictions of the dynamics model selectively track only the task-relevant component, the red target block, and ignore irrelevant details.
04局限性 LimitationsLimitations
Note:以下局限性均由作者在论文 Discussion 节明确陈述（stated by the authors）。All limitations below are explicitly stated by the authors in the Discussion section of the paper.
对辅助奖励预测的依赖Dependence on auxiliary reward prediction
"the performance of our approach heavily depends on the auxiliary reward prediction task. This might be because our autoencoder is not learning temporal information, which is crucial for learning task-relevant information." 如果奖励信号不可用（如无监督探索场景），当前框架将难以编码任务相关信息。"the performance of our approach heavily depends on the auxiliary reward prediction task. This might be because our autoencoder is not learning temporal information, which is crucial for learning task-relevant information." If the reward signal is unavailable, as in unsupervised exploration settings, the current framework can hardly encode task-relevant information.
仅使用单视角 RGB 输入Only single-viewpoint RGB input
"our model operates only on RGB pixels from a single camera viewpoint"，尚未整合本体感受状态（proprioceptive states）、点云等其他模态输入，限制了在更复杂真实场景中的应用。"our model operates only on RGB pixels from a single camera viewpoint"; it does not yet integrate other input modalities such as proprioceptive states or point clouds, which limits its application in more complex real-world scenarios.
样本效率仍不足以部署于真实世界Sample efficiency still too low for real-world deployment
RLBench 实验结果"are preliminary because they are still too sample-inefficient to be used in real-world scenarios."未来方向包括：利用少量演示数据、引入带路径规划的动作模式、在视频数据集上预训练世界模型等。The RLBench results "are preliminary because they are still too sample-inefficient to be used in real-world scenarios." Future directions include exploiting a few demonstrations, introducing action modes with path planning, and pre-training the world model on video datasets.
未利用时序视频表征学习Temporal video representation learning not exploited
自编码器在单帧图像上训练，不包含时序信息。作者指出探索视频 ViT 表征学习是值得研究的未来方向（inferred from design）。The autoencoder is trained on single-frame images and carries no temporal information. The authors point out that exploring video ViT representation learning is a future direction worth studying (inferred from design).
