モデルベースRLに潜在空間一貫性を導入したVIGOR、未見の視覚変化へゼロショットで汎化

Abstract

Model-based reinforcement learning (MBRL) achieves strong sample efficiency by planning within learned latent dynamics, yet its performance degrades substantially under unseen visual distractions such as background variations, lighting changes, or camera shifts. Unlike model-free RL, where encoder perturbations affect only single-step predictions, MBRL suffers from a two-level vulnerability: visual distractions first push encoder outputs out of distribution, and these errors then compound through recursive latent rollouts over the planning horizon. We propose VIsual Generalization via latent-space cOnsistency in model-based RL (VIGOR), a framework that enables zero-shot generalization to unseen visual distractions while retaining the sample efficiency of its MBRL backbone. VIGOR integrates three interdependent components: (i) asymmetric weak-to-strong augmentation, which pairs weak-only and weak-to-strong latent views within a single batch; (ii) dynamics-level consistency, which enforces augmentation-invariant transition predictions through direct latent regression; and (iii) encoder-level stabilization, which prevents encoder drift under the cross-augmentation supervision imposed by dynamics-level consistency. Evaluations on the DeepMind Control Suite (DMC) and Robosuite show that VIGOR outperforms state-of-the-art model-free and model-based baselines, surpassing the second-best baseline by 3.4\% on DMC and 43.6\% on Robosuite. Ablations further show that VIGOR’s robustness is augmentation-agnostic: replacing the default augmentation with alternatives from distinct perturbation families preserves strong generalization, confirming that latent-space consistency, not the augmentation choice, drives robustness.

1 Introduction

Visual reinforcement learning (visual RL) trains control policies directly from pixels, eliminating task-specific state estimation [26, 25, 45, 44]. Despite success in locomotion and manipulation [8, 23, 17], visual RL remains sensitive to unseen visual distractions: even minor shifts in backgrounds, lighting, or camera position can cause substantial policy degradation.

In model-free RL (MFRL), visual generalization has been addressed through data augmentation [16, 20], robust representation learning [18, 4, 48], and stabilized value-function training [29, 19]. Model-based RL (MBRL) offers superior sample efficiency by planning in learned latent dynamics [17, 12], but achieving visual generalization in MBRL is structurally more challenging. In MFRL, encoder perturbations affect only a single-step prediction. In MBRL, the problem is two-level: visual distractions first push encoder outputs out of the training distribution, and these out-of-distribution (OOD) latents are then recursively propagated by the dynamics model over a planning horizon H, causing errors to compound across steps (Figure 1). This two-level vulnerability, which does not arise in MFRL, means that techniques effective for single-step robustness are insufficient when the encoder simultaneously serves a recursive dynamics model, Q-value networks, and a planner.

We propose VIsual Generalization via latent-space cOnsistency in model-based RL (VIGOR), a framework that enables zero-shot generalization while retaining the sample efficiency of its TD-MPC2 backbone. VIGOR addresses the two-level vulnerability through three interdependent components: (i) asymmetric weak-to-strong augmentation (AW), an asymmetric batch-splitting strategy that produces paired weak-only and weak-to-strong latents; (ii) dynamics-level consistency (DC), which enforces augmentation-invariant transition predictions by anchoring dynamics outputs to stable targets; and (iii) encoder-level stabilization (ES), which prevents the shared encoder from drifting under the cross-augmentation pressure imposed by DC.

We evaluate VIGOR on the DeepMind Control Suite (DMC) [40] and Robosuite [54] under the RL-ViGen zero-shot protocol [49], covering seven visual distraction types—a substantially broader evaluation than the two types (background-color, background-video) typically considered in prior work [18, 20, 4, 7] (Appendix D.3). VIGOR outperforms state-of-the-art MFRL and MBRL baselines, improving zero-shot generalization over the second-best baseline by 3.4\% on DMC and 43.6\% on Robosuite while retaining the sample efficiency of its TD-MPC2 backbone. Ablation studies further reveal that naive transfer of MFRL strategies fails in MBRL: applying strong augmentation uniformly across the batch, without the asymmetric splitting of AW, achieves only 48\% of VIGOR’s performance on DMC (Section 4.3). Furthermore, VIGOR’s robustness is augmentation-agnostic: latent-space consistency, rather than any specific augmentation choice, drives generalization (Section 4.4; Appendix E.5).

Our contributions are:

• Problem identification and analysis. We identify a two-level failure mode in MBRL: visual distractions first produce OOD encoder outputs, which then compound through recursive rollouts—a failure mode unique to recursive latent dynamics (Figure 1; Section 4.5).
• Integrated framework for MBRL robustness. We propose VIGOR, which addresses this failure mode through three interdependent components (AW, DC, and ES) designed for the recursive world-model setting. We show that naive transfer of MFRL strategies fails in MBRL, that VIGOR’s integration is necessary, and that the resulting robustness is augmentation-agnostic (Sections 4.3–4.5; Appendix E.5).
• State-of-the-art zero-shot generalization. VIGOR outperforms MFRL and MBRL baselines, surpassing the second-best baseline by 3.4\% on DMC and 43.6\% on Robosuite while preserving the strong sample efficiency of its TD-MPC2 backbone.

2 Related Work

2.1 Visual Model-Based Reinforcement Learning

Latent-space MBRL learns compact world models for planning or imagination in latent space [6, 10, 12, 14, 17]. The Dreamer series adopts a decoder-based architecture that reconstructs pixel observations [9, 11, 12], whereas TD-MPC uses a decoder-free architecture that plans without reconstruction [14, 17]. Both families assume i.i.d. training and test distributions; in practice, unseen visual distractions push encoder outputs out of distribution, causing compounding rollout errors (Figure 1). VIGOR equips latent-space MBRL with three interdependent components (AW, DC, and ES) designed to mitigate this failure mode. An extended discussion of MBRL architectures and our backbone selection rationale is provided in Appendix F.

2.2 Visual Generalization in Model-Free Reinforcement Learning

MFRL approaches address visual robustness through data augmentation [16, 20, 37], self-supervised representation learning [18, 4], and pretrained visual encoders [48, 32]. These methods operate at the single-step level: the encoder maps observations directly to Q-values or policy outputs without recursive rollouts. In MBRL, by contrast, encoder perturbations propagate through multi-step latent rollouts and compound across the planning horizon, motivating VIGOR’s two-level design. A detailed survey is provided in Appendix G.1.

2.3 Visual Generalization in Model-Based Reinforcement Learning

Invariant MBRL methods [50, 41, 53, 39, 52] learn representations that discard task-irrelevant visual factors but typically require test-time adaptation. VIGOR targets the stricter zero-shot setting where no adaptation data are available. Dr. G [7] is the most directly comparable prior work, as it explicitly targets zero-shot generalization in latent-space MBRL. It extends Dreamer [9] with dual contrastive objectives and an inverse-dynamics auxiliary loss. VIGOR differs in mechanism: rather than contrastive negatives and auxiliary supervision, it enforces dynamics-level consistency (DC) through direct latent regression, supported by asymmetric weak-to-strong augmentation (AW) and encoder-level stabilization (ES). Built on TD-MPC2 [17], VIGOR targets the same zero-shot setting through a fundamentally different inductive bias. A detailed survey is provided in Appendix G.2.

3 VIGOR

VIGOR addresses the two-level vulnerability identified in Section 1: visual distractions first perturb encoder outputs, and these perturbations then compound through recursive latent rollouts. The framework comprises three interdependent components, each targeting a different level of the world-model pipeline: input augmentation, latent dynamics, and encoder representation. Figure 2 provides an overview. The three components depend on each other: AW produces the paired latents on which DC operates; DC stabilizes dynamics but can cause encoder drift; ES counteracts this drift. Because AW modifies only the input pipeline while DC and ES act through auxiliary losses, VIGOR can be integrated into latent-space MBRL frameworks without modifying the underlying planner or optimization procedure.

3.1 Asymmetric Weak-to-Strong Augmentation (AW)

We distinguish between weak augmentations (minor perturbations that preserve most visual information) and strong augmentations that impose substantial visual changes. We refer to the sequential composition \tau^{s}\circ\tau^{w} as weak-to-strong augmentation, emphasizing application order rather than a progressive curriculum. We adopt random-shift [44] as \tau^{w} and random-overlay [18] as \tau^{s}.

AW produces paired weak-only and weak-to-strong samples within a single mini-batch, providing the structural foundation for DC while preserving the original batch size. As illustrated in Figure 2, observations are processed through two augmentation branches: the weak-only branch applies \tau^{w} alone, while the weak-to-strong branch applies \tau^{s}\circ\tau^{w}. The two branches are then concatenated along the batch dimension and jointly encoded by the shared encoder h_{\theta} into the mixed representation z_{t}^{m}, which is rolled forward by the dynamics model.

Batch split. Given a mini-batch \mathcal{B} with index set \mathcal{I}=\{1,\dots,|\mathcal{B}|\}, we partition \mathcal{I} uniformly at random into two disjoint subsets:

\displaystyle\mathcal{I}=\mathcal{I}^{w}\cup\mathcal{I}^{ws},\quad\mathcal{I}^{w}\cap\mathcal{I}^{ws}=\emptyset,\quad\zeta=|\mathcal{I}^{w}|/|\mathcal{I}^{ws}|.

Samples in \mathcal{I}^{w} are routed through the weak-only branch, while samples in \mathcal{I}^{ws} are routed through the weak-to-strong branch. We set \zeta=1 throughout our experiments, splitting the batch evenly.

Augmented representations. At the first time step of each trajectory segment of horizon H, we construct weak-only and weak-to-strong observation stacks for the two subsets:

\displaystyle s_{t}^{w}=\tau^{w}\bigl(\{s_{t,i}\}_{i\in\mathcal{I}^{w}},\upsilon^{w}\bigr),\quad s_{t}^{ws}=\tau^{s}\bigl(\tau^{w}\bigl(\{s_{t,j}\}_{j\in\mathcal{I}^{ws}},\upsilon^{w}\bigr),\upsilon^{s}\bigr),

where \upsilon^{w}\sim\Upsilon^{w} and \upsilon^{s}\sim\Upsilon^{s} parameterize the respective augmentations. The mixed observation s_{t}^{m}=s_{t}^{w}\oplus s_{t}^{ws}, where \oplus denotes batch-dimension concatenation, is encoded into z_{t}^{m}=h_{\theta}(s_{t}^{m}) in a single forward pass. Further details on AW are provided in Appendix E.5.

3.2 Dynamics-Level Consistency (DC)

The key insight is that task-relevant dynamics should remain invariant to visual distractions: the world model should produce consistent next-state predictions regardless of the visual perturbation applied to its input. The DC loss enforces this invariance at the dynamics level, a constraint specific to latent-space MBRL, where the encoder output feeds a recursive transition model.

For reference, the standard dynamics loss in latent-space MBRL [17] is:

\displaystyle\mathcal{L}_{\text{dyn}}(\theta)=\bigl\lVert d_{\theta}(z_{t},a_{t})-\operatorname{sg}\bigl(h_{\theta}(s_{t+1})\bigr)\bigr\rVert_{2}^{2},

where z_{t}=h_{\theta}(s_{t}) is the encoded observation, d_{\theta} is the dynamics model, and \operatorname{sg}(\cdot) denotes the stop-gradient operator.

DC modifies this objective in two ways: at the first time step of each trajectory segment, the dynamics model receives the mixed representation z_{t}^{m} from AW; the target is the encoding of a weakly augmented successor. The single-step DC loss is:

\displaystyle\mathcal{L}_{\text{DC}}(\theta)=\bigl\lVert d_{\theta}(\hat{z}_{t}^{m},a_{t})-\operatorname{sg}(z_{t+1}^{w})\bigr\rVert_{2}^{2},\quad\text{where }z_{t+1}^{w}=h_{\theta}\bigl(\tau^{w}(s_{t+1},\upsilon^{w})\bigr).

