Abstract
Learned dynamics models combined with both planning and policy learningalgorithms have shown promise in enabling artificial agents to learn to performmany diverse tasks with limited supervision. However, one of the fundamentalchallenges in using a learned forward dynamics model is the mismatch betweenthe objective of the learned model (future state reconstruction), and that ofthe downstream planner or policy (completing a specified task). This issue isexacerbated by vision-based control tasks in diverse real-world environments,where the complexity of the real world dwarfs model capacity. In this paper, wepropose to direct prediction towards task relevant information, enabling themodel to be aware of the current task and encouraging it to only model relevantquantities of the state space, resulting in a learning objective that moreclosely matches the downstream task. Further, we do so in an entirelyself-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 sceneconditioned on the goal, and as a result outperforms standard task-agnosticdynamics models and model-free reinforcement learning.