Distributionally Robust Multi-Task Reinforcement Learning via Adaptive Task Sampling

TL;DR

This paper introduces DRATS, a new adaptive task sampling method for multi-task reinforcement learning that improves data efficiency and worst-task performance by focusing on tasks that are furthest from being solved.

Contribution

The paper proposes a novel distributionally robust adaptive task sampling algorithm that addresses data imbalance in multi-task reinforcement learning.

Findings

DRATS outperforms existing sampling algorithms on MetaWorld benchmarks.

DRATS improves data efficiency in multi-task reinforcement learning.

DRATS increases worst-task performance in benchmark tests.

Abstract

Multi-task reinforcement learning (MTRL) aims to train a single agent to efficiently optimize performance across multiple tasks simultaneously. However, jointly optimizing all tasks often yields imbalanced learning: agents quickly solve easy tasks but learn slowly on harder ones. While prior work primarily attributes this imbalance to conflicting task gradients and proposes gradient manipulation or specialized architectures to address it, we instead focus on a distinct and under-explored challenge: imbalanced data allocation. Standard MTRL allocates an equal number of environment interactions to each task, which over-allocates data to easy tasks that require relatively few interactions to solve and under-allocates data to hard tasks that require substantially more experience to solve. To address this challenge, we introduce Distributionally Robust Adaptive Task Sampling (DRATS), an algorithm that adaptively prioritizes sampling tasks furthest from being solved. We derive DRATS by formalizing MTRL as a feasibility problem from which we derive a minimax objective for minimizing the worst-case return gap, the difference between a desired target return and the agent's return on a task. In benchmarks like MetaWorld-MT10 and MT50, DRATS improves data efficiency and increases worst-task performance compared to existing task sampling algorithms.