The mixed representation z_{t}^{m} is injected only at the initial step of each trajectory segment of horizon H (\hat{z}_{t}^{m}=z_{t}^{m}), after which the dynamics model rolls forward recursively: \hat{z}_{t+h}^{m}=d_{\theta}(\hat{z}_{t+h-1}^{m},a_{t+h-1}) for h=1,\dots,H. At every step, the rollout prediction \hat{z}_{t+h}^{m} is aligned with the weakly augmented target z_{t+h}^{w}, propagating the initial perturbation through the entire horizon while requiring augmentation-invariant transitions.

The weak target z_{t+h}^{w} provides a stable supervisory signal. Applying \tau^{w} universally, rather than maintaining a clean path, is a deliberate choice: consistent with prior work on visual RL [45, 44], weak augmentation improves learning and provides a stable yet slightly diversified anchor for the subsequent objectives. The model must produce consistent predictions from the perturbed mixed input, thereby enforcing augmentation-invariant latent dynamics.

3.3 Encoder-Level Stabilization (ES)

While DC enforces invariance at the dynamics level, the encoder may still produce inconsistent representations across augmentations. In VIGOR, a single shared encoder feeds the dynamics model, Q-value networks, and model predictive control (MPC) planner. The cross-augmentation supervision from DC can distort this shared representation, causing encoder drift that propagates through multi-step rollouts. ES adapts the alignment objective of SODA [18], originally proposed for MFRL, to a different role in VIGOR: counteracting the encoder drift induced by DC’s cross-augmentation supervision, a failure mode specific to MBRL, where the encoder feeds recursive latent rollouts.

At each time step, the ES loss minimizes the \ell_{2} distance between \ell_{2}-normalized weak-only and weak-to-strong encodings of the same frame:

\displaystyle\mathcal{L}_{\text{ES}}(\theta)=\Bigl\|\frac{z_{t}^{ws}}{\lVert z_{t}^{ws}\rVert_{2}}-\frac{\operatorname{sg}(z_{t}^{w})}{\lVert\operatorname{sg}(z_{t}^{w})\rVert_{2}}\Bigr\|_{2}^{2}.

ES anchors weak-to-strong encodings to their weak-only counterparts, counteracting the encoder drift induced by DC’s cross-augmentation supervision. A detailed discussion of ES’s structural role is provided in Appendix H.

3.4 Overall Training Objective

VIGOR augments the TD-MPC2 [17] objective with \mathcal{L}_{\text{DC}} and \mathcal{L}_{\text{ES}}. For a trajectory segment \Gamma=(s_{i},a_{i},r_{i},s_{i+1})_{i=t}^{t+H} sampled from replay buffer \mathcal{B}, the overall objective is:

\mathcal{L}_{\text{total}}(\theta)=\mathbb{E}_{\Gamma\sim\mathcal{B}}\Biggl[{\color[rgb]{0.9609,0,0}\mathbb{E}_{\upsilon^{w}\sim\Upsilon^{w},\upsilon^{s}\sim\Upsilon^{s}}}\Bigl[\sum_{i=t}^{t+H}\!\lambda^{\,i-t}\underbrace{\underbrace{\mathcal{L}_{\text{TD-MPC2}}(\theta;\mathcal{L}_{\text{rew}},\mathcal{L}_{Q},{\color[rgb]{0.9609,0,0}\mathcal{L}_{\text{DC}}},\Gamma_{i})}_{\text{world-model losses}}\;+\;{\color[rgb]{0.9609,0,0}\mathcal{L}_{\text{ES}}(\theta;\upsilon^{w},\upsilon^{s})}}_{{\color[rgb]{0.9609,0,0}\mathcal{L}_{\text{VIGOR}}(\theta;\upsilon^{w},\upsilon^{s})}}\Bigr]\Biggr],

where \Gamma_{i}=(s_{i},a_{i},r_{i},s_{i+1}) denotes the per-step transition tuple within the segment, and \mathcal{L}_{\text{TD-MPC2}} denotes the base world-model objective comprising the reward prediction loss (\mathcal{L}_{\text{rew}}) and the Q-value loss (\mathcal{L}_{Q}), augmented with the DC loss (\mathcal{L}_{\text{DC}}). The discount factor \lambda\in(0,1] downweights losses at later horizon steps. Here, \mathcal{L}_{\text{VIGOR}}(\theta;\upsilon^{w},\upsilon^{s}) denotes the per-step VIGOR loss. The terms highlighted in red are the only additions to the original TD-MPC2 objective. The training procedure is summarized in Algorithm 1 (Appendix B); implementation details and hyperparameters are given in Appendix C; and a detailed discussion of each component appears in Appendix H.

4 Experiments

We evaluate VIGOR’s zero-shot generalization and sample efficiency against state-of-the-art MFRL and MBRL methods on six continuous control tasks from DMC [40] and four manipulation tasks from Robosuite [54]. All agents are trained in clean environments and evaluated zero-shot under unseen visual distractions (Figure 3; representative evaluation scenes are catalogued in Appendix D.5). Our experiments address the following questions:

• Does VIGOR achieve superior zero-shot generalization while preserving sample efficiency?
• How does AW contribute to the efficiency–robustness trade-off, and why is its batch-splitting design necessary for MBRL?
• How do DC and ES individually contribute, and how do they interact?
• How do specific design choices (augmentation scheduling, augmentation family, and auxiliary objective) affect generalization?
• Does VIGOR preserve latent-space consistency at both the encoder and dynamics levels under unseen visual shifts?

4.1 Experimental Setup

Environments. We benchmark VIGOR on six DMC tasks (walker-walk, finger-spin, cheetah-run, walker-stand, reacher-easy, and cartpole-swingup) and four Robosuite manipulation tasks (door-opening, nut-assembly, peg-in-hole, and lifting). Details are provided in Appendix D.

Baselines. We compare VIGOR against eight baselines spanning both paradigms: four MFRL methods (DrQ-v2 [44], SVEA [16], SRM [20], SimGRL [37]) and four MBRL methods (TD-MPC2 [17], DreamerV3 [12], HRSSM [39], Dr. G [7]). Baseline descriptions are provided in Appendix D.2.

Distractions. We follow the RL-ViGen zero-shot protocol [49] and evaluate all methods under seven visual distraction types, substantially broader than the two types (background-color, background-video) used in most prior work [18, 20, 4, 7]. Details are provided in Appendix D.3.

4.2 Experimental Results

To address (Q1), we evaluate VIGOR along two axes: zero-shot generalization and sample efficiency. All results are averaged over ten random seeds.

Zero-shot generalization. All agents are trained in a clean setting and evaluated zero-shot under unseen visual distractions across multiple difficulty levels (easy and hard in DMC; easy, hard, and extreme in Robosuite). Table 1 reports mean and interquartile mean (IQM) [1] aggregated across all tasks, distraction types, and difficulty levels; detailed breakdowns are provided in Appendix E.1.

VIGOR achieves the strongest zero-shot generalization on both benchmarks. On DMC, it leads on both metrics (mean: 777.6; IQM: 856.9), surpassing the second-best baseline (SRM) by 3.4\% in mean and 2.6\% in IQM. The gap widens substantially on Robosuite, where most baselines degrade under severe distribution shift: VIGOR reaches mean 131.8 and IQM 58.7, exceeding the second-best baseline (SRM) by 43.6\% in mean and 16.2\% in IQM. Notably, our backbone TD-MPC2 ranks lowest among all baselines on both benchmarks and both metrics, indicating that latent-space MBRL without explicit consistency mechanisms is particularly fragile under visual shifts. VIGOR’s strong performance suggests that jointly enforcing encoder-level stabilization and dynamics-level consistency enables structured rollouts under unseen distractions.

Sample efficiency. Figure 4 reports aggregated IQM normalized scores; per-task curves and statistical analyses are provided in Appendix E.2. On DMC, VIGOR closely matches TD-MPC2’s rapid convergence while requiring fewer samples than all other baselines. On Robosuite, the efficiency gap widens further: TD-MPC2 and VIGOR converge at comparable rates, while other baselines fall behind in both speed and asymptotic returns. These results confirm that VIGOR preserves the strong sample efficiency of its TD-MPC2 backbone while substantially improving robustness—a combination that prior methods achieve only in part. We corroborate this under a tighter budget on DMC 100K (Appendix E.3) and report wall-clock cost in Appendix E.4.

4.3 Ablation Studies

We conduct ablation studies addressing (Q2) and (Q3), isolating the contribution of each component to VIGOR’s efficiency–robustness balance.

AW strategy (Q2). We compare AW against three variants: weak-only (WO), strong-only (SO), and weak-to-strong-only (WTSO), which applies both transforms to every sample without batch splitting. We also sweep the mixture ratio \zeta=|\mathcal{I}^{w}|/|\mathcal{I}^{ws}| to probe the sensitivity of the weak-only to weak-to-strong balance. Table 2 and Figure 5 report zero-shot generalization and sample efficiency, respectively; further analyses are provided in Appendix E.5.

The results reveal complementary failure modes that illuminate why AW’s specific design is necessary for MBRL. On DMC, SO degrades to 48\% of VIGOR, confirming that applying strong augmentations to an entire batch triggers the two-level vulnerability: perturbed encoder outputs compound through recursive rollouts, destabilizing world-model learning. WTSO partially recovers (62\%) but remains far below AW, indicating that persistent strong exposure without a stable weak-only anchor disrupts optimization. WO converges fast (Figure 5, left), but under-generalizes (83\%), as weak shifts alone are insufficient for harder distractions. On Robosuite, where distribution shifts are more severe, the ranking inverts: WO becomes the weakest (47\%), SO improves relatively (62\%), and WTSO reaches (79\%) yet still falls short of AW. AW is the only strategy that succeeds on both benchmarks, because its asymmetric splitting simultaneously provides optimization stability and perturbation diversity.

The \zeta-sweep reveals a complementary cross-benchmark pattern: on DMC, increasing the weak-only fraction improves performance (\zeta{=}4: 96\% > \zeta{=}2: 94\% > \zeta{=}0.5: 88\%), while on Robosuite the trend reverses (\zeta{=}0.5: 87\% > \zeta{=}2: 80\% > \zeta{=}4: 75\%). Only the balanced mixture (\zeta{=}1) achieves the best performance on both, confirming that this ratio balances stability and generalization coverage.

DC and ES contributions (Q3). Table 3 isolates the effect of each loss. Removing DC reduces generalization to 89\% of full VIGOR on DMC and to 46\% on Robosuite; removing ES reduces it to 93\% and 88\%, respectively. Removing both losses jointly reduces performance further (76\% on DMC; 33\% on Robosuite).

The asymmetry between DC and ES is more pronounced under severe distribution shifts. On DMC, both components contribute modestly and similarly (89\% vs. 93\%). On Robosuite, the gap widens sharply: DC alone (w/o ES) retains 88\%, whereas ES alone (w/o DC) retains only 46\%. DC constrains the recursive prediction pathway: without it, the dynamics model produces inconsistent rollouts under visual shifts. ES alone cannot compensate, because encoder-level alignment does not prevent the dynamics model from amplifying residual discrepancies across multi-step rollouts. Conversely, the 12\% gap between DC alone and full VIGOR on Robosuite shows that ES’s encoder-level alignment is also needed for DC’s dynamics-level constraint to function effectively. The two losses are thus complementary: DC anchors transition predictions, while ES anchors the representations on which DC operates. We further visualize these failure modes in Appendix E.6.

Summary. VIGOR’s gains stem from the interplay of all three components. AW provides the efficiency–robustness balance through asymmetric batch splitting; DC and ES are jointly necessary for reliable generalization, with DC playing the dominant role under severe distribution shifts by directly constraining the recursive dynamics pathway.

4.4 Ablation on Key Design Choices

We analyze (Q4) by ablating three design choices, each isolated through a single-variable variant: VIGOR_CONST_AUG (fixed augmentation parameters across the horizon), VIGOR_CONV (random-conv replacing random-overlay), and VIGOR_CONTRAST (contrastive ES objective replacing the alignment-based one). Figure 6 summarizes results on finger-spin, cheetah-run, and walker-walk. VIGOR outperforms all variants on both axes, indicating that (i) dynamic augmentation resampling is important for stable rollouts; (ii) random-overlay provides stronger robustness than convolution-based perturbations; and (iii) alignment-based ES complements AW and DC more effectively than contrastive ES. Formal definitions and detailed results are provided in Appendix E.7.

