[Paper Review] Multi-Task Reinforcement Learning with Context-based Representations

Shagun Sodhani, Amy Zhang+1 more|arXiv (Cornell University)|Feb 11, 2021

Reinforcement Learning in Robotics59 references20 citations

TL;DR

This paper proposes CARE (Contextual Attention-based Representation learning), a multi-task reinforcement learning framework that leverages task metadata—such as natural language descriptions—as context to learn composable, interpretable representations across a family of tasks. By conditioning representation selection on context, CARE reduces negative interference and achieves state-of-the-art performance on Meta-World, a challenging 50-task robotic manipulation benchmark.

ABSTRACT

The benefit of multi-task learning over single-task learning relies on the ability to use relations across tasks to improve performance on any single task. While sharing representations is an important mechanism to share information across tasks, its success depends on how well the structure underlying the tasks is captured. In some real-world situations, we have access to metadata, or additional information about a task, that may not provide any new insight in the context of a single task setup alone but inform relations across multiple tasks. While this metadata can be useful for improving multi-task learning performance, effectively incorporating it can be an additional challenge. We posit that an efficient approach to knowledge transfer is through the use of multiple context-dependent, composable representations shared across a family of tasks. In this framework, metadata can help to learn interpretable representations and provide the context to inform which representations to compose and how to compose them. We use the proposed approach to obtain state-of-the-art results in Meta-World, a challenging multi-task benchmark consisting of 50 distinct robotic manipulation tasks.

Motivation & Objective

• To address the challenge of negative interference in multi-task reinforcement learning by enabling fine-grained, context-dependent knowledge transfer.
• To leverage task metadata—such as natural language descriptions—as contextual signals to guide representation learning and policy composition.
• To develop a method that learns interpretable, task-specific representations that can be dynamically composed based on context for improved generalization.
• To improve sample efficiency and asymptotic performance in complex, real-world multi-task environments like Meta-World.

Proposed method

• The method introduces a Block Contextual MDP (BC-MDP) framework that extends standard MDPs by incorporating task-specific context, such as natural language descriptions.
• It employs a mixture of task- and object-specific encoders to learn multiple, disentangled representations of the state space.
• A context-aware attention mechanism selects and combines relevant representations based on the given task context, enabling dynamic policy composition.
• The policy network uses the selected representations to produce actions, allowing the agent to adapt its behavior per task context.
• The framework is trained end-to-end using off-policy deep RL algorithms, with context used to condition representation selection.
• The approach supports high-level, unstructured, or under-specified metadata, enabling flexibility in real-world application.

Experimental results

Research questions

• RQ1 Can task metadata be effectively used to guide the selection and composition of shared representations in multi-task reinforcement learning?
• RQ2 How can context-aware representation learning reduce negative interference when transferring knowledge across diverse tasks?
• RQ3 Can a mixture of interpretable, task-specific encoders improve sample efficiency and final performance in complex multi-task benchmarks?
• RQ4 To what extent can context-based representation selection generalize across tasks with varying state spaces and objectives?

Key findings

• CARE achieves state-of-the-art performance on the Meta-World benchmark, outperforming prior methods in both sample efficiency and final task success rates.
• The use of context—such as natural language descriptions—significantly improves generalization and reduces negative interference in multi-task learning.
• The learned representations are interpretable, with individual encoders specializing in distinct aspects like object types or skill components.
• The method generalizes across tasks with different state spaces by focusing only on relevant subspaces, enabling effective transfer even in partial observability settings.
• The ablation study confirms that context-aware representation selection is crucial for performance, with performance degrading when context is ignored or not used effectively.