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

Motivation

• Visual distractions push encoder outputs out of distribution, and recursive latent rollouts compound the resulting errors across the planning horizon.This two-level vulnerability makes single-step robustness methods insufficient for an encoder that also serves the dynamics model, Q-value networks, and planner.Source: §1 · pp. 1–2

Approach

• VIGOR combines asymmetric weak-to-strong augmentation, dynamics-level consistency, and encoder-level stabilization to target the two-level vulnerability.The components produce paired latent views, anchor dynamics outputs to stable targets, and prevent encoder drift under cross-augmentation supervision.Source: §1 · p. 2

Ablations

• Uniform strong augmentation without asymmetric batch splitting achieves only 48% of VIGOR’s performance on DMC.The reported ablations also find that VIGOR’s robustness is augmentation-agnostic: latent-space consistency, rather than a particular augmentation choice, drives generalization.Source: §1 · p. 2

Results

• VIGOR surpasses the second-best baseline by 3.4% on DMC and 43.6% on Robosuite while retaining TD-MPC2’s sample efficiency.The evaluation covers seven visual distraction types under the RL-ViGen zero-shot protocol.Source: §1 · p. 2

2.1 Visual Model-Based Reinforcement Learning

• Latent-space MBRL learns compact world models for planning, with Dreamer using pixel reconstruction and TD-MPC using decoder-free planning.Both families assume i.i.d. training and test distributions; unseen visual distractions can push encoder outputs out of distribution and compound rollout errors.Source: §2.1 · p. 3

2.2 Visual Generalization in Model-Free Reinforcement Learning

• Model-free visual-robustness methods use augmentation, self-supervised representation learning, and pretrained encoders, operating without recursive latent rollouts.The distinction motivates VIGOR’s two-level design for MBRL, where encoder perturbations propagate through multi-step rollouts.Source: §2.2 · p. 3

2.3 Visual Generalization in Model-Based Reinforcement Learning

• Invariant MBRL methods typically require test-time adaptation, whereas VIGOR targets zero-shot generalization without adaptation data.Unlike Dr. G’s contrastive objectives and inverse-dynamics auxiliary loss, VIGOR uses direct latent regression for dynamics-level consistency, supported by asymmetric augmentation and encoder stabilization.Source: §2.3 · p. 3

3 VIGOR

• VIGOR’s three interdependent components target input augmentation, latent dynamics, and encoder representation without modifying the underlying planner or optimization procedure.AW supplies paired latents for DC, while ES counteracts encoder drift that DC can induce.Source: §3 · p. 3

3.1 Asymmetric Weak-to-Strong Augmentation (AW)

• AW places weak-only and weak-to-strong views in one mini-batch, concatenates them, and jointly encodes them with the shared encoder.The two branches apply weak augmentation alone or weak followed by strong augmentation, preserving the original batch size.Source: §3.1 · p. 3

3.2 Dynamics-Level Consistency (DC)

• DC trains dynamics predictions from mixed inputs to match weakly augmented successor targets, enforcing augmentation-invariant transitions over the rollout horizon.Weak augmentation supplies a stable supervisory signal while the initial perturbation is propagated through recursive predictions.Source: §3.2 · pp. 4–5

3.3 Encoder-Level Stabilization (ES)

• ES aligns weak-only and weak-to-strong encodings of the same frame to counter encoder drift induced by DC’s cross-augmentation supervision.The encoder is shared with the dynamics model, Q-value networks, and MPC planner, so drift can propagate through multi-step rollouts.Source: §3.3 · p. 5

3.4 Overall Training Objective

• The overall objective adds VIGOR losses to the TD-MPC2 world-model objective, which includes reward and Q-value losses.The discount factor downweights losses at later horizon steps.Source: §3.4 · p. 5

4.1 Experimental Setup

• Evaluation covers six DMC tasks and four Robosuite manipulation tasks, with agents trained on clean observations and tested zero-shot under seven distraction types.The protocol spans easy and hard difficulty levels in DMC and easy, hard, and extreme levels in Robosuite.Source: §4 · p. 5, §4.1 · p. 6, §4.2 · p. 6

4.2 Experimental Results