Augmentation-agnostic robustness. A central question is whether VIGOR’s robustness depends on the specific augmentation choice or on latent-space consistency. We replace random-overlay with random-conv (textural perturbations) and SRM (spectral masking), two families that produce qualitatively distinct perturbations (Appendix E.5, Figure 19). As shown in Figure 7, all three variants maintain strong zero-shot generalization, indicating that VIGOR’s robustness is driven by latent-space consistency rather than by the choice of augmentation.

4.5 Latent-Space Consistency Analysis

To address (Q5), we compare VIGOR against TD-MPC2 (its backbone) and Dr. G (the most directly comparable prior method) by rolling out each agent’s world model from identical initial inputs and projecting the predicted latent states into 2D via UMAP [31] (Figure 8).

VIGOR maintains a compact latent manifold across difficulty levels, consistently aligning perturbed trajectories with their clean counterparts. This consistency holds at both levels of the vulnerability identified in Figure 1: the initial encoder outputs under visual distractions remain close to their clean counterparts (encoder level), and this alignment is preserved throughout the rollout horizon (dynamics level). The resulting structural stability translates into stable returns. TD-MPC2 and Dr. G exhibit scattered or divergent rollouts, leading to degraded performance, despite Dr. G’s contrastive objectives being designed for generalization. This comparison suggests that direct dynamics-level regression (DC) based on asymmetric batch-splitting (AW), combined with encoder stabilization (ES), provides a more effective inductive bias for MBRL robustness than contrastive objectives with inverse-dynamics supervision, because it jointly constrains both levels of the prediction pathway that drives planning. Additional visualizations are provided in Appendix E.8; we further verify that this consistency holds under deeper planning horizons up to H{=}10 (Appendix E.9).

5 Conclusion

We introduced VIGOR, a framework that achieves robust zero-shot visual generalization in latent-space MBRL while preserving the strong sample efficiency of its TD-MPC2 backbone. VIGOR addresses the two-level vulnerability of latent-space MBRL: visual distractions perturb encoder outputs, and these perturbations compound through recursive rollouts. The framework integrates three interdependent components (AW, DC, and ES) to mitigate this failure mode. Across challenging OOD benchmarks, VIGOR surpasses the second-best baseline by 3.4\% on DMC and 43.6\% on Robosuite. Ablation studies confirm that all three components are jointly necessary and that VIGOR’s robustness is augmentation-agnostic: latent-space consistency, not any particular augmentation choice, is what drives generalization.

Limitations and future work. VIGOR targets appearance-based perturbations and does not address geometric distortions; evaluations under unseen camera rotations (Appendix E.10) reveal that such shifts remain challenging, motivating geometry-aware representations as a primary future direction. Beyond visual shifts, extending VIGOR to dynamics- and task-level shifts [36, 5] is a natural next step, and validating it in real-robot deployments would assess its practical applicability.

References

1. [1] R. Agarwal, M. Schwarzer, P. S. Castro, A. C. Courville, and M. Bellemare (2021) Deep reinforcement learning at the edge of the statistical precipice. Advances in Neural Information Processing Systems 34, pp. 29304–29320.
2. [2] E. Alonso, A. Jelley, V. Micheli, A. Kanervisto, A. Storkey, T. Pearce, and F. Fleuret (2024) Diffusion for world modeling: visual details matter in Atari. Advances in Neural Information Processing Systems 37, pp. 58757–58791.
3. [3] S. Batra and G. S. Sukhatme (2025) Zero shot generalization of vision-based rl without data augmentation. In International Conference on Machine Learning, pp. 3251–3273.
4. [4] D. Bertoin, A. Zouitine, M. Zouitine, and E. Rachelson (2022) Look where you look! saliency-guided q-networks for generalization in visual reinforcement learning. Advances in Neural Information Processing Systems 35, pp. 30693–30706.
5. [5] M. Beukman, D. Jarvis, R. Klein, S. James, and B. Rosman (2023) Dynamics generalisation in reinforcement learning via adaptive context-aware policies. Advances in Neural Information Processing Systems 36, pp. 40167–40203.
6. [6] D. Ha and J. Schmidhuber (2018) World models. Advances in Neural Information Processing Systems 31.
7. [7] J. Ha, K. Kim, and Y. Kim (2025) Dream to generalize: zero-shot model-based reinforcement learning for unseen visual distractions. arXiv preprint arXiv:2506.05419.
8. [8] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine (2018) Soft actor-critic: off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning, pp. 1861–1870.
9. [9] D. Hafner, T. Lillicrap, J. Ba, and M. Norouzi (2020) Dream to control: learning behaviors by latent imagination. In International Conference on Learning Representations,
10. [10] D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson (2019) Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2555–2565.
11. [11] D. Hafner, T. P. Lillicrap, M. Norouzi, and J. Ba (2021) Mastering atari with discrete world models. In International Conference on Learning Representations,
12. [12] D. Hafner, J. Pasukonis, J. Ba, and T. Lillicrap (2025) Mastering diverse domains through world models. Nature 640.
13. [13] D. Hafner, W. Yan, and T. Lillicrap (2025) Training agents inside of scalable world models. arXiv preprint arXiv:2509.24527.
14. [14] N. A. Hansen, H. Su, and X. Wang (2022) Temporal difference learning for model predictive control. In International Conference on Machine Learning, pp. 8387–8406.
15. [15] N. Hansen, R. Jangir, Y. Sun, G. Alenyà, P. Abbeel, A. A. Efros, L. Pinto, and X. Wang (2021) Self-supervised policy adaptation during deployment. In International Conference on Learning Representations,
16. [16] N. Hansen, H. Su, and X. Wang (2021) Stabilizing deep q-learning with convnets and vision transformers under data augmentation. Advances in Neural Information Processing Systems 34, pp. 3680–3693.
17. [17] N. Hansen, H. Su, and X. Wang (2024) TD-MPC2: scalable, robust world models for continuous control. In International Conference on Learning Representations,
18. [18] N. Hansen and X. Wang (2021) Generalization in reinforcement learning by soft data augmentation. In IEEE International Conference on Robotics and Automation, pp. 13611–13617.
19. [19] S. Huang, Y. Sun, J. Hu, S. Guo, H. Chen, Y. Chang, L. Sun, and B. Yang (2024) Learning generalizable agents via saliency-guided features decorrelation. Advances in Neural Information Processing Systems 36.
20. [20] Y. Huang, P. Peng, Y. Zhao, G. Chen, and Y. Tian (2022) Spectrum random masking for generalization in image-based reinforcement learning. Advances in Neural Information Processing Systems 35, pp. 20393–20406.
21. [21] M. Jaderberg, K. Simonyan, A. Zisserman, et al. (2015) Spatial transformer networks. Advances in neural information processing systems 28.
22. [22] L. Kaiser, M. Babaeizadeh, P. Milos, B. Osinski, R. H. Campbell, K. Czechowski, D. Erhan, C. Finn, P. Kozakowski, S. Levine, et al. (2019) Model-based reinforcement learning for Atari. arXiv preprint arXiv:1903.00374.
23. [23] D. Kalashnikov, A. Irpan, P. Pastor, J. Ibarz, A. Herzog, E. Jang, D. Quillen, E. Holly, M. Kalakrishnan, V. Vanhoucke, et al. (2018) Scalable deep reinforcement learning for vision-based robotic manipulation. In Conference on Robot Learning, pp. 651–673.
24. [24] K. Kim, J. Lanier, and R. Fox (2025) Make the pertinent salient: task-relevant reconstruction for visual control with distractions. In Reinforcement Learning Conference,
25. [25] M. Laskin, A. Srinivas, and P. Abbeel (2020) CURL: contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning, pp. 5639–5650.
26. [26] M. Laskin, K. Lee, A. Stooke, L. Pinto, P. Abbeel, and A. Srinivas (2020) Reinforcement learning with augmented data. Advances in Neural Information Processing Systems 33, pp. 19884–19895.
27. [27] J. W. Lee and H. Hwang (2025) Diffusion guided adaptive augmentation for generalization in visual reinforcement learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 880–889.
28. [28] K. Lee, K. Lee, J. Shin, and H. Lee (2020) Network Randomization: a simple technique for generalization in deep reinforcement learning. In Conference on Learning Representations,
29. [29] S. Liu, Z. Chen, Y. Liu, Y. Wang, D. Yang, Z. Zhao, Z. Zhou, X. Yi, W. Li, W. Zhang, et al. (2023) Improving generalization in visual reinforcement learning via conflict-aware gradient agreement augmentation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 23436–23446.
30. [30] V. Mai, K. Mani, and L. Paull (2022) Sample efficient deep reinforcement learning via uncertainty estimation. In International Conference on Learning Representations,
31. [31] L. McInnes, J. Healy, and J. Melville (2018) UMAP: uniform manifold approximation and projection for dimension reduction. arXiv preprint arXiv:1802.03426.
32. [32] S. Nair, A. Rajeswaran, V. Kumar, C. Finn, and A. Gupta (2023) R3M: a universal visual representation for robot manipulation. In Conference on Robot Learning, pp. 892–909.
33. [33] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer (2022) High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10684–10695.
34. [34] M. Schwarzer, A. Anand, R. Goel, R. D. Hjelm, A. Courville, and P. Bachman (2020) Data-efficient reinforcement learning with self-predictive representations. arXiv preprint arXiv:2007.05929.
35. [35] M. Schwarzer, J. S. O. Ceron, A. Courville, M. G. Bellemare, R. Agarwal, and P. S. Castro (2023) Bigger, better, faster: human-level atari with human-level efficiency. In International Conference on Machine Learning, pp. 30365–30380.
36. [36] Y. Seo, K. Lee, I. Clavera Gilaberte, T. Kurutach, J. Shin, and P. Abbeel (2020) Trajectory-wise multiple choice learning for dynamics generalization in reinforcement learning. Advances in Neural Information Processing Systems 33, pp. 12968–12979.
37. [37] W. Song, H. Choi, K. Sohn, and D. Min (2024) A simple framework for generalization in visual RL under dynamic scene perturbations. Advances in Neural Information Processing Systems 37, pp. 121790–121826.
38. [38] A. Stone, O. Ramirez, K. Konolige, and R. Jonschkowski (2021) The distracting control suite–a challenging benchmark for reinforcement learning from pixels. arXiv preprint arXiv:2101.02722.
39. [39] R. Sun, H. Zang, X. Li, and R. Islam (2024) Learning latent dynamic robust representations for world models. In Forty-first International Conference on Machine Learning,
40. [40] Y. Tassa, Y. Doron, A. Muldal, T. Erez, Y. Li, D. d. L. Casas, D. Budden, A. Abdolmaleki, J. Merel, A. Lefrancq, et al. (2018) DeepMind Control Suite. arXiv preprint arXiv:1801.00690.
41. [41] T. Wang, S. Du, A. Torralba, P. Isola, A. Zhang, and Y. Tian (2022) Denoised MDPs: learning world models better than the world itself. In International Conference on Machine Learning, pp. 22591–22612.
42. [42] R. Yang, J. Wang, Q. Peng, R. Guo, G. Wu, and B. Li (2025) Learning robust representations with long-term information for generalization in visual reinforcement learning. In The Thirteenth International Conference on Learning Representations,
43. [43] S. Yang, Y. Ze, and H. Xu (2024) MoVie: visual model-based policy adaptation for view generalization. Advances in Neural Information Processing Systems 36.
44. [44] D. Yarats, R. Fergus, A. Lazaric, and L. Pinto (2022) Mastering visual continuous control: improved data-augmented reinforcement learning. In International Conference on Learning Representations,
45. [45] D. Yarats, I. Kostrikov, and R. Fergus (2021) Image augmentation is all you need: regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations,
46. [46] W. Ye, S. Liu, T. Kurutach, P. Abbeel, and Y. Gao (2021) Mastering atari games with limited data. Advances in neural information processing systems 34, pp. 25476–25488.
47. [47] Z. Yuan, T. Wei, S. Cheng, G. Zhang, Y. Chen, and H. Xu (2024) Learning to manipulate anywhere: a visual generalizable framework for reinforcement learning. arXiv preprint arXiv:2407.15815.
48. [48] Z. Yuan, Z. Xue, B. Yuan, X. Wang, Y. Wu, Y. Gao, and H. Xu (2022) Pre-trained image encoder for generalizable visual reinforcement learning. Advances in Neural Information Processing Systems 35, pp. 13022–13037.
49. [49] Z. Yuan, S. Yang, P. Hua, C. Chang, K. Hu, and H. Xu (2024) RL-ViGen: a reinforcement learning benchmark for visual generalization. Advances in Neural Information Processing Systems 36.
50. [50] A. Zhang, R. T. McAllister, R. Calandra, Y. Gal, and S. Levine (2021) Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations,
51. [51] D. Zhang, B. Lv, H. Zhang, F. Yang, J. Zhao, H. Yu, C. Huang, H. Zhou, C. Ye, and C. Jiang (2024) Focus on what matters: separated models for visual-based RL generalization. Advances in Neural Information Processing Systems 37, pp. 116960–116986.
52. [52] X. Zhou, C. Ying, Y. Feng, H. Su, and J. Zhu (2025) Self-consistent model-based adaptation for visual reinforcement learning. arXiv preprint arXiv:2502.09923.
53. [53] C. Zhu, M. Simchowitz, S. Gadipudi, and A. Gupta (2023) RePo: resilient model-based reinforcement learning by regularizing posterior predictability. Advances in Neural Information Processing Systems 36, pp. 32445–32467.
54. [54] Y. Zhu, J. Wong, A. Mandlekar, R. Martín-Martín, A. Joshi, S. Nasiriany, and Y. Zhu (2020) Robosuite: a modular simulation framework and benchmark for robot learning. arXiv preprint arXiv:2009.12293.

