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

Abstract

Model-based reinforcement learning (MBRL) achieves strong sample efficiency by planning in learned latent dynamics but degrades substantially under unseen visual distractions such as background, lighting, or camera changes. Unlike model-free RL, where encoder perturbations affect only single-step predictions, MBRL suffers from a two-level vulnerability: distractions push encoder outputs out of distribution, and these errors compound through recursive latent rollouts. We propose visual generalization via latent-space consistency in model-based RL (VIGOR), which enables zero-shot generalization to unseen distractions while retaining its backbone's sample efficiency. VIGOR comprises three interdependent components: (i) asymmetric weak-to-strong augmentation, which pairs weak-only and weak-to-strong latent views within a batch; (ii) dynamics-level consistency, which enforces augmentation-invariant transition predictions via direct latent regression; and (iii) encoder-level stabilization, which prevents encoder drift under the cross-augmentation supervision of (ii). On the DeepMind Control Suite (DMC) and Robosuite, VIGOR instantiated on TD-MPC2 outperforms state-of-the-art model-free and model-based baselines in zero-shot generalization, surpassing the second-best baseline by 3.4% on DMC and 43.6% on Robosuite, with sample efficiency comparable to its backbone and fewer environment interactions than the other baselines. Ablations show that VIGOR is augmentation-agnostic: augmentations from distinct perturbation families preserve strong generalization, confirming that latent-space consistency, not the augmentation choice, drives robustness. A second instantiation on DreamerV3, a structurally distinct recurrent state-space model, also improves zero-shot generalization on both benchmarks, indicating backbone-agnostic benefits. Code and videos: https://vigor-gen.github.io/.