Learning Task Informed Abstractions
This addresses the challenge of visual distractions in reinforcement learning for control tasks, representing a novel method for a known bottleneck.
The paper tackles the problem of model-based reinforcement learning struggling with complex visual scenes by proposing Task Informed Abstractions (TIA) to separate reward-correlated features from distractors, resulting in significant performance gains over state-of-the-art methods on visual control tasks.
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.