Appendix Contents

Appendix A Impact Statement

This work presents VIGOR, an MBRL framework that achieves zero-shot visual generalization to unseen distractions while preserving sample efficiency. By integrating AW, DC, and ES, VIGOR addresses a critical bottleneck in deploying visual RL agents to environments with visual distribution shifts.

The primary impact is methodological: VIGOR provides a robust approach to stabilize both encoder representations and latent dynamics under distribution shifts. These advances contribute to the broader goal of developing RL systems that are less reliant on task-specific data collection and more resilient to the visual variability in real-world environments.

As our evaluations are conducted within simulated benchmarks (DMC and Robosuite), this work does not pose immediate direct societal or ethical risks. However, as these techniques move toward real-world deployment in robotics and autonomous systems, we emphasize that reliability and safety become paramount. Future applications in safety-critical domains should undergo rigorous validation and incorporate safeguards to prevent catastrophic decision errors arising from misaligned latent representations or unexpected environmental shifts.

Appendix B Training Procedure

This section details the training procedure of VIGOR. Since VIGOR leaves the planning mechanism unchanged, we focus on world model learning and omit planner updates in Algorithm 1.

For a trajectory segment \Gamma=(s_{i},a_{i},r_{i},s_{i+1})_{i=t}^{t+H} sampled from replay buffer \mathcal{B}, the reward and Q-value losses are defined as:

\mathcal{L}_{\text{rew}}(\theta)=\operatorname{CE}\bigl(\hat{r}_{t},r_{t}\bigr),\qquad\mathcal{L}_{Q}(\theta)=\operatorname{CE}\bigl(\hat{q}_{t},q_{t}\bigr),

where \hat{r}_{t}=R_{\theta}(z_{t},a_{t}) and \hat{q}_{t}=Q_{\theta}(z_{t},a_{t}) are the predicted reward and Q-value, respectively. The TD target is q_{t}=r_{t}+\bar{Q}_{\theta}\bigl(z_{t+1},\pi_{\theta}(z_{t+1})\bigr), where \bar{Q}_{\theta} is an exponential moving average of the Q-value, z_{t}=h_{\theta}(s_{t}) is the latent representation, and \pi_{\theta} is a policy trained via entropy maximization. \operatorname{CE} denotes the cross-entropy loss. For details on the world model architecture and planning procedure, we refer readers to Hansen et al. [17].

As shown in Algorithm 1, integrating VIGOR into a latent-space MBRL framework is straightforward: only the steps highlighted in yellow need to be added. In this work, we instantiate and validate VIGOR within the TD-MPC2 framework, without modifying its planner or optimization procedures. Extending VIGOR to alternative backbones (e.g., Dreamer-style world models) is a promising direction for future work.

Appendix C Implementation Details

C.1 Code Base

We evaluate VIGOR against four MFRL baselines (DrQ-v2 [44], SVEA [16], SRM [20], and SimGRL [37]) and four MBRL baselines (TD-MPC2 [17], DreamerV3 [12], HRSSM [39], and Dr. G [7]). We use the reference implementations from RL-ViGen [49]1 for DrQ-v2, SVEA, and SRM, and the official implementations from the respective authors for SimGRL2, DreamerV33, HRSSM4, and Dr. G5. VIGOR is built on the official TD-MPC2 repository6, with modifications restricted to the additional components introduced in Section 3.

C.2 Computational Resources

All experiments were conducted on a single NVIDIA A5000 GPU. On average, a complete training and evaluation run took approximately 48 hours for the DMC tasks and 72 hours for the Robosuite tasks.

C.3 Hyperparameters

To ensure a fair comparison, we reuse common settings across all baselines whenever possible and follow each method’s official configuration for method-specific values. Table 4 reports hyperparameters across baselines (with method-specific values listed where they differ), and Table 5 reports the full hyperparameter set for VIGOR. VIGOR inherits TD-MPC2 defaults [17] unless otherwise noted; VIGOR-specific parameters are highlighted in yellow.

Appendix D Evaluation Details

We follow the zero-shot evaluation protocol of Yuan et al. [49] and define seven distraction categories for DMC, each evaluated at two difficulty levels (easy and hard), resulting in a total of fourteen DMC evaluation settings. For Robosuite, we consider three difficulty levels (easy, hard, and extreme). These settings induce distribution shifts in appearance factors such as color, lighting, background, and camera position, providing a rigorous testbed for zero-shot visual generalization.

D.1 Environments and Tasks

We benchmark VIGOR on two widely used continuous-control suites: DMC and Robosuite. Representative environments are illustrated in Figure 9.

In DMC [40], we consider six visuomotor control tasks: cartpole-swingup, where the agent must swing up and balance a cartpole; finger-spin, which requires continuous spinning of a planar finger; walker-walk, which trains a bipedal walker to move forward stably; walker-stand, which focuses on maintaining balance in an upright posture; cheetah-run, where the agent learns high-speed forward locomotion; and reacher-easy, which involves controlling a 2-DoF arm to reach a target location.

For Robosuite [54], we evaluate four robotic manipulation tasks: door-opening, where the robot arm must open a hinged door; nut-assembly, which requires placing a round nut onto a peg; peg-in-hole, where a peg must be inserted into a narrow slot; and lifting, which involves grasping and lifting a cube.

D.2 Baseline Descriptions

This subsection provides detailed descriptions of the baselines used in our experiments, complementing the summary in Section 4.1.

Selection criteria.

We select baselines according to the following criteria: (i) The method must provide an official codebase in its original paper or in a recognized benchmarking framework (e.g., RL-ViGen [49]). We exclude methods whose scores can only be reproduced from unofficial or re-implemented code, as this may introduce unreliable or biased results. (ii) The method must explicitly target visual robustness or zero-shot visual generalization in RL. (iii) Reproduced scores must be broadly consistent with those reported in the original publication; methods with large unexplained discrepancies are excluded.

We additionally include three strong baselines (DrQ-v2, TD-MPC2, and DreamerV3) that do not explicitly target zero-shot visual generalization but serve essential reference roles: DrQ-v2 as a widely adopted MFRL baseline for sample efficiency, TD-MPC2 as VIGOR’s backbone, and DreamerV3 as a state-of-the-art reconstruction-based MBRL method. Including these methods enables direct assessment of how much zero-shot robustness VIGOR adds over its backbone and how general-purpose methods perform under distribution shifts without explicit robustness mechanisms.

MFRL baselines.

(i) DrQ-v2 [44] combines random-shift augmentation with an off-policy actor–critic architecture, achieving strong sample efficiency on pixel-based continuous control. (ii) SVEA [16] stabilizes off-policy Q-learning by maintaining a weighted mixture of augmented and clean Q-targets, preventing the value function from being corrupted by augmentation noise. (iii) SRM [20] applies spectrum-based augmentations in the frequency domain to enhance robustness to spatial corruptions. (iv) SimGRL [37] combines random-overlay with random-shift augmentation, injecting task-irrelevant background perturbations to mitigate observational overfitting.

MBRL baselines.

(v) TD-MPC2 [17] achieves strong performance via decoder-free latent-space planning with temporal-difference learning. It serves as VIGOR’s backbone. (vi) DreamerV3 [12] employs latent imagination with pixel-level reconstruction and achieves competitive performance across diverse control tasks. (vii) HRSSM [39] extends DreamerV3 with a bisimulation-based hybrid RSSM to extract task-relevant latent representations without pixel-level reconstruction. (viii) Dr. G [7] extends Dreamer [9] with dual contrastive objectives and a recurrent inverse-dynamics loss, explicitly targeting zero-shot visual generalization. It is the most directly comparable prior method.

D.3 Visual Distraction Categories

We consider seven types of distractions, each with two difficulty levels for DMC and three difficulty levels for Robosuite:

• object-color: Change the body color of the agent. Uniformly sample the parameters of the color (i.e., (r,g,b)) from the pre-defined distribution for each difficulty.
• background-color: Change the background color of the environment (e.g., terrain grid or sky). Uniformly sample the parameters of the color (i.e., (r,g,b)) from the pre-defined distribution for each difficulty.
• background-video: Overlay the background with the randomly sampled natural video. Sample a random video with the same width and height as the original image from a set of natural videos [38]. Overlay the video only to the background sky for easy and to all backgrounds, including the ground terrain other than the agent, for hard difficulty.
• cam-pos: Change the position of the tracking camera’s focus by randomly adding noise offset. Let the initial position of the tracking camera’s focus be X_{\text{cam}}=(c_{x},c_{y},c_{z}) in Euclidean space. Sample a random offset \delta\in\mathbb{R}^{3} from the uniform distribution with different bounds: \mathcal{U}(-0.08,0.08) for easy and \mathcal{U}(-0.15,0.15) for hard difficulty. Inject the offset to the initial position of the camera; X_{\text{cam}}=(c_{x}+\delta_{x},c_{y}+\delta_{y},c_{z}+\delta_{z}).
• light-position: Change the position and orientation of the tracking light of the agent. Following the approach used in [38], the tracking light’s coordinate is parameterized as the spherical coordinate; (\phi,\vartheta,\rho) where \phi is azimuth, \vartheta is inclination, and \rho is the radius of the sphere. Sample \phi from the normal distribution \mathcal{N}(\pi/6,1) for easy and \mathcal{N}(\pi/3,1) for hard difficulty. Sample \vartheta\sim\mathcal{N}(2\pi,1) and transform the initial pose of the tracking light X_{\text{light}} to (\phi,\vartheta,\rho), where \rho=\lVert X_{\text{light}}\rVert is the radius (Euclidean norm of the initial position).
• light-color: Change the color of the tracking light of the agent. Uniformly sample the parameters of the color (i.e., (r,g,b)) from the pre-defined distribution for each difficulty.
• moving-light: Rotate the tracking light around the agent. As in light-position, the spherical coordinate of the tracking light is randomly initialized as (\phi,\vartheta,\rho). Let the speed of azimuth rotation as \Delta_{\phi}=\pi/200 for easy and \Delta_{\phi}=\pi/100 for hard difficulty. Rotate the tracking light counterclockwise along the azimuth axis at every time-step; (\phi,\vartheta,\rho)\leftarrow(\phi,\vartheta,\rho)+(\Delta_{\phi},0,0).

