Learning Task Informed Abstractions

Anonymous

Keywords: learning under distraction, factored dynamics, model-based RL, representation learning

TL;DR: Accelerate model-based RL in complex visual domains, by factoring task-relevant and task-irrelevant features into two world models, where the secondary model is dissociated from the reward.

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 separates reward-correlated visual features from background distractions. For learning TIA, we introduce the formalism of Task Informed MDP (TiMDP), which 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.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2106.15612/code)