Learning Invariant Visual Representations for Planning with Joint-Embedding Predictive World Models

Deep learning-driven visual perception has transformed autonomous control, with models like DINO-WM and other JEPAs enabling predictive planning in high-dimensional environments. These models leverage pretrained visual encoders to extract features, then predict future states without pixel reconstruction. Despite successes, they suffer from sensitivity to environmental variations, especially background changes and distractors, which introduce slow features that dominate the learned representations. This leads to poor generalization and brittle planning in real-world scenarios. The challenge is to develop methods that can learn control-relevant, invariant features from pretrained encoders, ensuring robustness across diverse visual conditions while maintaining computational efficiency.

Existing JEPAs, including DINO-WM, tend to encode slow, task-irrelevant visual features such as backgrounds and distractors, which vary gradually over time. This causes the predictive models to focus on these nuisance features, degrading robustness when environmental conditions change at test time. The core problem is how to learn latent representations that are invariant to such slow features, preserving only the control-relevant dynamics necessary for planning. Achieving this without sacrificing the predictive power and while maintaining computational efficiency remains a significant challenge, especially given the high dimensionality and complexity of pretrained features.

The paper introduces a bisimulation encoder that enforces control-relevant invariance by leveraging the bisimulation metric, which measures behavioral equivalence based on reward and transition similarities. This encoder is trained jointly with a latent predictive model, using Wasserstein distance to approximate transition similarity, and incorporates PCA-based regularization to prevent feature collapse. Unlike prior work, it does not rely on reward prediction, simplifying training. The approach effectively filters out slow, irrelevant features from fixed pretrained visual encoders, enabling a compressed latent space that retains essential dynamics while being robust to environmental variations. This combination of bisimulation, PCA regularization, and pretrained features constitutes a novel contribution to invariant representation learning.

• �� Extract visual features using pretrained encoders (DINOv2, SimDINOv2, iBOT).
• �� Design a bisimulation encoder hη that maps high-dimensional features to a low-dimensional, control-relevant latent space.
• �� Train the encoder jointly with a latent transition model Tϕ, minimizing a bisimulation loss based on Wasserstein distance between predicted and actual transition behaviors.
• �� Incorporate PCA-based variance regularization to prevent feature collapse, ensuring diversity in the latent space.
• �� Use a contrastive predictive loss to train the transition model, with the bisimulation encoder enforcing invariance.
• �� Optimize planning via CEM in the learned invariant latent space, testing robustness under environmental shifts.

The experimental setup involves a PointMaze navigation task with varying background conditions and distractors. The models are trained on fixed visual environments and tested under diverse background changes. Baselines include DINO-WM and variants with different pretrained encoders. Key metrics include success rate, latent space size, and robustness to visual shifts. Hyperparameters involve the regularization weights for bisimulation and PCA regularization. Ablation studies analyze the impact of each component, confirming that the bisimulation loss and PCA regularization are critical for invariance. Results show the proposed model outperforms baselines in success rate and robustness, with a tenfold reduction in latent dimensionality.

The proposed approach achieves over 80% success rate under severe background shifts, outperforming DINO-WM by 20%. The latent space is compressed to 1/10 of the original size, maintaining predictive accuracy. Across different pretrained encoders, the model consistently demonstrates enhanced robustness against slow features and distractors. Ablation results highlight the importance of the bisimulation loss and PCA regularization, with the combined method providing the best invariance and compactness. These findings validate the effectiveness of the invariant representation learning framework in real-world-like environments.

This method is directly applicable to autonomous navigation, robotic manipulation, and visual control in unpredictable environments. It requires only pretrained visual encoders and minimal task-specific tuning, making it suitable for resource-constrained robotic systems. The invariant representations enable more reliable planning in dynamic, cluttered settings, facilitating deployment in real-world scenarios such as warehouse robots, autonomous vehicles, and service robots, where environmental variability is high.

The approach relies heavily on the quality of pretrained features; if slow features dominate, invariance may be limited. Wasserstein distance approximation introduces computational overhead, potentially affecting real-time applications. The method's performance in highly dynamic or complex environments with non-stationary visual features remains to be fully validated. Future work should focus on scalability, adaptive regularization, and multi-modal integration to address these challenges.