# Goal-Aware Prediction: Learning to Model What Matters

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

## Key facts

- Cited by 45 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 14 Jul 2020.

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

- [R3M: A Universal Visual Representation for Robot Manipulation](https://mlsanity.com/paper/2203.12601/) (2022)
- [Goal-Conditioned Reinforcement Learning: Problems and Solutions](https://mlsanity.com/paper/2201.08299/) (2022)
- [Physion: Evaluating Physical Prediction from Vision in Humans and Machines](https://mlsanity.com/paper/2106.08261/) (2021)
- [Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models](https://mlsanity.com/paper/2502.14819/) (2025)
- [DREAMWALKER: Mental Planning for Continuous Vision-Language Navigation](https://mlsanity.com/paper/2308.07498/) (2023)
- [Greedy Hierarchical Variational Autoencoders for Large-Scale Video Prediction](https://mlsanity.com/paper/2103.04174/) (2021)
- [Language Grounding with 3D Objects](https://mlsanity.com/paper/2107.12514/) (2021)
- [Contrastive Difference Predictive Coding](https://mlsanity.com/paper/2310.20141/) (2023)
- [Generalizing Goal-Conditioned Reinforcement Learning with Variational Causal Reasoning](https://mlsanity.com/paper/2207.09081/) (2022)
- [Value Gradient weighted Model-Based Reinforcement Learning](https://mlsanity.com/paper/2204.01464/) (2022)
- [Deciding What to Model: Value-Equivalent Sampling for Reinforcement Learning](https://mlsanity.com/paper/2206.02072/) (2022)
- [VisuoSpatial Foresight for Physical Sequential Fabric Manipulation](https://mlsanity.com/paper/2102.09754/) (2021)
- [Control-Aware Prediction Objectives for Autonomous Driving](https://mlsanity.com/paper/2204.13319/) (2022)
- [Plan To Predict: Learning an Uncertainty-Foreseeing Model for Model-Based Reinforcement Learning](https://mlsanity.com/paper/2301.08502/) (2023)
- [Learning to Discern: Imitating Heterogeneous Human Demonstrations with Preference and Representation Learning](https://mlsanity.com/paper/2310.14196/) (2023)
- [Solving Minimum-Cost Reach Avoid using Reinforcement Learning](https://mlsanity.com/paper/2410.22600/) (2024)
- [Efficient Robotic Policy Learning via Latent Space Backward Planning](https://mlsanity.com/paper/2505.06861/) (2025)
- [MapGo: Model-Assisted Policy Optimization for Goal-Oriented Tasks](https://mlsanity.com/paper/2105.06350/) (2021)
- [Learning Domain Invariant Representations in Goal-conditioned Block MDPs](https://mlsanity.com/paper/2110.14248/) (2021)
- [LS3: Latent Space Safe Sets for Long-Horizon Visuomotor Control of Sparse Reward Iterative Tasks](https://mlsanity.com/paper/2107.04775/) (2021)
- [DisCo RL: Distribution-Conditioned Reinforcement Learning for General-Purpose Policies](https://mlsanity.com/paper/2104.11707/) (2021)
- [Pre-trained Visual Dynamics Representations for Efficient Policy Learning](https://mlsanity.com/paper/2411.03169/) (2024)
- [Beyond Next-Observation Prediction: Agent-Authored World Modeling for Sequential Decision Making](https://mlsanity.com/paper/2606.25421/) (2026)
- [Goal-Aware Cross-Entropy for Multi-Target Reinforcement Learning](https://mlsanity.com/paper/2110.12985/) (2021)
- [On Rate-Distortion Theory in Capacity-Limited Cognition & Reinforcement Learning](https://mlsanity.com/paper/2210.16877/) (2022)
- [Learning World Models for Unconstrained Goal Navigation](https://mlsanity.com/paper/2411.02446/) (2024)
- [Policy-shaped prediction: avoiding distractions in model-based reinforcement learning](https://mlsanity.com/paper/2412.05766/) (2024)
- [Tell Me What's Next: Textual Foresight for Generic UI Representations](https://mlsanity.com/paper/2406.07822/) (2024)
- [Make the Pertinent Salient: Task-Relevant Reconstruction for Visual Control with Distractions](https://mlsanity.com/paper/2410.09972/) (2024)
- [Learning Intuitive Physics with Multimodal Generative Models](https://mlsanity.com/paper/2101.04454/) (2021)

## Abstract (arXiv)

Learned dynamics models combined with both planning and policy learning algorithms have shown promise in enabling artificial agents to learn to perform many diverse tasks with limited supervision. However, one of the fundamental challenges in using a learned forward dynamics model is the mismatch between the objective of the learned model (future state reconstruction), and that of the downstream planner or policy (completing a specified task). This issue is exacerbated by vision-based control tasks in diverse real-world environments, where the complexity of the real world dwarfs model capacity. In this paper, we propose to direct prediction towards task relevant information, enabling the model to be aware of the current task and encouraging it to only model relevant quantities of the state space, resulting in a learning objective that more closely matches the downstream task. Further, we do so in an entirely self-supervised manner, without the need for a reward function or image labels. We find that our method more effectively models the relevant parts of the scene conditioned on the goal, and as a result outperforms standard task-agnostic dynamics models and model-free reinforcement learning.