• VIGOR beat SRM by 3.4% in DMC mean and 43.6% in Robosuite mean, leading both benchmarks in zero-shot generalization.Table 1 aggregates episode returns across tasks, distraction types, and difficulty levels.Source: Table 1 · p. 6, §4.2 · p. 6
• VIGOR closely matches TD-MPC2’s rapid DMC convergence and converges comparably on Robosuite, preserving its backbone’s sample efficiency.On DMC, VIGOR requires fewer samples than all other baselines; on Robosuite, other baselines lag in speed and asymptotic returns.Source: §4.2 · pp. 6–7, Figure 4 · p. 7

4.3 Ablation Studies

• Removing DC reduces generalization to 89% on DMC and 46% on Robosuite, while removing ES yields 93% and 88%, respectively.Removing both losses reduces performance to 76% on DMC and 33% on Robosuite; the authors describe DC and ES as complementary, with DC dominant under severe shifts.Source: Table 3 · p. 8, §4.3 · p. 8
• AW succeeds on both benchmarks, while the preferred weak-only fraction reverses across them; the balanced mixture, ζ=1, performs best on both.Strong-only reaches 48% of VIGOR on DMC, while weak-to-strong-only reaches 79% on Robosuite; the reported comparisons use VIGOR (AW, ζ=1) as reference.Source: §4.3 · pp. 7–8, Table 2 · p. 7

4.4 Ablation on Key Design Choices

• VIGOR outperforms variants with fixed augmentation parameters, random-conv replacing random-overlay, or contrastive ES replacing alignment-based ES.The authors attribute these comparisons to the importance of dynamic resampling, stronger robustness from random-overlay than convolution-based perturbations, and alignment-based ES complementing AW and DC.Source: §4.4 · pp. 8–9, Figure 6 · p. 8
• Replacing random-overlay with random-conv or SRM preserves strong zero-shot generalization, indicating robustness is augmentation-agnostic.The three variants use qualitatively distinct perturbation families.Source: Figure 7 · p. 9, §4.4 · p. 9

4.5 Latent-Space Consistency Analysis

• Across difficulty levels, VIGOR’s perturbed latent trajectories remain compact and close to clean trajectories, while TD-MPC2 and Dr. G show scattered or divergent rollouts.The reported consistency holds both for initial encoder outputs and throughout rollouts, with stable returns and verification up to planning horizon H=10.Source: Figure 8 · p. 9, §4.4 · p. 9, §4.5 · p. 9

5 Conclusion

• VIGOR surpasses the second-best baseline by 3.4% on DMC and 43.6% on Robosuite while preserving TD-MPC2 sample efficiency.Ablations indicate that all three components are jointly necessary and that robustness is augmentation-agnostic.Source: §5 · p. 10
• VIGOR targets appearance-based perturbations; unseen camera rotations remain challenging, motivating geometry-aware representations as future work.Extending the framework to dynamics- and task-level shifts and validating it on real robots are also identified as future directions.Source: §5 · p. 10

A Impact Statement

• VIGOR stabilizes encoder representations and latent dynamics under distribution shifts, addressing a bottleneck in deploying visual RL agents in visually variable environments.Source: A Impact Statement · p. 15
• The work contributes to the broader goal of RL systems less reliant on task-specific data collection and more resilient to real-world visual variability.Source: A Impact Statement · p. 15
• Evaluations are limited to simulated DMC and Robosuite benchmarks; safety-critical real-world applications would require rigorous validation and safeguards.Source: A Impact Statement · p. 15

B Training Procedure

• VIGOR leaves the planning mechanism unchanged, so the training procedure focuses on world-model learning and omits planner updates.Its integration into TD-MPC2 does not modify the planner or optimization procedures.Source: B Training Procedure · p. 15
• The base world-model objective includes reward-prediction and Q-value losses, augmented with the dynamics-consistency loss.The per-step reward and Q-value are predicted from latent state and action; the text identifies the latent representation and policy training via entropy maximization.Source: B Training Procedure · p. 15
• Each training iteration samples a replay-buffer trajectory segment and accumulates VIGOR losses across horizon steps, downweighted at later steps by λ.The algorithm inputs include horizon H, weak and strong augmentation functions, learning rate η, and momentum coefficient κ.Source: B Training Procedure · p. 15
• The algorithm updates online network parameters by gradient descent and updates target parameters using a momentum coefficient.Source: B Training Procedure · p. 15