D.4 Evaluation Protocol

For DMC, each of the seven distraction types listed above includes two difficulty levels: easy and hard. Robosuite follows the predefined evaluation splits (eval-easy, eval-hard, and eval-extreme) as introduced in RL-ViGen [49].

D.5 Evaluation Scenes

Representative evaluation scenes for DMC and Robosuite are shown in Figures 10(a) and 10(b), respectively. Agents are trained in the default clean setting and then evaluated zero-shot under diverse unseen distractions. Figure 11 illustrates the four Robosuite tasks (door-opening, nut-assembly, peg-in-hole, and lifting) under increasingly severe visual distractions, from eval-easy to eval-extreme.

Appendix E Supplementary Results

This appendix complements Section 4.2 with additional figures, tables, and analyses.

E.1 Zero-Shot Generalization

This subsection provides a detailed breakdown of our zero-shot generalization results, complementing the aggregate scores in Section 4.2. Full results across distraction types and tasks are reported in Figures 12 and 13, with per-task averages summarized in Table 6.

Performance across distraction types and difficulty levels.

Figure 12 reports zero-shot generalization on DMC across fourteen conditions (seven distraction types \times two difficulty levels: easy and hard) and on Robosuite across three difficulty levels (easy, hard, and extreme); each panel additionally shows the original (distraction-free) condition as a reference. Across all methods, returns tend to degrade monotonically as distraction difficulty increases (original \rightarrow easy \rightarrow hard \rightarrow extreme), reflecting the inherent challenge of distribution shifts in visual control. However, the magnitude of this degradation differs substantially across methods.

On DMC, baseline behavior diverges sharply across distraction types. TD-MPC2, despite being one of the strongest latent-space MBRL methods under clean settings, exhibits sharp drops under hard distractions on background-color ({\approx}0.32), background-video ({\approx}0.05), and light-color ({\approx}0.32), highlighting its sensitivity to severe visual shifts. DreamerV3 and HRSSM follow similar trajectories, with HRSSM degrading even under easy cam-pos perturbations. SimGRL maintains relatively stable performance across most distractions and achieves the highest score on background-video hard ({\approx}0.79), reflecting its training strategy that injects diverse background perturbations. In contrast, VIGOR sustains high returns across nearly all distraction types and difficulty levels, ranking first on every distraction except background-video. On Robosuite, where most baselines collapse to near-zero performance under hard and extreme shifts, VIGOR maintains the highest returns at every difficulty level, demonstrating its robustness under severe distribution shifts.

Task-level performance breakdown.

Table 6 and Figure 13 provide a detailed breakdown of zero-shot generalization performance across all tasks in DMC and Robosuite. VIGOR ranks first on every individual task in both benchmarks, demonstrating that its advantage is not driven by a few favorable cases but by uniformly strong performance across the task set.

On DMC, VIGOR achieves the highest average return (777.6), surpassing the second-best baseline (SRM, 752.0) by 3.4\%. At the task level, VIGOR attains the strongest scores on all six tasks, exceeding the respective second-best methods by 2.0\% on walker-walk (811.3 vs. DreamerV3, 795.5), 1.0\% on finger-spin (809.1 vs. SRM, 801.5), 4.0\% on cheetah-run (461.6 vs. SRM, 443.7), 0.5\% on walker-stand (938.6 vs. SVEA, 933.5), 2.7\% on reacher-easy (899.0 vs. SimGRL, 875.5), and 0.0\% on cartpole-swingup (746.0 vs. SRM, 745.7).

On Robosuite, the performance gains are substantially more pronounced. VIGOR achieves the highest average return (131.8), surpassing the second-best baseline (SRM, 91.8) by 43.6\%, and attains the best scores on all four tasks: door-opening (143.9, +55.6\% over SRM, 92.5), lifting (35.4, +31.1\% over DrQ-v2, 27.0), nut-assembly (15.5, +31.4\% over SVEA, 11.8), and peg-in-hole (332.6, +39.1\% over SRM, 239.1). Many baselines degrade to near-random performance in this setting, whereas VIGOR maintains robustness under visually challenging manipulation conditions.

Summary.

These findings confirm that VIGOR’s three components (AW, DC, and ES) are essential for mitigating severe distribution shifts. By jointly promoting augmentation-invariant latent dynamics, stabilizing transition predictions under perturbed inputs, and preventing encoder drift, VIGOR achieves the highest zero-shot generalization on every task across both benchmarks, not only in aggregate but also at the individual task level. This consistent dominance corroborates that conventional latent-space MBRL is fragile to OOD perturbations, whereas VIGOR’s latent-space consistency enables robust generalization without test-time adaptation.

E.2 Sample Efficiency

This subsection provides a detailed per-task comparison of sample efficiency, complementing the aggregate results in Section 4.2.

Training curves.

Figure 14 presents per-task learning curves for DMC and Robosuite, reporting IQM normalized scores across all methods. While no single method dominates across all tasks, VIGOR consistently ranks among the top performers. Notably, VIGOR maintains this level of efficiency despite being trained with more challenging objectives that incorporate visual perturbations, confirming that robustness does not come at the cost of efficiency.

Evaluation protocol.

To quantify sample efficiency, we adopt the evaluation protocol of Mai et al. [30]. We define TD-MPC2 as the oracle baseline, since it serves as VIGOR’s backbone and provides the most direct reference for assessing the efficiency cost of the added robustness components. We then measure the number of environment episodes required for each non-oracle method to reach 25\%, 50\%, and 75\% of TD-MPC2’s asymptotic return. If an agent fails to reach a threshold within the training budget, we assign the maximum episode count as a penalty (1{,}000 for DMC and 2{,}000 for Robosuite).

Quantitative results.

As shown in Table 7, VIGOR reaches every threshold in the fewest episodes among all non-oracle baselines on both benchmarks, demonstrating that the added robustness components incur minimal cost in sample efficiency relative to the TD-MPC2 backbone.

On DMC, VIGOR reaches the 25\%, 50\%, and 75\% thresholds in 395, 412, and 422 episodes, respectively. The second-best baseline at each threshold is Dr. G (415 at 25\%; 540 at 50\%) or DreamerV3 (553 at 75\%); VIGOR requires 4.8\% to 23.7\% fewer episodes than the second-best baseline at each threshold. On Robosuite, the efficiency advantage is more pronounced: VIGOR requires 655, 720, and 740 episodes to reach the three thresholds, while the second-best baseline (SRM) requires 1{,}240, 1{,}340, and 1{,}370 episodes, corresponding to roughly 1.9\times more interactions across all thresholds. Many baselines fail to reach the higher thresholds within the training budget, indicating that VIGOR’s sample-efficiency advantage amplifies under visually complex manipulation tasks. These results confirm that AW, DC, and ES preserve the strong sample efficiency of the TD-MPC2 backbone while enabling robustness to unseen distractions.

Statistical analysis.

Figure 15 presents empirical cumulative distribution functions (ECDFs) of episode returns across tasks. ECDFs capture the full distribution of training efficiency and illustrate how quickly each method reaches target performance thresholds. VIGOR exhibits ECDF profiles that closely track TD-MPC2 while surpassing all other baselines, further validating its efficiency–robustness balance.

Summary.

These results demonstrate that AW, DC, and ES preserve the sample efficiency of latent-space MBRL while enabling robust zero-shot generalization, incurring only modest overhead relative to the TD-MPC2 backbone.

E.3 DMC 100K

The main experiments (Section 4.2) train all methods for 1M gradient steps, a budget sufficient for most methods to reach asymptotic performance. To stress-test VIGOR’s efficiency–robustness trade-off under severe data constraints, we conduct an additional evaluation on DMC with only 100K gradient steps (10% of the standard budget), following prior work [22, 25, 34, 46, 35, 2].

Motivation.

Evaluating sample efficiency requires assessing two conditions: (i) how quickly an agent reaches its asymptotic performance, and (ii) how high that asymptote is. Under the full 1M-step budget, most methods converge, making it difficult to distinguish whether strong generalization stems from robust representations or simply from prolonged training. The 100K budget isolates early-phase learning dynamics: methods that build robust latent structures efficiently should generalize well even before convergence, while methods that rely on extended training to compensate for fragile representations will degrade disproportionately. Furthermore, zero-shot generalization performance is bounded by the training-time return, so any efficiency gap directly tightens the robustness ceiling, exacerbating the efficiency–robustness trade-off.

Experimental setup.

All methods are trained for 100K gradient steps on six DMC tasks under the same RL-ViGen protocol used in the main experiments (Section 4.1). Hyperparameters are identical to the 1M-step setting; no method-specific tuning is performed for the reduced budget. We evaluate both sample efficiency (IQM normalized scores as a function of training steps) and zero-shot generalization (IQM normalized scores under unseen visual distractions after 100K steps). All results are averaged over ten random seeds.

Results.

Figure 16 reports sample-efficiency curves (left) and zero-shot generalization (right).

• Sample efficiency. TD-MPC2 achieves the highest IQM normalized score ({\approx}0.80), with VIGOR closely following ({\approx}0.78). DreamerV3 ranks third ({\approx}0.61), followed by Dr. G ({\approx}0.58) and SVEA ({\approx}0.56). Two patterns are notable. First, the strongest sample-efficient methods are dominated by MBRL backbones (TD-MPC2, VIGOR, DreamerV3, and Dr. G), confirming that latent-space world models provide an efficiency advantage even under severe data constraints. However, not all MBRL methods benefit equally: HRSSM ({\approx}0.25) falls below all MFRL baselines under this tight budget, suggesting that its more complex world-model objective (bisimulation-based hybrid RSSM) is harder to optimize under limited interaction. Second, TD-MPC2 and VIGOR substantially outpace DreamerV3, Dr. G, and HRSSM, highlighting the efficiency of TD-MPC2’s MPC-based planning over recurrent world-model rollouts. VIGOR closely matches TD-MPC2’s trajectory throughout training despite the addition of AW, DC, and ES, indicating that these components introduce minimal training overhead (Appendix E.4).
• Zero-shot generalization. VIGOR achieves the strongest zero-shot generalization, substantially outperforming all baselines. The contrast with the sample-efficiency ranking is striking: TD-MPC2, the most efficient method, drops to mid-tier generalization performance. This confirms that sample efficiency alone does not guarantee robustness: without explicit consistency mechanisms, the representations learned under a constrained budget remain fragile to unseen visual shifts. VIGOR closes this gap: it retains near-TD-MPC2 efficiency while achieving the highest generalization, demonstrating that latent-space consistency is established early in training rather than requiring extended training to converge.

Summary.

These results demonstrate that VIGOR’s efficiency–robustness advantage is not an artifact of extended training. Even under a 10\times reduction in training budget, VIGOR achieves the strongest trade-off between sample efficiency and zero-shot generalization, confirming that latent-space consistency is built into the world-model representations from the early stages of training.

E.4 Wall-Clock Training Efficiency

To complement the sample-efficiency analysis, we evaluate VIGOR’s wall-clock training efficiency, which is important for assessing practical usability.

Experimental setup.

VIGOR and TD-MPC2 are trained under an identical RL-ViGen protocol on a single NVIDIA A5000 (24 GB) GPU, using the same dataloaders, environment settings, and hyperparameters described in Appendices C and D.

Overall results.

As shown in Figure 17 (left), VIGOR closely matches TD-MPC2 in sample efficiency: IQM normalized scores as a function of environment steps are nearly identical, and both methods reach asymptotic performance at similar sample budgets. Figure 17 (right) analyzes performance as a function of wall-clock time aggregated across all tasks. TD-MPC2 trains slightly faster during the early stages, but the difference remains modest throughout. Over a 24-hour window, the average normalized-score gap is approximately 0.17, and this gap generally narrows as training progresses.

Wall-clock overhead analysis.

