Learning Task Informed Abstractions

TL;DR

This paper introduces Task Informed Abstractions (TIA) for reinforcement learning, which explicitly separate task-relevant features from distractors in complex visual scenes, improving performance.

Contribution

The paper proposes the formalism of Task Informed MDP (TiMDP) and a method to learn visual features through cooperative reconstruction and adversarial dissociation from rewards.

Findings

TIA significantly outperforms state-of-the-art methods on visual control tasks.

Learning task-informed features improves robustness to visual distractions.

Empirical results demonstrate the effectiveness of the proposed approach.

Abstract

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.