Source-linked AI summary
Out-of-Distribution Generalisation with Sequence Models in Offline Multi-Agent Reinforcement Learning
Oussama Hidaoui, Omer Ebead, Ulrich Armel Mbou Sob, Siddarth Singh, Juan Claude Formanek, Felix Chalumeau, Omayma Mahjoub, Sasha Abramowitz, Ruan John de Kock, Wiem Khlifi, Louay Ben Nessir, Simon Verster Du Toit, Daniel Rajaonarivonivelomanantsoa, Asim Awad Osman, Arnol Manuel Fokam, Refiloe Shabe, Arnu Pretorius
TL;DR
Generalising offline MARL agents to unseen tasks is under-explored despite the cost and risk of collecting new online data. The paper extends sequence models for heterogeneous multi-task settings and evaluates how task diversity, dataset size and model capacity affect transfer. It finds that task diversity is the dominant factor, with multi-task training improving held-out performance and outperforming behaviour cloning.
Problem
Prior offline MARL research largely evaluates on the training task, leaving generalisation to unseen tasks under-explored.
Method
The paper develops a multi-task offline MARL benchmark and extends centralised sequence models to multiple tasks with varying agent counts, analysing task diversity, dataset size and model capacity.
Results
3.2x average zero-shot scaling is observed as training task count increases, while task diversity—not sheer dataset size—is the main driver of test performance; sequence models also outperform behaviour cloning.
Takeaways & Limitations
Offline MARL progress should prioritise large, diverse multi-task datasets and tune model capacity for the available data budget.
Takeaways & Limitations
The analysis is limited to centralised sequence model architectures and does not yet cover decentralised or CTDE algorithms.
Abstract
from arXiv · showhide
Generalising to unseen tasks remains a fundamental challenge in offline multi-agent reinforcement learning (MARL). In this work, we present a principled analysis of zero-shot task generalisation in the offline setting and conduct an extensive empirical investigation into the scaling behaviour governing task diversity, dataset size, and network capacity. To facilitate this study, we extend offline sequence modelling architectures to handle multi-task observation and action spaces alongside variable agent counts across tasks. Our primary finding is that scaling task diversity---rather than sheer dataset size is the dominant factor in achieving robust zero-shot transfer. Through large-scale experiments across four challenging environments (Connector, RWARE, SMAX, and LBF), we demonstrate that our multi-task approach achieves a mean improvement of 3.2x on held-out test tasks compared to single-task models and consistently outperforms strong behaviour cloning baselines. These results suggest that the development of generalisable MARL agents should prioritise the diversity of the training distribution with varying numbers of agents, providing a roadmap for scaling offline MARL effectively.
1 INTRODUCTION
The paper addresses the under-explored problem of zero-shot generalisation to unseen tasks in offline MARL. Across four environments, it finds that diverse multi-task training improves transfer, with task diversity and model capacity more important than simply increasing dataset size.
- Motivation: Prior offline MARL work largely trains and evaluates on the same task, leaving generalisation to unseen tasks under-explored.The paper frames unseen-task generalisation as important for deploying agents beyond their training distribution.
- Findings: Offline MARL sequence models trained on diverse multi-task datasets generalise better to unseen tasks than single-task alternatives.The paper reports this pattern across its empirical studies and model objectives.
- Findings: 3.2x average zero-shot improvement is obtained as the number of training tasks increases, while sheer dataset size is not the main driver of test performance.The analysis varies task diversity, data size and model size to separate their effects on unseen-task performance.
- Analysis: The paper combines empirical scaling experiments with a theoretical analysis identifying task coverage as a driver of out-of-distribution performance gains.Its contributions include formal grounding for the observed roles of task diversity and model capacity.
2 MULTI-TASK SEQUENCE MODELLING FOR OFFLINE MARL
The paper formulates offline multi-task MARL as learning one joint policy from fixed datasets for unseen test tasks, then adapts sequence models to heterogeneous tasks and agent counts. The design uses autoregressive policies, task-balanced training, and agent padding, masking and shuffling.
- Problem formulation: The objective is to learn one joint policy from fixed multi-task training data that maximises zero-shot performance on unseen test tasks.Training uses offline datasets collected from the training-task set without additional online interaction.
- Sequence modelling: Centralised joint-action policies are theoretically optimal but scale poorly, motivating autoregressive factorisation across agents.The joint distribution is represented as a product of conditional action distributions.
- Multi-task architecture: The multi-task sequence models support multiple tasks and varying agent counts without explicit task IDs or task-specific output heads.This design requires the model to infer task information while enabling transfer to new tasks.
- Variable-agent handling: Dynamic padding, masking and randomised agent ordering handle absent agents and encourage shared representations across agents.Absent-agent inputs are zero-padded, their loss contributions are masked, and active and inactive agents are shuffled during training.
- Training objectives: The multi-task models use autoregressive BC, CQL or ICQ losses, with value-function learning optionally framed as classification using HL-Gauss.The paper states that its design choices are validated through ablation studies.
- Training: Task-balanced batching samples evenly across tasks, producing gradients for a uniform task mixture rather than a size-weighted mixture.The method allocates batch samples across tasks using quotient, remainder and round-robin assignment.
3 EMPIRICAL ANALYSIS
The experiments evaluate offline multi-task MARL across four environments, varying task diversity, dataset size, and model capacity. Results show that adding tasks improves held-out generalisation more reliably than adding transitions, while larger models improve both training and test performance.
- Experimental design: The benchmark spans LBF, RWARE, Connector, and SMAX, with tasks partitioned into training and held-out test sets.Offline datasets contain task-specific rollouts, with task-balanced batching used because episode lengths differ.
- Experimental design: Models are evaluated on held-out tasks using normalised episode returns averaged over independent episodes and repeated random seeds.The best-performing checkpoint is evaluated over 320 episodes per test task, with mean and standard deviation reported across three seeds.
- Task diversity: As training task count increases, test-task performance improves nearly monotonically across environments, despite relatively high training performance and some RWARE training degradation.The authors attribute the RWARE decline to task complexity and insufficient capacity as diversity grows.
- Task diversity: 5.4x on RWARE, 1.3x on LBF, 2.9x on Connector, and 3.2x on SMAX are the reported maximum test-task gains from multi-task training.These gains are averaged across all three algorithms and environments separately.
- Dataset size: Increasing transitions improves training performance but provides little evidence of comparable generalisation gains to adding more tasks.The dataset-size sweep fixes the number of RWARE tasks while varying transition count.
4 WHY DOES MULTI-TASK TRAINING HELP MARL SEQUENCE MODELS TO GENERALISE?
The analysis explains generalisation through a bound combining training error, task-space coverage, and value-function smoothness. It connects the theory to empirical evidence that diverse nearby tasks and greater model capacity improve unseen-task performance.
- The generalisation bound combines training error, task-space coverage radius, and a Lipschitz constant for value functions over tasks.
- The theory assumes tasks share a common embedded state-action space and initial distribution, with value functions smooth under a task metric.Smoothness may fail for task families with abrupt structural changes.
- Task diversity improves generalisation by shrinking the distance between held-out tasks and nearby training tasks.The bound decreases as minzi d(zi, ztest) becomes smaller.
- Model capacity improves generalisation through lower training error while leaving the coverage term fixed for a given training-task set.This matches the reported monotonic test-performance improvement with model size.
- For finite task suites, test-task regret is bounded by training error when test tasks are included, and otherwise increases with their maximum distance from training tasks.If every test task lies within distance r of training data, the bound is εtrain + 2Lr.
- Proxy coverage estimates theoretical coverage using hand-crafted task descriptors, with higher values indicating better test-task coverage.Across environments, proxy coverage rises monotonically as training tasks are added, although saturation rates differ.
5 RELATED WORK
Prior work spans offline MARL, sequence modelling, multi-task reinforcement learning, and formal multi-task MARL generalisation. The paper positions itself at their intersection, addressing task generalisation in offline multi-agent sequence models.
- Offline MARL research has largely focused on single-task settings and challenges such as extrapolation error, coordination, and constrained value estimation.
- Sequence-modelling research progressed from Decision Transformer to multi-task and multi-agent extensions, including MADT and offline multi-agent methods.
- Multi-task reinforcement learning has mainly studied representation and transfer in single-agent continuous-control and robotics, with evidence that adding data can sometimes reduce downstream performance.
- Related work suggests high-capacity models trained on diverse datasets may generalise broadly, while prior zero-shot offline studies chiefly examined single-agent settings.
- Multi-task MARL introduces architectural and evaluation challenges, alongside formal theories of zero-shot generalisation in goal-based settings.
6 CONCLUSION
The paper concludes that task diversity and model capacity are central to zero-shot generalisation in offline MARL sequence models. It recommends diverse datasets and capacity tuning, while identifying broader architectures and environments as future directions.
- Task diversity is a key driver of reducing the train–test generalisation gap in offline MARL sequence models.
- Model scaling produces significant zero-shot out-of-distribution improvements across different sequence-model objectives.
- Future offline MARL progress should prioritise large, diverse multi-task datasets and model-capacity tuning for the available data budget.
- The study is limited to centralised sequence-model architectures and does not analyse decentralised or CTDE algorithms.
- Future work includes transfer across environments and faster fine-tuning in safety-critical, data-scarce real-world domains.
A.1 PROOF OF THEOREM 4.4
The theorem’s proof bounds performance differences between tasks using task-descriptor distance, with proxy coverage quantifying how closely test tasks are represented in training. It applies across both state- and history-dependent policies and relies on explicit assumptions about task similarity.
- For unseen tasks, the theorem uses the nearest training task and its distance to bound the generalisation error; seen tasks have zero nearest-neighbour distance.The general case is obtained by taking a supremum over test tasks.
- The value difference between the same policy evaluated on two tasks is bounded by L d(z, z′), where L combines reward and transition sensitivity.The bound is stated as |J_z(π) − J_z′(π)| ≤ L d(z, z′).
- The proof handles history-dependent policies by replacing states with histories because task transitions depend only on the current state-action pair.The policy’s history-conditioned action distribution is held fixed across the two tasks.
- The simulation lemma supplies sufficient conditions for the task-similarity assumption, while Theorem 4.4 requires only that assumption itself.The conditions include bounded rewards and Lipschitz relationships for rewards and transition kernels.
- Proxy coverage is computed from hand-crafted task descriptors, with higher coverage indicating smaller nearest-neighbour error after descriptor normalization.Descriptor components are normalized to [0, 1], and reported coverage is additionally normalized per environment at N = 1.
- The benchmark represents tasks across LBF, Connector, RWARE, and SMAX, whose descriptors capture environment-specific task configurations and agent settings.The environments differ in cooperation structure, observation conditions, action spaces, and task compositions.
C MULTI-TASK OFFLINE MARL CAN GENERALISE BETTER THAN BEHAVIOUR CLONING
The study compares offline objectives for zero-shot transfer across four MARL environments. Multi-task offline RL methods generally outperform behaviour cloning on mixed-quality data, although the strongest objective varies by environment.
- Oryx performs best on LBF and Connector, whereas CQL performs best on RWARE and SMAX.On SMAX, BC closely follows CQL, while Oryx performs worst.
- Mixed-data offline MARL methods exhibit stronger zero-shot generalisation than behaviour cloning.The authors relate this pattern to mixed replay datasets rather than purely expert data.
- Offline RL is the best-performing algorithm on each environment, with CQL or Oryx leading depending on the benchmark.The comparison aggregates normalized episode returns over held-out tasks across LBF, RWARE, Connector, and SMAX.
D DATASET QUALITY ABLATION
Trajectory quality affects training performance more reliably than held-out-task performance. High-quality data can improve training, especially for CQL-Sable, but those gains do not transfer consistently to unseen tasks.
- High-quality trajectories improve training performance, particularly for MT-CQL-Sable, but the gains do not transfer to test tasks.The comparison uses trajectory subsets for MT-Oryx, MT-CQL-Sable, and MT-BC-Sable on RWARE.
- Low-quality trajectories consistently produce the worst results.This pattern is reported across the compared algorithms and trajectory subsets.
- The ablation tests whether trajectory quality explains generalisation after dataset-size increases fail to yield significant unseen-task improvements.The experiment samples trajectories from specific dataset-quality subsets.
E SCALING ANALYSIS ON LBF AND CONNECTOR
Increasing model size improves both training and test performance until a critical threshold, after which performance plateaus. Easier environments reach saturation at smaller models than RWARE.
- LBF, Connector, and SMAX saturate at much smaller model sizes than RWARE because they are considerably easier.The reported scaling trend remains visible for BC-Sable on LBF, although it is more marginal.
- Performance improves with model size up to a critical point, then plateaus across LBF, Connector, and SMAX.The same qualitative scaling behaviour is observed for both training and test performance.
F VISUALISATION OF MULTI-TASK POLICY
Visual roll-outs show that one multi-task model learns distinct team strategies for sparse and congested warehouse tasks, while the figures specify how to read the frames.
- Team strategy visualisation: One multi-task model simultaneously learns distinct strategies for medium-2ag and medium-32ag2.It rapidly explores sparse warehouses in medium-2ag2 and moves completed agents aside in congested medium-32ag2.
- medium-2ag2: In medium-2ag2, two agents rapidly traverse the sparse warehouse to explore and find the shelf.
- medium-32ag2: In medium-32ag2, agents move completed agents to the bottom-right corner to avoid collisions in the congested warehouse.
- Reading the figures: Figures 13 and 14 should be read left to right, then top to bottom.
J ARCHITECTURE DESIGN ABLATIONS
Ablations identify task-balanced batching and agent masking and shuffling as important design choices, while training-task performance remains high as task count increases.
- HL-Gauss: Approximately 8% higher test performance occurs with HL-Gauss for MT Oryx, while its effect on MT CQL-Sable is marginal.The ablation targets multi-task value-function learning on RWARE, where episode-return variance across tasks is significant.
- Agent shuffling and masking: 16% lower test performance occurs on average when agent masking and shuffling are disabled for both algorithms.
- Task-balanced batching: 37% lower test performance occurs without task-balanced batching for MT Oryx and MT CQL-Sable on average.The alternative samples a full batch from one randomly selected task at each update.
- Task scaling: Training-task performance remains high across environments as the number of tasks increases.