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

Latent-space MBRL plans inside a learned world model. An encoder \(h_\theta\) maps pixels to a latent state, and a dynamics model \(d_\theta\) rolls that state forward to evaluate future actions. This works well as long as test observations look like training observations. The video below walks through what happens when they do not.

In model-free RL, a visual shift perturbs only a single-step prediction. In model-based RL the same shift causes two failures in sequence:

1. Encoder-level drift. Visual distractions yield out-of-distribution latents \(z'_t \notin \mathcal{Z}^{\text{train}}\), even though the physical state is unchanged.
2. Dynamics-level error accumulation. The dynamics model propagates these latents recursively over the planning horizon \(H\), so errors compound across steps and the rollouts used for planning diverge.

Techniques designed for single-step robustness are therefore insufficient. Our goal is zero-shot generalization: no test-time data and no adaptation, and only appearance-based distractions change between training and deployment.

If the failure happens at two levels, the fix has to act at two levels. The world model should predict the same transitions regardless of the appearance-based perturbation applied to its input, and the shared encoder should remain stable while this constraint is imposed.

Simply transferring model-free strategies does not work: applying strong augmentation uniformly across the batch, without AW's asymmetric split, reaches only 48% of VIGOR's performance on DMC.

AW splits each mini-batch evenly at random. One half passes through weak augmentation \(\tau^w\) (random-shift), the other through \(\tau^s \circ \tau^w\) (random-overlay after random-shift). The two branches are concatenated and encoded in a single forward pass, which preserves the batch size and provides the pairs that DC operates on.

\[ s^m_t = s^w_t \oplus s^{ws}_t, \qquad z^m_t = h_\theta(s^m_t) \]

DC re-targets the standard dynamics loss. The mixed latent is injected at the first step of each trajectory segment, and every rollout prediction \(\hat{z}^m_{t+h} = d_\theta(\hat{z}^m_{t+h-1}, a_{t+h-1})\) is aligned with a weakly augmented target, propagating the initial perturbation through the whole horizon while requiring augmentation-invariant transitions.

\[ \mathcal{L}_{\text{DC}}(\theta) = \big\| d_\theta(\hat{z}^m_t, a_t) - \text{sg}(z^w_{t+1}) \big\|_2^2, \quad \text{where } z^w_{t+1} = h_\theta\big(\tau^w(s_{t+1}, \upsilon^w)\big) \]

A single shared encoder feeds the dynamics model, Q-value networks, and the MPC planner, and DC's cross-augmentation supervision can distort it. ES adapts the alignment objective of SODA to counteract this drift, anchoring weak-to-strong encodings to their weak-only counterparts.

\[ \mathcal{L}_{\text{ES}}(\theta) = \left\| \frac{z^{ws}_t}{\|z^{ws}_t\|_2} - \frac{\text{sg}(z^w_t)}{\|\text{sg}(z^w_t)\|_2} \right\|_2^2 \]

AW modifies only the input pipeline, DC only re-targets the existing dynamics loss, and ES adds a single auxiliary loss, so the planner and optimizer stay unchanged. We instantiate VIGOR on two structurally different backbones.

Both VIGOR instantiations exceed every baseline on every metric, on both benchmarks. The default backbone TD-MPC2 ranks lowest among all baselines; VIGOR lifts it to the top of the table.

Robustness comes at negligible cost. VIGOR's learning curves closely track those of the unmodified backbone on both TD-MPC2 and DreamerV3. On Robosuite, TD-MPC2 and TD-MPC2+VIGOR converge at comparable rates, while the other baselines fall behind in both speed and asymptotic returns.

VIGOR is consistent at both levels of the vulnerability. Rolling out each world model from identical inputs shows that VIGOR's initial encodings stay close to their clean counterparts, and this alignment is preserved through the rollout horizon. TD-MPC2 and Dr. G produce scattered initial encodings or divergent rollouts, with degraded returns.

Each component fails differently when removed. Without DC, performance on Robosuite falls to 46%: encoder-level alignment alone cannot stop the dynamics model from amplifying residual discrepancies. Among augmentation strategies, only AW performs well on both benchmarks, because the asymmetric split provides both optimization stability and perturbation diversity.

VIGOR targets appearance-based perturbations. Geometric viewpoint changes, dynamics shift, and task-level shift are outside its scope: under unseen camera rotations, all evaluated methods, including VIGOR, degrade sharply. Our backbone-generality evidence is limited to two world-model families, and validation on real robots remains future work.