Technical Summary: Distributionally Robust Adaptive Task Sampling (DRATS)

Problem Statement

Multi-task Reinforcement Learning (MTRL) aims to train a single agent to optimize performance across multiple tasks simultaneously, leveraging knowledge transfer to improve data efficiency. However, standard MTRL approaches often suffer from imbalanced learning: agents rapidly solve easy tasks while making slow progress on harder ones.

Prior work has largely attributed this imbalance to conflicting task gradients, proposing solutions involving gradient manipulation (e.g., gradient surgery) or specialized network architectures. This paper identifies a distinct, under-explored cause of imbalance: imbalanced data allocation. Standard MTRL allocates an equal interaction budget to all tasks. Since hard tasks require significantly more experience to solve than easy ones, uniform allocation results in over-allocating data to easy tasks (which need little more data) and under-allocating data to hard tasks. While Curriculum Learning (CL) methods adapt sampling distributions, they typically prioritize easy tasks first to facilitate transfer, which can exacerbate data imbalance in settings where positive transfer is weak or tasks conflict.

Methodology: Distributionally Robust Adaptive Task Sampling (DRATS)

The authors propose DRATS, an algorithm that adaptively prioritizes sampling tasks that are furthest from being solved. The method is derived through the following steps:

1. Formalization as a Feasibility Problem

Instead of maximizing the average return across tasks, MTRL is formalized as a feasibility problem: finding a policy $\theta$ such that the return $J_i(\theta)$ in every task $i$ meets a reference target $J^{ref}_i$.
The objective is transformed into minimizing the return gap $g_i(\theta) = \max(J^{ref}_i - J_i(\theta), 0)$. To handle cases where constraints cannot be simultaneously satisfied (due to capacity limits or conflicts), the goal becomes minimizing the maximum return gap across all tasks (a minimax objective):
$$\min_{\theta} \max_{i \in [k]} g_i(\theta)$$

2. Minimax Optimization and Distributional Robustness

The inner maximization is recast as a maximization over a task distribution $q \in \Delta_k$ (the probability simplex). Solving this directly leads to a non-smooth, greedy solution that concentrates all probability on the single worst-performing task, causing instability and data inefficiency.

To address this, DRATS employs Distributionally Robust Optimization (DRO) by restricting the task distribution $q$ to lie within a Kullback-Leibler (KL) ball of radius $\epsilon$ centered at a base distribution $p_0$ (uniform). This yields the constrained optimization:
$$\min_{\theta} \max_{q \in \mathcal{Q}} \mathbb{E}_{i \sim q}[g_i(\theta)]$$
where $\mathcal{Q} = \{q \in \Delta_k \mid \text{KL}(q \| p_0) \leq \epsilon\}$.

3. Closed-Form Solution and Algorithm

Using Lagrangian relaxation, the optimal task distribution $q^*$ is derived in closed form as a softmax of the return gaps:
$$q^*_i = \frac{\exp(\eta g_i(\theta))}{\sum_j \exp(\eta g_j(\theta))}$$
Here, $\eta$ acts as an inverse temperature parameter controlling the sharpness of the prioritization.

DRATS Algorithm (Algorithm 1):

1. Data Collection: Sample tasks from the current distribution $q_t$ and collect trajectories.
2. Gap Estimation: Compute Monte Carlo estimates of returns $\hat{J}_i$ and update reference returns $J^{ref}_i$ (estimated online as the maximum observed return).
3. Gap Normalization: Compute normalized return gaps $\hat{g}_i$.
4. Distribution Update: Update $q_t$ via mirror ascent (exponentiated gradient) on the objective $L(\theta, q) = \mathbb{E}_{i \sim q}[g_i(\theta)] - \frac{1}{\eta}\text{KL}(q \| p_0)$.
5. Projection: Enforce a minimum sampling probability $\epsilon$ to ensure all tasks are sampled sufficiently for stable gap estimation.
6. Policy Update: Update the policy $\theta$ using an on-policy algorithm (e.g., PPO) to maximize the weighted return.

4. Convergence Guarantees

The authors prove that under standard assumptions (including a low-regret guarantee for the policy update), DRATS converges to the optimal worst-case return gap at a rate of $O(1/\sqrt{T})$, where $T$ is the number of iterations.

Key Contributions

1. DRATS Algorithm: Introduction of a multi-task RL algorithm that adaptively prioritizes tasks based on the return gap, derived from a minimax formulation of MTRL.
2. Theoretical Convergence: Proof that DRATS converges to within $\epsilon$ error of the optimal worst-task return gap at a standard $O(1/\sqrt{T})$ rate.
3. Empirical Validation: Demonstration that DRATS achieves higher data efficiency, higher aggregate return, and significantly higher worst-task return compared to existing curriculum learning baselines (e.g., Uniform, Learning Progress, Learning Potential, Hard First).
4. Orthogonality: Demonstration that DRATS is compatible with existing multi-task network architectures (e.g., Soft Modularization, MOORE) and improves performance regardless of whether tasks exhibit positive transfer or conflict.

Experimental Results

The authors evaluated DRATS on four benchmarks: a 4-task Gridworld, MuJoCo6 (6 continuous control tasks), MetaWorld-MT10, and MetaWorld-MT50.

• Aggregate Performance: DRATS achieved the highest aggregate return across all benchmarks.
• Worst-Task Performance: The improvement was driven primarily by superior performance on the hardest tasks. For example, in MuJoCo6, DRATS was the only method to successfully solve the difficult Ant and Humanoid tasks, whereas baselines failed to make progress on Ant.
• Data Efficiency: DRATS reached the performance levels of baselines with significantly fewer interactions. In MT10, it surpassed the final return of the Uniform baseline on hard tasks (e.g., pick-place, push) after just 6 million timesteps.
• Comparison to Baselines:
  • Uniform: Over-allocates to easy tasks; fails to solve hard tasks.
  • Learning Progress/Potential: Tend to prioritize tasks with high learning rates or high value estimation error, which often correlates with easy tasks, neglecting hard tasks that are stuck.
  • Hard First (SMT): Suffers from performance collapse on hard tasks; once a task is declared "solved" based on a threshold, sampling probability drops abruptly, causing the agent to forget the task and requiring relearning.
• Ablations:
  • Adaptive Sampling vs. Reweighting: Adaptive sampling (DRATS) is significantly more data-efficient than reweighting the objective of uniformly sampled data. Reweighting leads to high-variance gradient estimates for hard tasks due to insufficient data collection.
  • Hyperparameters: DRATS is robust to the choice of temperature $\eta$ and step size $\alpha$.

Significance and Claims

The paper claims that DRATS addresses a fundamental inefficiency in MTRL: the misallocation of data resources. By shifting the focus from gradient conflict to data allocation, DRATS provides a data-efficient mechanism to balance performance across tasks.

The authors emphasize that:

• DRATS is the first application of Distributionally Robust Optimization to multi-task RL in continuous control settings.
• Unlike many Curriculum Learning methods that assume positive transfer from easy to hard tasks, DRATS explicitly targets the worst-case return gap, making it effective even when tasks conflict or transfer is weak.
• The method is modular; it modifies only the data collection process and can be combined with specialized architectures or gradient manipulation techniques.

The paper concludes that DRATS offers a principled, theoretically grounded approach to achieving balanced multi-task learning, improving both the efficiency of learning and the robustness of the final policy across all tasks.