C.1 Code Base

• VIGOR is built on the official TD-MPC2 repository, with modifications restricted to the components introduced by the method.The evaluation includes four model-free and four model-based baselines, using reference or official implementations.Source: §C.1 · p. 16

C.2 Computational Resources

• Training and evaluation runs took approximately 48 hours for DMC tasks and 72 hours for Robosuite tasks on one NVIDIA A5000 GPU.Source: §C.2 · p. 16

C.3 Hyperparameters

• Experiments reuse common settings across baselines where possible, with method-specific values reported separately; VIGOR otherwise inherits TD-MPC2 defaults.Tables 4 and 5 describe baseline and VIGOR hyperparameters, respectively.Source: §C.3 · p. 16, Table 4 · p. 16, Table 5 · p. 17

D Evaluation Details

• DMC evaluation covers seven distraction categories at easy and hard levels, while Robosuite uses easy, hard, and extreme levels.The settings shift appearance factors including color, lighting, background, and camera position.Source: D Evaluation Details · p. 17

D.1 Environments and Tasks

• The benchmarks comprise six DMC visuomotor tasks and four Robosuite manipulation tasks, with agents trained in clean environments and evaluated zero-shot under unseen distractions.Robosuite scenes show distraction severity increasing from eval-easy through eval-extreme.Source: §D.1 · p. 17, §D.5 · p. 19, Figure 10 · p. 20

D.2 Baseline Descriptions

• The baselines include methods selected for official implementations, visual-robustness relevance, and reproducible scores, alongside DrQ-v2, TD-MPC2, and DreamerV3 as reference methods.TD-MPC2 is VIGOR’s backbone, and Dr. G is the most directly comparable prior method; the set also spans augmentation-based MFRL and latent- or reconstruction-based MBRL approaches.Source: §D.2 · p. 18

D.3 Visual Distraction Categories

• Evaluation uses seven visual-distraction types, with easy and hard levels on DMC and eval-easy, eval-hard, and eval-extreme splits on Robosuite.Distractions include changes to agent or background color, background video, camera position, and lighting.Source: §D.3 · pp. 18–19, §D.4 · p. 19

E.1 Zero-Shot Generalization

• VIGOR ranks first across nearly all DMC distraction types and difficulty levels, and has the highest Robosuite returns at every difficulty level.The exception is DMC background-video hard, where SimGRL scores highest; baseline returns generally decline as distraction severity increases.Source: §E.1 · p. 21, Figure 12 · p. 21
• VIGOR’s average return is 777.6 on DMC, 3.4% above SRM (752.0), and 131.8 on Robosuite, 43.6% above SRM (91.8).It also records the strongest scores on all six DMC tasks and all four Robosuite tasks.Source: §E.1 · pp. 21–22, Figure 13 · p. 22, Table 6 · p. 22

E.2 Sample Efficiency

• VIGOR reaches DMC 25%, 50%, and 75% thresholds in 395, 412, and 422 episodes, respectively, and reaches every threshold in the fewest episodes among non-oracle methods on both benchmarks.On Robosuite, it takes 655, 720, and 740 episodes, versus SRM’s 1,240, 1,340, and 1,370; many baselines miss higher thresholds within budget.Source: Table 7 · p. 23, §E.2 · p. 23

E.3 DMC 100K

• With only 100K training steps, VIGOR’s DMC sample-efficiency score is approximately 0.78, close to TD-MPC2’s approximately 0.80, while VIGOR achieves the strongest zero-shot generalization.The results indicate that VIGOR’s robustness appears early in training rather than requiring extended convergence.Source: §E.3 · pp. 24–25, Figure 16 · p. 25

E.4 Wall-Clock Training Efficiency

