# Focus On What Matters: Separated Models For Visual-Based RL Generalization

arXiv: https://arxiv.org/abs/2410.10834 · HTML: https://mlsanity.com/paper/2410.10834/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in RL for robot locomotion and manipulation, one of 15 directions in Deep reinforcement learning methods.
- Builds on 42 of its 47 references that are in the map.
- First posted on arXiv 29 Sep 2024.

## Placed in

- Deep reinforcement learning methods / [RL for robot locomotion and manipulation](https://mlsanity.com/deeprl/rl-for-robot-locomotion-and-manipulation/)

## Builds on (most cited first, 30 of 42)

- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [End-to-End Training of Deep Visuomotor Policies](https://mlsanity.com/paper/1504.00702/) (2015)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [On Variational Bounds of Mutual Information](https://mlsanity.com/paper/1905.06922/) (2019)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Reinforcement Learning with Augmented Data](https://mlsanity.com/paper/2004.14990/) (2020)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Target-driven Visual Navigation in Indoor Scenes using Deep Reinforcement Learning](https://mlsanity.com/paper/1609.05143/) (2016)
- [Quantifying Generalization in Reinforcement Learning](https://mlsanity.com/paper/1812.02341/) (2018)
- [Improving Sample Efficiency in Model-Free Reinforcement Learning from Images](https://mlsanity.com/paper/1910.01741/) (2019)
- [Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Model](https://mlsanity.com/paper/1907.00953/) (2019)
- [Learning to Navigate in Complex Environments](https://mlsanity.com/paper/1611.03673/) (2016)
- [A Study on Overfitting in Deep Reinforcement Learning](https://mlsanity.com/paper/1804.06893/) (2018)
- [Sim-to-Real via Sim-to-Sim: Data-efficient Robotic Grasping via Randomized-to-Canonical Adaptation Networks](https://mlsanity.com/paper/1812.07252/) (2018)
- [Variational Information Maximisation for Intrinsically Motivated Reinforcement Learning](https://mlsanity.com/paper/1509.08731/) (2015)
- [Generalization in Reinforcement Learning by Soft Data Augmentation](https://mlsanity.com/paper/2011.13389/) (2020)
- [Learning Agile Soccer Skills for a Bipedal Robot with Deep Reinforcement Learning](https://mlsanity.com/paper/2304.13653/) (2023)
- [Self-Supervised Policy Adaptation during Deployment](https://mlsanity.com/paper/2007.04309/) (2020)
- [Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2101.05265/) (2021)
- [Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation](https://mlsanity.com/paper/2107.00644/) (2021)
- [Network Randomization: A Simple Technique for Generalization in Deep Reinforcement Learning](https://mlsanity.com/paper/1910.05396/) (2019)
- [Observational Overfitting in Reinforcement Learning](https://mlsanity.com/paper/1912.02975/) (2019)
- [Improving Generalization in Reinforcement Learning with Mixture Regularization](https://mlsanity.com/paper/2010.10814/) (2020)
- [Denoised MDPs: Learning World Models Better Than the World Itself](https://mlsanity.com/paper/2206.15477/) (2022)
- [Learning Task Informed Abstractions](https://mlsanity.com/paper/2106.15612/) (2021)

## Built on by (most cited first, 6 of 6)

- [DPMambaIR: All-in-One Image Restoration via Degradation-Aware Prompt State Space Model](https://mlsanity.com/paper/2504.17732/) (2025)
- [Dreamer-CDP: Improving Reconstruction-free World Models Via Continuous Deterministic Representation Prediction](https://mlsanity.com/paper/2603.07083/) (2026)
- [Rejecting Hallucinated State Targets during Planning](https://mlsanity.com/paper/2410.07096/) (2024)
- [SegDAC: Visual Generalization in Reinforcement Learning via Dynamic Object Tokens](https://mlsanity.com/paper/2508.09325/) (2025)
- [StarWM: Self-Supervised Trained Attention Routing for Robust World Models](https://mlsanity.com/paper/2609.30667/) (2026)
- [VIGOR: Zero-Shot Visual Generalization via Latent-Space Consistency in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2610.02801/) (2026)

## Abstract (arXiv)

A primary challenge for visual-based Reinforcement Learning (RL) is to generalize effectively across unseen environments. Although previous studies have explored different auxiliary tasks to enhance generalization, few adopt image reconstruction due to concerns about exacerbating overfitting to task-irrelevant features during training. Perceiving the pre-eminence of image reconstruction in representation learning, we propose SMG (Separated Models for Generalization), a novel approach that exploits image reconstruction for generalization. SMG introduces two model branches to extract task-relevant and task-irrelevant representations separately from visual observations via cooperatively reconstruction. Built upon this architecture, we further emphasize the importance of task-relevant features for generalization. Specifically, SMG incorporates two additional consistency losses to guide the agent's focus toward task-relevant areas across different scenarios, thereby achieving free from overfitting. Extensive experiments in DMC demonstrate the SOTA performance of SMG in generalization, particularly excelling in video-background settings. Evaluations on robotic manipulation tasks further confirm the robustness of SMG in real-world applications.