# Model-Based Reinforcement Learning for Heterogeneous Multi-Robot Task Assignment Under Distribution Shifts

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in RL for games and combinatorial optimization, one of 15 directions in Deep reinforcement learning methods.
- Builds on 41 of its 111 references that are in the map.
- First posted on arXiv 21 Aug 2026.

## Placed in

- Deep reinforcement learning methods / [RL for games and combinatorial optimization](https://mlsanity.com/deeprl/rl-for-games-and-combinatorial-optimization/)

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

- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [Adaptive Conformal Inference Under Distribution Shift](https://mlsanity.com/paper/2106.00170/) (2021)
- [Data-driven Distributionally Robust Optimization Using the Wasserstein Metric: Performance Guarantees and Tractable Reformulations](https://mlsanity.com/paper/1505.05116/) (2015)
- [Transformer Hawkes Process](https://mlsanity.com/paper/2002.09291/) (2020)
- [Intensity-Free Learning of Temporal Point Processes](https://mlsanity.com/paper/1909.12127/) (2019)
- [Conformal Inference for Online Prediction with Arbitrary Distribution Shifts](https://mlsanity.com/paper/2208.08401/) (2022)
- [Fully Neural Network based Model for General Temporal Point Processes](https://mlsanity.com/paper/1905.09690/) (2019)
- [Safe Planning in Dynamic Environments using Conformal Prediction](https://mlsanity.com/paper/2210.10254/) (2022)
- [Scalable Multi-Robot Collaboration with Large Language Models: Centralized or Decentralized Systems?](https://mlsanity.com/paper/2309.15943/) (2023)
- [Transformer Embeddings of Irregularly Spaced Events and Their Participants](https://mlsanity.com/paper/2201.00044/) (2021)
- [Lifelong Multi-Agent Path Finding for Online Pickup and Delivery Tasks](https://mlsanity.com/paper/1705.10868/) (2017)
- [Combating the Compounding-Error Problem with a Multi-step Model](https://mlsanity.com/paper/1905.13320/) (2019)
- [COHERENT: Collaboration of Heterogeneous Multi-Robot System with Large Language Models](https://mlsanity.com/paper/2409.15146/) (2024)
- [Sample-Efficient Safety Assurances using Conformal Prediction](https://mlsanity.com/paper/2109.14082/) (2021)
- [Trust the Model When It Is Confident: Masked Model-based Actor-Critic](https://mlsanity.com/paper/2010.04893/) (2020)
- [Fast and Flexible Temporal Point Processes with Triangular Maps](https://mlsanity.com/paper/2006.12631/) (2020)
- [Integrated Task Assignment and Path Planning for Capacitated Multi-Agent Pickup and Delivery](https://mlsanity.com/paper/2110.14891/) (2021)
- [Wasserstein Distributionally Robust Stochastic Control: A Data-Driven Approach](https://mlsanity.com/paper/1812.09808/) (2018)
- [Lifelong Path Planning with Kinematic Constraints for Multi-Agent Pickup and Delivery](https://mlsanity.com/paper/1812.06355/) (2018)
- [A Graph Regularized Point Process Model For Event Propagation Sequence](https://mlsanity.com/paper/2211.11758/) (2022)
- [When to Update Your Model: Constrained Model-based Reinforcement Learning](https://mlsanity.com/paper/2210.08349/) (2022)
- [Deep Continuous-Time State-Space Models for Marked Event Sequences](https://mlsanity.com/paper/2412.19634/) (2024)
- [Generalized Mission Planning for Heterogeneous Multi-Robot Teams via LLM-constructed Hierarchical Trees](https://mlsanity.com/paper/2501.16539/) (2025)
- [Robust Task Scheduling for Heterogeneous Robot Teams under Capability Uncertainty](https://mlsanity.com/paper/2106.12111/) (2021)
- [Multi-Robot Task Allocation -- Complexity and Approximation](https://mlsanity.com/paper/2103.12370/) (2021)
- [Online Multi-Agent Pickup and Delivery with Task Deadlines](https://mlsanity.com/paper/2403.12377/) (2024)
- [Is Score Matching Suitable for Estimating Point Processes?](https://mlsanity.com/paper/2410.04037/) (2024)
- [Speeding up Routing Schedules on Aisle-Graphs with Single Access](https://mlsanity.com/paper/2102.05733/) (2021)
- [Autonomous Delivery of Multiple Packages Using Single Drone in Urban Airspace](https://mlsanity.com/paper/2209.01580/) (2022)

## Abstract (arXiv)

Heterogeneous multi-robot service systems must assign requests to compatible robots, construct feasible schedules, and adapt as new tasks arrive online. Historical data can help anticipate future demand, but relying too heavily on inaccurate predictions can degrade performance under distribution shifts. We develop a prediction-aware adaptive rollout framework for heterogeneous multi-robot task assignment with scheduled and real-time requests. The problem is formulated as a finite-horizon stochastic dynamic program incorporating robot-task compatibility, ordered service requirements, routing constraints, service windows, and end-of-horizon return requirements. The proposed policy evaluates current assignments using sampled future request scenarios while restricting immediate commitments to requests already observed. To enable online use, the framework combines pruned candidate controls, wait actions, and an interaction-aware base policy for efficient future-cost estimation. Robustness to forecast error is provided by adaptively reweighting predicted requests based on recent prediction mismatch and selectively re-optimizing assigned but unstarted requests. We also introduce a historical-data-driven procedure for selecting the heterogeneous fleet composition before deployment. In a case study using real nursing-task requests from hospital inpatient floors, the proposed approach achieves near-complete service and reduces serviced-request wait times relative to reactive, token-passing, prediction-positioning, and myopic greedy baselines, with the largest improvements in tail-delay metrics.