The wall-clock overhead of VIGOR is modest and stems primarily from strong augmentations, which introduce additional preprocessing (e.g., background-video composition and external image loading). VIGOR mitigates this cost through its batch-splitting design: AW partitions a single batch into weak-only and weak-to-strong subsets rather than creating duplicate views, and all augmentations are applied before encoding. As a result, the encoder processes the same total number of samples as in the original TD-MPC2 pipeline, and the DC and ES losses add only lightweight computations on top of the existing latent representations. This design explains why VIGOR incurs only a minor slowdown despite incorporating substantially stronger generalization mechanisms.

Per-task analysis.

Figure 18 provides a per-task breakdown across all ten environments. TD-MPC2 maintains a small speed advantage in wall-clock time, but VIGOR remains close throughout training and the difference narrows as both methods converge. Table 8 summarizes the IQM normalized-score differences (TD-MPC2 minus VIGOR) at fixed wall-clock checkpoints. Most values are small, and the differences diminish toward later stages (at hour 14 the gap drops to 0.06), confirming that VIGOR’s overhead becomes less pronounced over time.

Summary.

VIGOR introduces only a modest wall-clock overhead relative to TD-MPC2, which diminishes as training progresses. These results confirm that VIGOR preserves practical training efficiency while delivering substantially improved zero-shot robustness.

E.5 Augmentation Ablation Studies

This subsection provides supplementary analyses for the augmentation ablations in Sections 4.3 and 4.4, which respectively present the AW ablation results (zero-shot generalization in Table 2; sample efficiency in Figure 5) and the augmentation-family experiment (Figure 7). Here, we focus on patterns not discussed in the main text: per-difficulty performance and variance characteristics of each AW variant, the \zeta-sweep’s cross-benchmark dynamics, sample-efficiency convergence patterns, and a qualitative comparison of the augmentation families used in the augmentation-agnostic experiment.

Per-difficulty and variance analysis of AW variants.

The main text reports aggregate performance for each variant (Table 2). Here, we examine the per-difficulty and variance patterns that reveal how each variant fails.

• DMC. The gap between variants widens from easy to hard. SO achieves 423.8\pm 421.9 (easy) and 374.9\pm 401.0 (hard), with standard deviations nearly equal to the means throughout, indicating that learning is highly unstable and seed-dependent. WTSO shows a similar pattern (550.6\pm 377.6 easy; 481.3\pm 358.3 hard) with large variance, confirming that persistent strong exposure without a stable anchor produces unreliable training. WO, by contrast, maintains lower relative variance (773.9\pm 246.0 easy; 606.6\pm 326.8 hard) and converges fast (Figure 5, left), but its generalization ceiling is bounded by limited perturbation diversity. VIGOR achieves both the highest returns and the lowest variance at both levels (863.9\pm 155.4 easy; 781.1\pm 261.2 hard), demonstrating that AW’s asymmetric splitting stabilizes optimization while maintaining robustness.
• Robosuite. The per-difficulty breakdown reveals a sharper separation, with all variants exhibiting monotonic degradation from easy to extreme. WO achieves moderate returns at the easy level (134.2\pm 143.8) but degrades steadily under harder shifts (112.6\pm 118.0 hard; 93.4\pm 96.9 extreme), confirming that weak augmentations alone provide insufficient robustness in visually complex manipulation. SO achieves higher easy-level returns (187.1\pm 182.2) but follows the same degradation pattern (137.3\pm 137.6 hard; 121.3\pm 122.6 extreme), with standard deviations consistently matching or exceeding the means, indicating highly unstable learning. WTSO improves at the easy level (297.6\pm 155.2) but similarly degrades under harder shifts (138.1\pm 131.3 hard; 130.1\pm 120.7 extreme; 79\% overall). Only AW maintains the highest returns at every difficulty level (332.0 easy; 236.7 hard; 146.0 extreme) with the strongest overall performance.

The cross-benchmark contrast is instructive: on DMC, WO preserves training efficiency but under-generalizes; on Robosuite, WO fails entirely because the visual complexity of manipulation scenes demands stronger perturbation exposure. Only AW’s asymmetric splitting simultaneously provides both signals: stability from the weak-only partition and diversity from the weak-to-strong partition.

Mixture-ratio sensitivity (\zeta-sweep).

The mixture ratio \zeta=|\mathcal{I}^{w}|/|\mathcal{I}^{ws}| controls the balance between weak-only and weak-to-strong samples within each batch. Table 2 shows that \zeta{=}4, \zeta{=}2, and \zeta{=}0.5 all underperform \zeta{=}1 on average, but the nature of the degradation differs across benchmarks:

• On DMC, increasing the weak-only fraction improves performance (\zeta{=}4: 96\% > \zeta{=}2: 94\% > \zeta{=}0.5: 88\%). Per-difficulty data show that this advantage stems primarily from the easy level, where \zeta{=}4 (854.4\pm 176.5) nearly matches VIGOR (863.9\pm 155.4). At the hard level, the gap widens (736.3\pm 274.3 vs. 781.1\pm 261.2), indicating that excessive weak-only bias sacrifices robustness under stronger shifts.
• On Robosuite, the trend reverses: \zeta{=}0.5 (87\%) outperforms \zeta{=}2 (80\%) and \zeta{=}4 (75\%). The advantage of \zeta{=}0.5 is most pronounced at the hard level (172.2\pm 157.4 vs. 144.7\pm 144.5 for \zeta{=}2 and 151.8\pm 156.4 for \zeta{=}4), confirming that severe shifts demand stronger perturbation exposure. Nevertheless, \zeta{=}0.5 still falls short of \zeta{=}1 at every level, indicating that excessive strong exposure impairs optimization even when robustness demands are high.

These results confirm that the balanced mixture (\zeta{=}1) is a deliberate design choice. The optimal ratio reflects a trade-off between training stability (favoring more weak-only samples) and generalization coverage (favoring more weak-to-strong samples), with \zeta{=}1 providing the best balance across both benchmarks.

Sample efficiency of AW variants.

Figure 5 reveals convergence patterns that complement the generalization analysis.

For the AW variants (left panel), AW and WO converge at comparable rates, confirming that the weak-only partition in AW preserves training stability. SO exhibits wide confidence intervals throughout training and fails to reach meaningful returns, while WTSO converges but to a substantially lower asymptote.

For the \zeta-sweep (right panel), all variants converge to similar asymptotic levels, but \zeta{=}0.5 exhibits slower early-phase convergence due to its higher proportion of strongly perturbed samples. These patterns reinforce the main-text finding that AW uniquely balances efficiency and robustness: alternative strategies sacrifice one or the other.

Qualitative comparison of augmentation families.

The augmentation-agnostic experiment in Section 4.4 replaces the default strong augmentation (random-overlay) with random-conv and SRM, two families from distinct perturbation domains. Figure 19 visualizes the qualitative differences on walker-walk:

• Random-overlay blends a randomly sampled natural image into the background, introducing contextual replacement that alters the global scene appearance while preserving the agent’s shape and texture.
• Random-conv applies a random convolution kernel to each frame, producing spatially smooth color and texture shifts without modifying scene structure. The perturbation is local and frequency-preserving.
• SRM (Spectrum Random Masking) masks contiguous bands in the Fourier spectrum, removing mid- and high-frequency image components. The resulting perturbation is global and frequency-destructive, qualitatively distinct from both overlay and convolution.

The rightmost column of Figure 19 shows the absolute pixel-wise difference between random-overlay and SRM, highlighting their distinct spectral characteristics. Despite these fundamental differences, all three variants maintain strong zero-shot generalization (Figure 7), confirming that latent-space consistency, not the augmentation choice, is the source of robustness.

Summary.

These supplementary analyses reinforce three conclusions from the main text. First, AW’s asymmetric batch splitting is the only augmentation strategy that preserves both training stability and zero-shot robustness across benchmarks with fundamentally different visual complexities; per-difficulty and variance analysis reveals that alternative strategies fail through either instability (SO, WTSO) or insufficient diversity (WO). Second, the balanced mixture ratio (\zeta{=}1) reflects a principled trade-off between stability and generalization coverage, with the optimal direction shifting across benchmarks but \zeta{=}1 providing the best compromise at every difficulty level. Third, VIGOR’s robustness is augmentation-agnostic: three qualitatively distinct perturbation families (contextual, textural, and spectral) all yield strong generalization, confirming that latent-space consistency, not the augmentation, drives robustness.

E.6 Latent-Space Consistency Analysis: DC and ES Ablation

To complement the quantitative ablation in Section 4.3 (Table 3), we visualize how removing DC and ES affects the structure of latent rollouts under visual perturbations. This analysis isolates the qualitative role of each component beyond aggregate return statistics.

Experimental setup.

We compare the full VIGOR with three ablated variants (w/o DC & ES, w/o DC, and w/o ES) on walker-walk under background-color and object-color distractions. For each variant, we roll out the learned world model from identical initial inputs and project the resulting latent states into 2D via UMAP [31]: the initial state z_{t} is encoded directly from the observation, while the subsequent states z_{t+1:t+H} are produced by the dynamics model. To make this distinction visible, we display z_{t} with red-edged markers in Figure 20. Episode returns and environment snapshots are shown alongside.

Results.

Figure 20 reveals distinct failure modes for each ablation that align with the quantitative patterns in Table 3.

• w/o DC & ES. Without either consistency mechanism, latent trajectories scatter widely even at the encoder level: the initial states z_{t} under easy and hard settings are far from those under original, and the rollouts diverge sharply. The return curves drop substantially under perturbations, with episode failures visible in the snapshots.
• w/o DC. ES alone aligns the encoder representations to some extent, but the dynamics rollouts diverge across difficulty levels. Although the initial states z_{t} are partially clustered, the dynamics model produces inconsistent multi-step predictions, illustrating that encoder-level alignment alone cannot prevent the dynamics from amplifying residual discrepancies across recursive rollouts.
• w/o ES. DC alone yields more compact rollout trajectories than the previous two variants, but the latent geometry is shifted across difficulty levels: the initial states z_{t} under perturbations are visibly displaced from their clean counterparts. While returns are higher than the previous variants, the misaligned encoder representations limit the benefit of DC’s dynamics-level constraint.
• VIGOR (full). Only the full model maintains both encoder-level alignment (initial states z_{t} across difficulty levels are tightly clustered) and dynamics-level consistency (subsequent rollouts z_{t+1:t+H} remain aligned across the full horizon). This two-level stability translates directly into stable returns and successful task execution across all evaluated difficulty levels.

Summary.

Only the full model maintains both encoder-level alignment (initial states z_{t} across difficulty levels are tightly clustered) and dynamics-level consistency (subsequent rollouts z_{t+1:t+H} remain aligned across the full horizon). This two-level stability translates directly into stable returns and successful task execution across all evaluated difficulty levels.

E.7 Ablation Study on Design Choices

This subsection provides formal definitions of the ablated variants introduced in Section 4.4 and detailed results across distraction types, tasks, and efficiency metrics.

Dynamic vs. consistent augmentations.

VIGOR applies augmentations at every step of the rollout horizon. For a trajectory segment s_{t:t+H}, the augmented views are:

\displaystyle s_{k}^{w}=\tau^{w}(s_{k},\upsilon_{k}^{w}),\quad s_{k}^{ws}=\tau^{s}(\tau^{w}(s_{k},\upsilon_{k}^{w}),\upsilon_{k}^{s}),\quad k=t,\dots,t{+}H.

If \upsilon_{k}^{w}=\upsilon_{t}^{w} and \upsilon_{k}^{s}=\upsilon_{t}^{s} for all k, the scheme is consistent: a single set of augmentation parameters is sampled and reused across all rollout steps. Otherwise, the scheme is dynamic: new parameters are sampled independently at every step. VIGOR uses the dynamic setting by default; the consistent variant is VIGOR_CONST_AUG.

Choice of strong augmentation.

VIGOR uses random-overlay by default. We compare against random-conv:

\displaystyle\tau^{s,\text{overlay}}(o,\tilde{o})=(1{-}\delta)\,o+\delta\,\tilde{o},\quad\tau^{s,\text{conv}}(o,w)=\mathrm{CONV}(o,w),

