# Distributionally Robust Multi-Task Reinforcement Learning via Adaptive Task Sampling

arXiv: https://arxiv.org/abs/2605.14350 · HTML: https://mlsanity.com/paper/2605.14350/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 50 of its 63 references that are in the map.
- First posted on arXiv 14 May 2026.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 30 of 50)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [Conflict-Averse Gradient Descent for Multi-task Learning](https://mlsanity.com/paper/2110.14048/) (2021)
- [Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning](https://mlsanity.com/paper/2310.06694/) (2023)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [Curriculum Learning for Reinforcement Learning Domains: A Framework and Survey](https://mlsanity.com/paper/2003.04960/) (2020)
- [Maximum a Posteriori Policy Optimisation](https://mlsanity.com/paper/1806.06920/) (2018)
- [Automatic Goal Generation for Reinforcement Learning Agents](https://mlsanity.com/paper/1705.06366/) (2017)
- [The Primacy Bias in Deep Reinforcement Learning](https://mlsanity.com/paper/2205.07802/) (2022)
- [Multi-Task Learning as a Bargaining Game](https://mlsanity.com/paper/2202.01017/) (2022)
- [Teacher-Student Curriculum Learning](https://mlsanity.com/paper/1707.00183/) (2017)
- [Prioritized Level Replay](https://mlsanity.com/paper/2010.03934/) (2020)
- [Paired Open-Ended Trailblazer (POET): Endlessly Generating Increasingly Complex and Diverse Learning Environments and Their Solutions](https://mlsanity.com/paper/1901.01753/) (2019)
- [Skew-Fit: State-Covering Self-Supervised Reinforcement Learning](https://mlsanity.com/paper/1903.03698/) (2019)
- [Multi-Task Reinforcement Learning with Context-based Representations](https://mlsanity.com/paper/2102.06177/) (2021)
- [Distributionally Robust Language Modeling](https://mlsanity.com/paper/1909.02060/) (2019)
- [Multi-Task Reinforcement Learning with Soft Modularization](https://mlsanity.com/paper/2003.13661/) (2020)
- [Goal-Conditioned Reinforcement Learning with Imagined Subgoals](https://mlsanity.com/paper/2107.00541/) (2021)
- [Rapid Exploration for Open-World Navigation with Latent Goal Models](https://mlsanity.com/paper/2104.05859/) (2021)
- [Teacher algorithms for curriculum learning of Deep RL in continuously parameterized environments](https://mlsanity.com/paper/1910.07224/) (2019)
- [Enhanced POET: Open-Ended Reinforcement Learning through Unbounded Invention of Learning Challenges and their Solutions](https://mlsanity.com/paper/2003.08536/) (2020)
- [C-Learning: Learning to Achieve Goals via Recursive Classification](https://mlsanity.com/paper/2011.08909/) (2020)
- [Automatic Curriculum Learning through Value Disagreement](https://mlsanity.com/paper/2006.09641/) (2020)
- [CURIOUS: Intrinsically Motivated Modular Multi-Goal Reinforcement Learning](https://mlsanity.com/paper/1810.06284/) (2018)
- [FAMO: Fast Adaptive Multitask Optimization](https://mlsanity.com/paper/2306.03792/) (2023)

## Abstract (arXiv)

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.