Nicholas Edward Corrado - Distributionally Robust Multi-Task Reinforcement Learning via Adaptive Task Sampling

​Talk Description: (Paper link: https://arxiv.org/abs/2605.14350) 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.

DRATS builds off AutoMixAlign (ACL 2025, https://arxiv.org/abs/2506.00569), a DPO-based adaptive data mixing algorithm for LLM alignment, developed with colleagues on Amazon’s Rufus team.

I will spend the first few minutes discussing why I chose to focus my dissertation on data collection, since I think that context helps explain how I approach research more broadly.


Bio: He recently received his PhD from the University of Wisconsin-Madison working under Josiah Hanna. His research focuses on improving the data efficiency of RL via adaptive data collection and data mixing, with applications to LLM alignment, robotics, and multi-agent learning. Previously, he was a research intern with Amazon’s Rufus team (working on multi-objective alignment) and Sandia National Laboratories (working on RL for power grid management). During the first year of his PhD, he worked in databases with Jignesh Patel. He received a BPhil in physics and a BS in mathematics from the University of Pittsburgh, where he studied high-energy physics with Vladimir Savinov.