• Over 24 hours, VIGOR’s average normalized-score gap from TD-MPC2 is approximately 0.17 and generally narrows as training progresses.The methods’ sample-efficiency curves are nearly identical, though TD-MPC2 trains slightly faster early on.Source: §E.4 · p. 25
• At hour 14, the IQM normalized-score difference (TD-MPC2 minus VIGOR) drops to 0.06; TD-MPC2 retains a small speed advantage, while VIGOR remains close as both converge.The per-task analysis covers all ten environments.Source: §E.4 · p. 26, Figure 18 · p. 26
• VIGOR’s modest wall-clock overhead diminishes over training, preserving practical training efficiency alongside substantially improved zero-shot robustness.Source: §E.4 · p. 26

• VIGOR achieves the highest returns and lowest variance on DMC at both easy and hard levels (863.9 ˘ 155.4 easy; 781.1 ˘ 261.2 hard).On Robosuite, AW maintains the highest returns at every difficulty level: 332.0 easy, 236.7 hard, and 146.0 extreme.Source: §E.5 · p. 27
• WO preserves training efficiency but under-generalizes on DMC, whereas Robosuite’s visual complexity demands stronger perturbation exposure.AW combines stability from weak-only samples with diversity from weak-to-strong samples.Source: §E.5 · p. 27
• The best mixture ratio differs in direction across benchmarks, but ζ=1 provides the best balance of training stability and generalization coverage.On DMC, more weak-only samples help; on Robosuite, more weak-to-strong samples help, while ζ=0.5 also converges more slowly early in training.Source: §E.5 · p. 28
• AW and WO converge at comparable rates, while SO fails to reach meaningful returns and WTSO converges to a substantially lower asymptote.SO also has wide confidence intervals throughout training.Source: §E.5 · p. 28
• Contextual random-overlay, textural random-conv, and spectral SRM all maintain strong zero-shot generalization, showing robustness across distinct augmentation families.The reported conclusion attributes robustness to latent-space consistency rather than the augmentation choice.Source: §E.5 · pp. 28–29

• Removing both DC and ES causes initial states to scatter and rollouts to diverge sharply, alongside substantial return drops and visible episode failures.Source: §E.6 · p. 29
• Without DC, ES partly clusters initial states, but the dynamics model still produces inconsistent multi-step predictions across difficulty levels.The residual discrepancies diverge through recursive rollouts.Source: §E.6 · p. 29
• Without ES, DC produces more compact rollouts, but perturbed initial states remain displaced from their clean counterparts, limiting DC’s benefit.Source: §E.6 · p. 29
• Only full VIGOR keeps initial encoder states and subsequent latent rollouts aligned across difficulty levels, with stable returns and successful task execution.Figure 20 visualizes this two-level consistency under background-color and object-color distractions.Source: §E.6 · p. 30, Figure 20 · p. 30

E.7 Ablation Study on Design Choices

• VIGOR_CONV performs competitively on color distractions but drops sharply on background-video; VIGOR_CONTRAST underperforms across most distraction types.VIGOR uses dynamic augmentation, resampling parameters independently at each rollout step, and random-overlay by default.Source: §E.7 · p. 31, Figure 22 · p. 32
• VIGOR achieves the best median, IQM, mean, and optimality gap, dominates the normalized-score distribution, and retains TD-MPC2’s sample efficiency.The ablation summary reports degradation when individual design choices are altered.Source: §E.7 · p. 32, Figure 23 · p. 32

E.8 Latent-Space Consistency Analysis

• Across background-color, background-video, and light-color distractions, VIGOR aligns both initial latent states and subsequent rollouts, unlike TD-MPC2 and Dr. G.The baselines’ encoder divergence is amplified through rollouts, with degraded returns and task failures.Source: §E.8 · p. 33, Figure 24 · p. 34

E.9 Performance under Longer Planning Horizons

• On finger-spin and door-opening, VIGOR maintains stable optimization and similar asymptotic returns for H=3, H=5, and H=10, though longer horizons slow convergence.The evaluation uses one representative task per benchmark because longer horizons increase wall-clock cost.Source: §E.9 · p. 35, Figure 25 · p. 35

E.10 Generalization to Unseen Camera Viewpoints