where o is the input image, \tilde{o}\sim\mathcal{D} is a task-irrelevant overlay, \delta=0.5, and w\sim\mathcal{N}(0,1) is a random convolution kernel. Both transforms are applied to every frame in the stacked observation s_{t}=\{o_{t},o_{t-1},\dots,o_{t-k+1}\}. We adopt random-overlay because it better approximates natural visual distractions at deployment, whereas random-conv applies synthetic pixel-level perturbations. The variant with random-conv is VIGOR_CONV. Figure 21 visualizes their qualitative differences.

Choice of auxiliary objective for ES.

VIGOR uses an alignment-based ES objective (adapted from SODA [18]) by default. As an alternative, we consider a contrastive ES objective (adapted from CURL [25]):

\displaystyle\mathcal{L}_{\text{contrast}}=-\log\frac{\exp(q^{\top}Wk_{+})}{\exp(q^{\top}Wk_{+})+\sum_{j=1}^{B}\exp(q^{\top}Wk_{j})},

where q=h_{\theta}(\tau^{w}(s_{t},\upsilon^{q})), k_{+}=h_{\theta^{-}}(\tau^{w}(s_{t},\upsilon^{k})), B is the size of mini-batch, and W is a learnable bilinear projection. The variant replacing alignment-based ES with contrastive ES is VIGOR_CONTRAST.

Zero-shot generalization.

Figure 22 reports zero-shot generalization across all distraction types. VIGOR exhibits smoother and more monotonic degradation (original \rightarrow easy \rightarrow hard) compared with the ablated variants, indicating stronger robustness to distribution shifts.

Among the variants, VIGOR_CONV performs competitively on color-related distractions (background-color, light-color, object-color), consistent with the color-invariance bias of random convolution. However, its performance drops sharply on background-video, where perturbations alter spatial and contextual cues rather than color. VIGOR_CONTRAST consistently underperforms across most distraction types.

Sample efficiency and aggregate metrics.

Figure 23 provides a comprehensive comparison. The aggregate statistics (top row) show that VIGOR achieves the best median, IQM, mean, and optimality gap, confirming its overall advantage. The performance profile (bottom-left) shows VIGOR’s dominance across the full distribution of normalized scores, and the learning curve (bottom-right) shows that VIGOR retains the sample efficiency of TD-MPC2 despite training under visual perturbations.

Qualitative comparison of strong augmentations.

Figure 21 illustrates the visual effects of the two strong augmentations on walker-walk. Random-overlay blends natural images into the background, introducing realistic domain shifts that approximate variations encountered in deployment. Random-conv applies convolutional filters that predominantly alter color statistics without structural changes. This distinction explains why VIGOR_CONV performs well on color-based perturbations but fails to generalize to shifts involving spatial or contextual variations.

Summary.

These ablations confirm that each design choice (dynamic augmentation, the choice of strong augmentation, and the auxiliary loss formulation) contributes meaningfully to VIGOR’s generalization. Altering any one of them leads to clear performance degradation, demonstrating that these choices collectively enable robust zero-shot generalization.

E.8 Latent-Space Consistency Analysis

This subsection provides detailed embedding visualizations that complement the latent-space analysis in Section 4.5. We examine whether the two-level consistency observed in the main text, at both the encoder and rollout levels, generalizes across multiple distraction types.

Experimental setup.

We compare VIGOR with its backbone TD-MPC2 [17] and Dr. G [7], the most directly comparable prior method for zero-shot visual generalization in latent-space MBRL. All agents are trained in clean environments and evaluated zero-shot under unseen visual distractions. For each agent, we roll out the learned world model from identical initial inputs to obtain a predicted latent trajectory z_{t:t+H}=(z_{t},z_{t+1},\dots,z_{t+H}), where each subsequent state is generated by the dynamics model as z_{t+h+1}=d_{\theta}(z_{t+h},a_{t+h}), and we record the episode return accumulated during evaluation. To visualize how each world model evolves over consecutive timesteps, we additionally depict the preceding rollout z_{t-1:t+H-1} as a lighter trajectory, providing temporal context. Comparing the two trajectories reveals whether the world model maintains consistent rollouts as the encoder processes successive observations under visual perturbations.

Results.

We project the high-dimensional latent trajectories into 2D via UMAP [31], with faded markers indicating earlier time steps. Episode returns and environment snapshots across evaluation levels (original, easy, hard) are shown alongside.

Figure 24 presents results for walker-walk under three perturbation types: background-color, background-video, and light-color. Across all three types, VIGOR exhibits consistency at both levels of the vulnerability identified in Figure 1. At the encoder level, the initial latent states extracted under original, easy, and hard settings are located within the same region of the latent manifold, indicating that the encoder maps visually distinct observations to similar task-relevant representations. At the rollout level, the subsequent latent trajectories z_{t+1:t+H} produced by the dynamics model remain tightly aligned across difficulty levels, confirming that the dynamics model preserves this consistency throughout the planning horizon. This two-level stability translates directly into stable episode returns across all settings.

In contrast, TD-MPC2 and Dr. G exhibit divergent patterns at both levels. Their initial encoder outputs scatter under perturbations, and this divergence is amplified through latent rollouts, producing trajectories that deviate substantially from the clean setting. The resulting return degradation and task failures (visible in both the episode-return curves and the environment snapshots) illustrate the two-level vulnerability in practice.

Summary.

These visualizations provide further empirical support for the hypothesis (Figure 1) that mapping OOD inputs onto the task-relevant latent manifold is essential for zero-shot generalization. Across all three perturbation types, VIGOR consistently preserves consistency at both the encoder and rollout levels, a property absent in the baselines.

E.9 Performance under Longer Planning Horizons

Our default experiments adopt a short MPC planning horizon of H{=}3, following standard practice in latent-space MPC. To examine whether VIGOR’s latent-consistency mechanism remains effective under deeper imagined rollouts, we conduct comparisons at extended planning horizons.

Experimental setup.

Longer planning horizons substantially increase wall-clock time because each optimization step requires deeper latent rollouts. Running H{=}5 or H{=}10 across all DMC and Robosuite tasks would therefore incur prohibitive computational cost. We thus evaluate H\in\{3,5,10\} on one representative task per benchmark: finger-spin on DMC and door-opening on Robosuite. All other hyperparameters remain fixed to isolate the effect of the planning horizon.

Results and discussion.

Figure 25 shows learning curves for VIGOR under the three horizons. Across H{=}3, H{=}5, and H{=}10, VIGOR maintains stable optimization and reaches similar asymptotic returns. As expected, increasing the horizon slows convergence because compounding prediction errors accumulate over longer imagined trajectories, a well-documented limitation of MBRL. Nevertheless, VIGOR’s latent-consistency regularization continues to stabilize the forward dynamics, preventing error accumulation from degrading policy learning even when the planning depth is tripled.

Summary.

These results confirm that VIGOR’s latent-consistency mechanism generalizes beyond the default short-horizon setting. By mitigating the accumulation of prediction errors, VIGOR maintains stable learning and competitive performance under deeper imagined rollouts.

E.10 Generalization to Unseen Camera Viewpoints

While RL-ViGen includes camera-position shifts (translations), these preserve the base camera orientation. To investigate whether models trained with appearance-focused perturbations can generalize to fundamentally different geometric transformations, we evaluate all methods on unseen camera rotations that modify projection geometry and the relative orientation between the agent and the camera.

Viewpoint-rotation settings.

We consider two rotational viewpoint distractions from DMControl-GB [18]:

• cam-ori (static rotation). At the start of each episode, the tracking camera’s spherical parameters c=(\phi,\theta,r,\theta_{\text{roll}}) are sampled uniformly within bounds controlled by a difficulty parameter \beta_{\text{cam}}\in\left[0,1\right] (easy: 0.3; hard: 0.8). Larger \beta_{\text{cam}} permits wider rotation and distance ranges ( -\frac{\pi}{2}\beta_{\text{cam}}\leq\phi,\theta,\theta_{\text{roll}}\leq\frac{\pi}{2}\beta_{\text{cam}} and 0.5\ r_{\text{original}}\beta_{\text{cam}}\leq r\leq 2.5\ r_{\text{original}}\beta_{\text{cam}}), producing substantial viewpoint deviations without temporal variation.
• cam-dyn (dynamic rotation). Camera parameters are initialized as in cam-ori and then evolve via a bounded random walk. Camera velocities are sampled as v_{0}\!\sim\!\mathcal{U}(v_{\min},v_{\max}),\quad v_{n}=v_{0}+\sum_{j=1}^{n}\mathcal{N}(0,\sigma\Sigma), (11) where n indexes the simulation step within the episode, \sigma is a scalar noise scale, and \Sigma is the per-axis covariance, generating continuous, temporally coherent viewpoint drift throughout the episode. Motion magnitudes scale with \beta_{\mathrm{cam}}, producing time-varying geometric distortions.

As illustrated in Figure 26, these viewpoint perturbations alter perspective, foreshorten limbs, reorder depth relationships, and in some cases cause the agent to leave the field of view entirely (e.g., cartpole-swingup).

Zero-shot evaluation.

All models are trained under the RL-ViGen protocol (without any viewpoint-rotation augmentations) and evaluated zero-shot on both settings.

Results and discussion.

As shown in Figure 27, performance drops sharply across all methods, and none maintains reliable zero-shot control under camera rotations. This degradation arises because rotational viewpoint changes induce geometric transformations that differ fundamentally from appearance-based perturbations (e.g., lighting, background, or color). Such transformations violate the pixel-level and feature-level invariances implicitly learned by current encoder-based methods, including VIGOR.

While VIGOR is designed to improve robustness against appearance-based perturbations (background videos, lighting changes, color shifts, and camera translations), these results indicate that geometric viewpoint variations pose a more challenging generalization problem. Recent work has explored Spatial Transformer Networks (STNs) [21] for few-shot viewpoint adaptation in visual RL [43, 47], demonstrating improved spatial robustness by explicitly modeling spatial transformations. However, achieving simultaneous robustness to both appearance and geometric perturbations in a fully zero-shot setting remains largely unexplored. Developing geometry-aware latent representations, multi-view consistent world models, or scene-structured encoders constitutes an important direction for future work.

Summary.

This evaluation reveals a clear limitation of current visual RL generalization methods. Although VIGOR significantly improves robustness to appearance-based perturbations, all evaluated methods struggle with the geometric projection changes induced by viewpoint rotation. These findings motivate future work on geometry-aware representations or multi-view consistent world models to achieve viewpoint-invariant control.

Appendix F MBRL Baseline Analysis

This section provides an extended discussion of the MBRL baselines used in our experiments. Since VIGOR is designed as a framework that can be integrated into existing latent-space MBRL methods, understanding the characteristics of representative backbones is important. We focus on two influential families of latent-space MBRL, Dreamer and TD-MPC, which are widely regarded as state-of-the-art for pixel-based visuomotor control.

Dreamer family.

The Dreamer family [9, 11, 12] learns a decoder-based latent world model for long-horizon control from pixels. The model is trained on latent dynamics, reward prediction, and pixel-level reconstruction, forcing the latent state to retain sufficient information for observation recovery. The initial latent state is inferred through a recurrent state-space model (RSSM [10]) conditioned on past states, actions, and current observations. Actions are generated by a learned actor policy optimized via actor–critic RL in latent space, with the critic trained on imagined rollouts. DreamerV2 [11] and DreamerV3 [12] scale the architecture and introduce improved optimization strategies, making the family increasingly competitive on high-dimensional pixel-based control tasks. More recently, DreamerV4 [13] extends the family with a scalable Transformer architecture for long-horizon, language-conditioned tasks; however, it is trained from purely offline interaction data and thus falls outside the online RL setting considered in this work.

TD-MPC family.

The TD-MPC family [14, 17] takes a different approach by discarding reconstruction losses and adopting a decoder-free latent world model. It jointly learns latent dynamics, reward, value, and a policy under a temporal-difference objective, operating entirely in latent space. Unlike the Dreamer family, which trains a separate actor for long-horizon imagination, the TD-MPC family integrates its policy into an online planning scheme: at each step, a model predictive controller searches over action sequences guided by Q-value predictions, using the learned policy as a prior, and executes the first action of the best trajectory. TD-MPC2 improves upon this formulation with architectural refinements and multi-task scalability, achieving strong sample efficiency on continuous visuomotor control benchmarks.

