Train-Small Deploy-Large: Leveraging Diffusion-Based Multi-Robot Planning

Abstract: Learning based multi-robot path planning methods struggle to scale or generalize to changes, particularly variations in the number of robots during deployment. Most existing methods are trained on a fixed number of robots and may tolerate a reduced number during testing, but typically fail when the number increases. Additionally, training such methods for a larger number of agents can be both time consuming and computationally expensive. However, analytical methods can struggle to scale computationally or handle dynamic changes in the environment. In this work, we propose to leverage a diffusion model based planner capable of handling dynamically varying number of agents. Our approach is trained on a limited number of agents and generalizes effectively to larger numbers of agents during deployment. Results show that integrating a single shared diffusion model based planner with dedicated inter-agent attention computation and temporal convolution enables a train small deploy-large paradigm with good accuracy. We validate our method across multiple scenarios and compare the performance with existing multi-agent reinforcement learning techniques and heuristic control based methods.

• The paper introduces MA-DBP, a conditional latent diffusion planner using axial attention, masking, curriculum learning, and moving-horizon execution to generalize from small robot teams to larger deployments.
• MA-DBP matches or approaches MAPPO in several navigation settings, remains effective when scaling from 3–4 trained agents to as many as 8 deployed agents, and requires up to 4× less total training time than comparable approaches.
• The method outperforms MADiff and CLF-QP in most tested cases but remains weaker than MAPPO in cluttered environments, with performance beyond eight agents and at higher robot densities still unresolved.

Overview

This paper addresses a persistent limitation of learning-based multi-robot path planning: the inability to generalize to team sizes larger than those seen during training. Most MARL and graph-based planners are trained with a fixed agent count, tolerate at best a reduced count at test time, and fail when the number of agents increases. The authors propose the Multi-Agent Diffusion Based Planner (MA-DBP), a conditional latent diffusion planner trained on a small number of agents that generalizes to larger teams at deployment — a "train-small, deploy-large" paradigm. The method is validated in simulation on three 2-D navigation scenarios against MAPPO, a CLF-QP controller, and MADiff, and is shown to achieve comparable success rates while requiring up to 4× less total training time for $n_a^{train}=8$.

Motivation and related work

Analytical methods such as MAPF solvers offer guarantees but scale poorly computationally and are brittle in dynamic or partially observed environments. Learning-based approaches adapt better but suffer from non-stationarity when agent counts change, and typically require retraining. Prior diffusion-based planning work either fixes the agent count (MADiff trains one denoising network per agent or a fixed-size model), offloads planning from the diffusion model (Shaoul et al.), incurs per-step projection costs (projected diffusion for continuous MAPF), or targets single-robot navigation (CoBL-Diffusion, NavFormer). MA-DBP differs by using a single shared diffusion model whose architecture supports a variable number of agents through masking and dimension-invariant embeddings.

Method

MA-DBP formulates goal-conditioned moving-horizon trajectory generation over a fixed window $[t:t_H]$ for $n_a$ active agents. Three components enable size generalization:

Axial attention pre-processing. Trajectories of shape $B \times H \times n_a \times d$ are projected to an embedding dimension $D$, augmented with learned positional embeddings over both the agent and temporal axes, and processed with attention along each axis separately (agent-axis attention followed by time-axis attention), then refined by an MLP. Because all embeddings share dimensionality regardless of $n_a$, the model can operate on varying agent counts. Critically, this attention computation occurs before the denoising U-Net rather than inside it, which the authors argue avoids the data hunger of transformer-based diffusion backbones and the need to replicate diffusion blocks per agent as in MADiff.

Conditional U-Net denoiser. A 1-D convolutional U-Net with skip connections predicts added noise. Context tokens encoding the scenario image $O_t$ (CNN encoder), start/goal poses (MLP), and agent count (learned embedding) are fused via multi-head attention and injected at every U-Net level through FiLM conditioning together with sinusoidal diffusion-timestep embeddings. Masking mechanisms allow training and execution with fewer active agents than the architectural maximum $n_{max}$, where $n^{max}_{train} < n_{max}$.

Multi-term loss. Beyond the standard noise-prediction loss, the objective includes boundary losses anchoring start and goal poses, temporal-consistency and jerk-minimization terms computed via finite differences, and a collision penalty against other agents and obstacle centers. The weighting vector $W = [0.85, 0.025, 0.025, 0.1]^\top$ emphasizes the noise loss, which the authors found necessary for stable training.

Training uses trajectories bootstrapped from pre-trained MAPPO policies for $[t:t_H]$0 agents, with curriculum learning increasing the agent count linearly every $[t:t_H]$1 epochs. At execution, a moving-window strategy replans over the horizon, and the denoised trajectory is tracked by a low-level controller decoupled from the planner, supporting heterogeneous platforms.

Experimental results

Validation uses VMAS across empty, obstacle, and barrier maps, with success defined as all agents reaching within 0.1 units of their goals within 100 steps, averaged over 20 trials.

When trained and evaluated on identical agent counts, MA-DBP achieves success rates competitive with MAPPO on the empty map (e.g., 0.95 vs. 0.94 at $[t:t_H]$2) but trails MAPPO in cluttered environments (e.g., 0.78 vs. 0.94 at $[t:t_H]$3 obstacles). It consistently outperforms MADiff, which reaches only 0.12–0.55 across all settings despite replanning full horizons, and outperforms CLF-QP in most cluttered cases, where the QP controller lacks collision foresight and becomes stuck. On the empty map, MA-DBP also exhibits faster inference than MAPPO.

The central scalability result shows that models trained with $[t:t_H]$4 or $[t:t_H]$5 degrade only modestly when deployed on up to 8 agents, whereas models trained with $[t:t_H]$6 upscale poorly — attributed to insufficient inter-agent interaction signal during training. Notably, denser goal placement at higher agent counts appears to ease planning, though it simultaneously increases collisions.

Ablations isolate the contributions of the axial attention processor (AAP) versus a linear encoder (LE) and complete-trajectory (CT) versus moving-window (MW) training. For models trained on 4 agents and evaluated on 6–7, MA-DBP attains 0.95/0.87 success versus 0.34/0.31 for CT+LE, 0.71/0.69 for CT+AAP, and 0.72/0.67 for MW+LE. The axial attention processor is identified as the more critical component; removing any of the three ingredients (attention pre-processing, curriculum, moving horizon) leads to failure to upscale.

Limitations and open questions

The paper concedes several constraints. MAPPO remains superior in obstacle-rich scenarios because its one-step sampling avoids the difficulty diffusion models face with short horizons; MA-DBP's advantage is confined to scalability, inference speed, and handling dynamically changing team sizes. Horizon length is a sensitive design parameter: too short starves the attention module of interaction context, too long yields overly static patterns and unstable training. Testing was capped at 8 agents due to environment size, so scaling behavior beyond this is unverified. Higher agent densities increase collision rates even as success rates hold. Finally, the authors note an inherent tension in diffusion planners between coarse long-horizon goal-seeking and fine short-horizon precision, stating that combining both behaviors currently leads to unstable training — leaving open how to unify these regimes.

Conclusion

MA-DBP demonstrates that a conditional diffusion planner with explicit axial attention pre-processing, curriculum training, and moving-horizon execution can generalize from small training teams to larger deployment teams with acceptable success rates, at substantially reduced training cost relative to MARL baselines. The results position diffusion-based planning as a viable route around the fixed-agent-count limitation of MARL, while the residual gaps in cluttered environments and the untested regime beyond eight agents delineate the boundaries of the current evidence.