• All evaluated methods struggle with unseen camera rotations that change projection geometry, despite VIGOR’s robustness to appearance-based perturbations.The evaluation covers static and dynamic viewpoint rotations without viewpoint-rotation training augmentations.Source: §E.10 · pp. 35, 37, Figure 27 · p. 37

F MBRL Baseline Analysis

• Dreamer uses a decoder-based world model and learns an actor through RL, whereas TD-MPC uses a decoder-free model and selects actions through MPC guided by Q-values.Dreamer generates imagined rollouts for training; TD-MPC generates them during both training and test-time planning.Source: F MBRL Baseline Analysis · p. 38
• TD-MPC2 outperforms Dreamer-family methods and TD-MPC across diverse DMC tasks, motivating its selection as VIGOR’s sample-efficient backbone.The comparison is reported as consistent with prior findings [17].Source: F MBRL Baseline Analysis · p. 38, Figure 28 · p. 39

G.1 Visual Generalization in MFRL

• MFRL visual-generalization methods use data augmentation, representation learning, pretrained encoders, or optimization strategies to address robustness.The survey includes methods such as RAD, CURL, PIE-G, and CG2A.Source: §G.1 · pp. 39–40
• Because MFRL methods address single-step predictions, VIGOR pairs encoder stabilization with dynamics-level consistency for errors that compound through recursive MBRL rollouts.AW supplies the paired latent structure for both constraints.Source: §G.1 · p. 40

G.2 Extended Discussion of MBRL Visual Generalization

• Some invariant MBRL methods require test-time adaptation or target-domain environment interactions, limiting their applicability to zero-shot settings.The cited methods include DBC, Denoised MDP, RePo, HRSSM, and SD.Source: §G.2 · p. 40
• Unlike Dr. G, VIGOR uses asymmetric batch splitting and direct latent regression with encoder stabilization, without contrastive negatives or auxiliary supervision.AW provides a weak-only anchor and weak-to-strong latents from one encoder pass per sample; Dr. G generates two independent views and uses contrastive and inverse-dynamics losses.Source: §G.2 · p. 40
• VIGOR builds on TD-MPC2’s decoder-free, MPC-based architecture, whereas Dr. G uses Dreamer’s recurrent world model without a decoder.Dr. G relies on contrastive and inverse-dynamics objectives to learn latent representations.Source: §G.2 · p. 40
• VIGOR reports stronger zero-shot generalization with higher sample efficiency than Dr. G.The authors suggest that direct dynamics-level regression combined with encoder stabilization provides a more effective inductive bias than contrastive learning for this setting.Source: §G.2 · p. 40

H Discussion of Method Components

• The design addresses MBRL’s two-level vulnerability: distractions perturb encoder outputs, and those perturbations compound through recursive world-model rollouts.The shared encoder feeds the dynamics model, reward predictor, Q-value networks, and MPC planner.Source: H Discussion of Method Components · p. 41
• Asymmetric weak-to-strong augmentation pairs weak-only and weak-to-strong latents in one mini-batch, avoiding a second encoder pass while balancing insufficient robustness against instability under heavy perturbation.Weak-only ablations showed insufficient augmentation diversity; strong-only and weak-to-strong-only ablations showed optimization instability.Source: H Discussion of Method Components · p. 41
• Dynamics-level consistency regresses predictions from perturbed mixed latents toward weakly augmented successor targets, stabilizing forward dynamics for planning under visual shifts.The passage describes this anchoring as mitigating compounding rollout errors over long horizons.Source: H Discussion of Method Components · p. 41
• Encoder stabilization aligns weak-only and weak-to-strong encodings of the same frame to limit drift caused by dynamics consistency’s cross-augmentation supervision.It is intended to preserve task-relevant features under clean and perturbed inputs so initial latents remain grounded for long-horizon planning.Source: H Discussion of Method Components · p. 41
• Removing any of VIGOR’s three components significantly degrades performance, with the largest drops on Robosuite manipulation tasks under severe distribution shifts.The components are complementary: asymmetric batching provides paired latents without additional overhead, dynamics consistency anchors perturbed-input predictions to weakly augmented targets, and encoder stabilization limits drift under cross-augmentation pressure.Source: H Discussion of Method Components · p. 41