Comparison of the two families.

Both families share the core principle of learning a transition model over compact latent representations. Their fundamental difference lies in how the world model is trained and how actions are selected: the Dreamer family relies on a decoder-based model and learns an actor policy via RL, while the TD-MPC family uses a decoder-free model and selects actions through MPC planning guided by Q-values. A further distinction concerns the role of imagined rollouts at deployment: the Dreamer family generates synthetic latent trajectories only during training for actor–critic optimization, while the actor directly selects actions at test time without rollouts. The TD-MPC family, by contrast, generates latent rollouts at both training and test time: during training for world-model learning and at test time for MPC planning. This distinction reflects two design philosophies, decoder-based actor learning versus decoder-free latent planning, with implications for sample efficiency and generalization.

Rationale for backbone selection.

Since VIGOR aims to enhance visual generalization without sacrificing sample efficiency, building upon an efficient backbone is essential. Dreamer family’s reconstruction-driven formulation increases model complexity and training cost, while TD-MPC family’s decoder-free design focuses directly on value learning and planning in latent space, yielding stronger sample efficiency in practice. As shown in Figure 28, TD-MPC2 outperforms both the Dreamer family and its predecessor TD-MPC across diverse tasks in DMC, consistent with prior findings [17]. We therefore adopt TD-MPC2 as VIGOR’s backbone. Section 4.2 demonstrates that VIGOR preserves this efficiency advantage while substantially improving zero-shot generalization.

Appendix G Extended Related Work

G.1 Visual Generalization in MFRL

This subsection provides a detailed survey of MFRL visual generalization methods, complementing the summary in Section 2. Baseline selection criteria are detailed in Appendix D.2; methods not selected as baselines are discussed below for context.

Data-augmentation approaches.

Data augmentation has emerged as a simple yet effective strategy for visual robustness in MFRL. RAD [26] systematically evaluates a suite of image transformations for RL, demonstrating that augmentations such as random crop and color jitter improve generalization. DrQ [45] and DrQ-v2 [44] integrate random-shift augmentation into off-policy actor–critic training, achieving strong sample efficiency on pixel-based control. SVEA [16] addresses the instability of augmented Q-learning by maintaining a weighted mixture of augmented and clean Q-targets, stabilizing training under diverse perturbations. SRM [20] applies spectral perturbations in the frequency domain to improve robustness to spatial corruptions. Network Randomization [28] passes observations through a randomly initialized convolutional layer, providing an implicit form of style augmentation; this principle inspires our random-conv ablation in Section 4.4. SimGRL [37] tackles imbalanced saliency via feature-level frame stacking and mitigates observational overfitting via a shifted random-overlay augmentation, which is applied uniformly to every sample without batch partitioning. DG2A [27] leverages a pretrained diffusion model (Stable Diffusion [33]) for adaptive data augmentation, generating augmented images guided by saliency maps.

Representation-learning approaches.

A complementary direction aims to learn visual representations that are inherently robust to task-irrelevant variations. CURL [25] applies contrastive learning to augmented observations, encouraging the encoder to capture task-relevant features. SODA [18] aligns the encoding of an augmented observation with that of its non-augmented counterpart through a soft alignment loss, promoting augmentation-invariant features without contrastive negatives. VIGOR’s ES component adapts SODA’s alignment objective to a structurally different role: in MFRL, alignment promotes augmentation-invariant features at the encoder, whereas in VIGOR it stabilizes the encoder against the cross-augmentation supervision induced by DC’s dynamics-level training (Section 3.3). SGQN [4] combines saliency-guided attention masks with self-supervised objectives, encouraging the encoder to focus on action-relevant regions. SMG [51] exploits an auxiliary image reconstruction task via separated model branches for encoders and decoders, forcing the encoder to extract task-relevant representations from foreground pixels. ALDA [3] disentangles the latent representation and projects OOD observations back to the training distribution using associative memory. ROUSER [42] learns robust action-value representations based on the information bottleneck framework, leveraging one-step rewards and compressed reward representations.

Pretrained-encoder approaches.

Recent work leverages visual representations pretrained on large-scale datasets. PIE-G [48] uses frozen pretrained encoders to provide generalization-friendly features, bypassing the need for augmentation-based robustness. R3M [32] pretrains visual representations on egocentric video data using time-contrastive learning for downstream robotic manipulation. MoVie [43] adapts pretrained vision models for RL via multi-view representation learning. These methods are orthogonal to VIGOR and could potentially be combined with our framework to further improve robustness.

Gradient- and optimization-based approaches.

CG2A [29] identifies conflicting gradients between augmented views and aggregates them via conflict-aware gradient surgery, improving generalization stability. PAD [15] uses a self-supervised auxiliary task at test time to adapt the encoder to new visual conditions, an approach that requires online adaptation and thus falls outside the zero-shot setting targeted by VIGOR.

Relationship to VIGOR.

All MFRL methods above address robustness at the single-step level: the encoder feeds into a Q-value or policy without recursive dynamics. As identified in Section 1, MBRL exhibits a two-level vulnerability: visual distractions first perturb encoder outputs, and these perturbations then compound through recursive latent rollouts. This structural difference means that augmentation or regularization applied to the encoder alone is insufficient: the dynamics model must also produce consistent predictions under distribution shifts. VIGOR addresses both levels: ES stabilizes the encoder (level one), while DC constrains the recursive dynamics pathway (level two), with AW providing the paired latent structure that enables both constraints. The closest MFRL counterpart, SimGRL, uses random-overlay augmentation but applies it uniformly to every sample without batch splitting, leaving both the dynamics-level vulnerability and cross-augmentation supervision unaddressed.

G.2 Extended Discussion of MBRL Visual Generalization

Invariant MBRL.

Several methods aim to learn distraction-invariant representations within the MBRL framework. DBC [50] uses bisimulation metrics to learn representations that group visually different but behaviorally equivalent states. Denoised MDP [41] identifies and removes noise variables from the latent representation. RePo [53] regularizes the encoder via a reconstruction-based objective that selectively preserves task-relevant information. HRSSM [39] learns a bisimulation-based hybrid RSSM in DreamerV3 [12] to extract endogenous, task-relevant latent representations without pixel-level reconstruction. SD [24] uses segmentation masks as a structural prior for robust representation learning, building upon the DreamerV3 architecture. While these methods improve robustness, several require test-time adaptation or additional environment interactions in the target domain, limiting their applicability in the zero-shot setting targeted by VIGOR.

Dr. G and its relationship to VIGOR.

Dr. G [7] is the most directly comparable prior work, as it explicitly targets zero-shot visual generalization in MBRL. Built on Dreamer [11], Dr. G applies weak and strong augmentations to create two independent views of every sample and trains dual contrastive losses (reality–reality at the encoder level and dream–reality at the dynamics level) alongside a recurrent inverse-dynamics objective.

VIGOR differs from Dr. G in three key aspects: (i) Efficiency: AW partitions a single batch into weak-only and weak-to-strong subsets via asymmetric splitting, where the weak-only subset directly serves as the stable anchor for both DC’s transition target and ES’s alignment target. This anchor-and-augment structure requires only a single encoder pass per sample. Dr. G, by contrast, generates two fully independent views per sample to construct contrastive pairs, doubling the encoder computation. (ii) Mechanism: Both methods address visual robustness in the latent-space MBRL pipeline, but through different means. Dr. G uses contrastive losses requiring negative samples and an inverse-dynamics auxiliary loss. VIGOR enforces consistency through direct latent regression (DC) based on asymmetric batch splitting (AW), combined with encoder stabilization (ES), without contrastive negatives or auxiliary supervision. (iii) Backbone: VIGOR builds on TD-MPC2’s decoder-free, MPC-based architecture. Dr. G builds on Dreamer’s recurrent world model (RSSM) but deliberately omits the decoder, relying instead on contrastive and inverse-dynamics objectives to learn latent representations.

Our experiments show that VIGOR achieves stronger zero-shot generalization with higher sample efficiency (Section 4.2), suggesting that direct dynamics-level regression combined with encoder stabilization provides a more effective inductive bias than contrastive learning for this setting.

Appendix H Discussion of Method Components

This section discusses why VIGOR’s three components (AW, DC, and ES) require careful adaptation for the MBRL setting, how they are implemented, and how they interact. While related ideas have appeared in MFRL, directly applying them to latent-space MBRL is non-trivial due to the two-level vulnerability identified in Section 1: visual distractions first perturb encoder outputs, and these perturbations then compound through recursive world-model rollouts. We highlight the resulting design challenges and clarify how each component addresses them.

AW.

The asymmetric batch-splitting design of AW, together with the rationale for routing every sample through the weak augmentation \tau^{w} before applying \tau^{s} to half the batch, is described in Section 3.1. Two consequences are worth highlighting in the context of the two-level vulnerability. First, the paired weak-only and weak-to-strong latents produced within a single mini-batch are precisely the supervision signal consumed by DC; without this pairing, DC would require a second encoder pass per sample, doubling encoder computation. Second, by applying strong augmentation only to a subset, AW balances two opposing failure modes observed in our ablations (Section 4.3; Appendix E.5): weak-only ablations exhibit insufficient robustness due to limited augmentation diversity, while strong-only and weak-to-strong-only ablations exhibit optimization instability under heavy perturbation.

DC.

DC directly addresses the dynamics level of the two-level vulnerability. Perturbation-induced drift in the latent state corrupts the entire synthetic trajectory produced by the world model, leading to erroneous planning decisions. This problem is amplified relative to MFRL because both Q-values and transition dynamics are conditioned on encoder outputs, so a single perturbed latent propagates errors across the full rollout horizon.

To mitigate this, VIGOR enforces latent consistency by anchoring the dynamics model to weakly augmented successor states. While the dynamics model is trained on perturbed mixed latents (z_{t}^{m}) generated by AW, the DC objective regresses these predictions toward targets produced by weak augmentations, which serve as nearly noise-free anchors. This ensures that the world model produces consistent next-state predictions regardless of input perturbations, preventing the compounding of rollout errors over long horizons. By aligning predictions from OOD inputs with their in-distribution counterparts, DC stabilizes the forward dynamics and supports robust planning under visual distribution shifts.

ES.

ES addresses the encoder level of the two-level vulnerability. While ES adapts the alignment objective of SODA [18], originally proposed for MFRL representation learning, its role in VIGOR is tailored to a specific failure mode in latent-space MBRL: encoder drift induced by DC’s cross-augmentation supervision. A single shared encoder produces the latent embeddings that concurrently feed the dynamics model, the reward predictor, Q-value networks, and the MPC planner. The cross-augmentation supervision imposed by DC, which trains the dynamics on mixed latents (z_{t}^{m}), can inadvertently distort the encoder’s representation. Without explicit constraints, encoder outputs tend to drift under distribution shifts, causing the dynamics-level compounding that DC is designed to prevent.

To mitigate this, VIGOR introduces an auxiliary ES loss that aligns weak-only and weak-to-strong encodings of the same frame. In VIGOR, ES acts as a structural stabilizer for the shared encoder, enabling DC to operate on consistent representations without distorting the latent geometry. The encoder thereby preserves task-relevant features under both clean and perturbed inputs, ensuring that initial latent states remain grounded for long-horizon planning in unseen environments.

Summary.

VIGOR’s effectiveness arises from the synergy between its three components, each targeting a different level of the world-model pipeline: AW establishes an asymmetric batching structure that provides paired latents without additional overhead; DC stabilizes the forward dynamics by anchoring predictions from perturbed mixed latents to weakly augmented targets, mitigating compounding rollout errors; and ES prevents the shared encoder from drifting under cross-augmentation pressure, ensuring that DC’s dynamics-level constraint operates on stable representations. The three components are complementary: our ablations (Sections 4.3–4.4; Appendix E.5–E.7) demonstrate that removing any single component results in significant performance degradation, with the largest drops observed on Robosuite manipulation tasks under severe distribution shifts.