# Learning Task Informed Abstractions

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

## Key facts

- Cited by 58 papers in the map (data through 14 Oct 2026).
- Builds on 55 of its 67 references that are in the map.
- First posted on arXiv 29 Jun 2021.

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

- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Unsupervised Domain Adaptation by Backpropagation](https://mlsanity.com/paper/1409.7495/) (2014)
- [The Kinetics Human Action Video Dataset](https://mlsanity.com/paper/1705.06950/) (2017)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Curiosity-driven Exploration by Self-supervised Prediction](https://mlsanity.com/paper/1705.05363/) (2017)
- [Learning Robust Global Representations by Penalizing Local Predictive Power](https://mlsanity.com/paper/1905.13549/) (2019)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Semi-Supervised Learning with Deep Generative Models](https://mlsanity.com/paper/1406.5298/) (2014)
- [Reinforcement Learning with Augmented Data](https://mlsanity.com/paper/2004.14990/) (2020)
- [Time-Contrastive Networks: Self-Supervised Learning from Video](https://mlsanity.com/paper/1704.06888/) (2017)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Deep Visual Foresight for Planning Robot Motion](https://mlsanity.com/paper/1610.00696/) (2016)
- [Don't Take the Easy Way Out: Ensemble Based Methods for Avoiding Known Dataset Biases](https://mlsanity.com/paper/1909.03683/) (2019)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [Improving Sample Efficiency in Model-Free Reinforcement Learning from Images](https://mlsanity.com/paper/1910.01741/) (2019)
- [Decoupling Representation Learning from Reinforcement Learning](https://mlsanity.com/paper/2009.08319/) (2020)
- [Variational Intrinsic Control](https://mlsanity.com/paper/1611.07507/) (2016)
- [DeepMDP: Learning Continuous Latent Space Models for Representation Learning](https://mlsanity.com/paper/1906.02736/) (2019)
- [Learning De-biased Representations with Biased Representations](https://mlsanity.com/paper/1910.02806/) (2019)
- [RUBi: Reducing Unimodal Biases in Visual Question Answering](https://mlsanity.com/paper/1906.10169/) (2019)

## Built on by (most cited first, 30 of 58)

- [Optimal Goal-Reaching Reinforcement Learning via Quasimetric Learning](https://mlsanity.com/paper/2304.01203/) (2023)
- [Denoised MDPs: Learning World Models Better Than the World Itself](https://mlsanity.com/paper/2206.15477/) (2022)
- [Causal Dynamics Learning for Task-Independent State Abstraction](https://mlsanity.com/paper/2206.13452/) (2022)
- [Bridging State and History Representations: Understanding Self-Predictive RL](https://mlsanity.com/paper/2401.08898/) (2024)
- [Contextualize Me -- The Case for Context in Reinforcement Learning](https://mlsanity.com/paper/2202.04500/) (2022)
- [Guaranteed Discovery of Control-Endogenous Latent States with Multi-Step Inverse Models](https://mlsanity.com/paper/2207.08229/) (2022)
- [Latent Action Learning Requires Supervision in the Presence of Distractors](https://mlsanity.com/paper/2502.00379/) (2025)
- [Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learning](https://mlsanity.com/paper/2209.09203/) (2022)
- [Inverse Dynamics Pretraining Learns Good Representations for Multitask Imitation](https://mlsanity.com/paper/2305.16985/) (2023)
- [On the Importance of Exploration for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2306.05483/) (2023)
- [Facing Off World Model Backbones: RNNs, Transformers, and S4](https://mlsanity.com/paper/2307.02064/) (2023)
- [INFOrmation Prioritization through EmPOWERment in Visual Model-Based RL](https://mlsanity.com/paper/2204.08585/) (2022)
- [The Generalization Gap in Offline Reinforcement Learning](https://mlsanity.com/paper/2312.05742/) (2023)
- [Learning Representations for Pixel-based Control: What Matters and Why?](https://mlsanity.com/paper/2111.07775/) (2021)
- [Learning World Models with Identifiable Factorization](https://mlsanity.com/paper/2306.06561/) (2023)
- [CARL: A Benchmark for Contextual and Adaptive Reinforcement Learning](https://mlsanity.com/paper/2110.02102/) (2021)
- [Learning Latent Dynamic Robust Representations for World Models](https://mlsanity.com/paper/2405.06263/) (2024)
- [RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability](https://mlsanity.com/paper/2309.00082/) (2023)
- [Structure in Deep Reinforcement Learning: A Survey and Open Problems](https://mlsanity.com/paper/2306.16021/) (2023)
- [Robust Representation Learning by Clustering with Bisimulation Metrics for Visual Reinforcement Learning with Distractions](https://mlsanity.com/paper/2302.12003/) (2023)
- [Generalizable Visual Reinforcement Learning with Segment Anything Model](https://mlsanity.com/paper/2312.17116/) (2023)
- [Task-Induced Representation Learning](https://mlsanity.com/paper/2204.11827/) (2022)
- [Building Minimal and Reusable Causal State Abstractions for Reinforcement Learning](https://mlsanity.com/paper/2401.12497/) (2024)
- [Ignorance is Bliss: Robust Control via Information Gating](https://mlsanity.com/paper/2303.06121/) (2023)
- [SeMAIL: Eliminating Distractors in Visual Imitation via Separated Models](https://mlsanity.com/paper/2306.10695/) (2023)
- [AD3: Implicit Action is the Key for World Models to Distinguish the Diverse Visual Distractors](https://mlsanity.com/paper/2403.09976/) (2024)
- [Combining Reconstruction and Contrastive Methods for Multimodal Representations in RL](https://mlsanity.com/paper/2302.05342/) (2023)
- [Focus On What Matters: Separated Models For Visual-Based RL Generalization](https://mlsanity.com/paper/2410.10834/) (2024)
- [Policy-shaped prediction: avoiding distractions in model-based reinforcement learning](https://mlsanity.com/paper/2412.05766/) (2024)
- [Reimagination with Test-time Observation Interventions: Distractor-Robust World Model Predictions for Visual Model Predictive Control](https://mlsanity.com/paper/2506.16565/) (2025)

## Abstract (arXiv)

Current model-based reinforcement learning methods struggle when operating from complex visual scenes due to their inability to prioritize task-relevant features. To mitigate this problem, we propose learning Task Informed Abstractions (TIA) that explicitly separates reward-correlated visual features from distractors. For learning TIA, we introduce the formalism of Task Informed MDP (TiMDP) that is realized by training two models that learn visual features via cooperative reconstruction, but one model is adversarially dissociated from the reward signal. Empirical evaluation shows that TIA leads to significant performance gains over state-of-the-art methods on many visual control tasks where natural and unconstrained visual distractions pose a formidable challenge.