CoFlow: Coordinated Few-Step Flow for Offline Multi-Agent Decision Making

How to achieve both fast joint trajectory generation and effective coordination in generative offline multi-agent reinforcement learning? Multi-agent diffusion models generate joint trajectories through iterative sampling but incur high inference latency. Distillation reduces sampling steps, but it may cause the student model to lose information about the cross-agent dependencies learned by the teacher, leading to a lower coordination success rate. A directly trained joint generator can model coordination among agents during one-step or few-step generation. Nonetheless, insufficient modeling of cross-agent dependencies may still impair coordination. Furthermore, its consistency training process introduces heavy computational and memory costs. In this paper, we propose Coordinated few-step Flow (CoFlow), which directly learns a joint averaged-velocity field for distillation-free one-step and few-step multi-agent trajectory generation. CoFlow uses Coordinated Velocity Attention (CVA) to incorporate teammate trajectory information for more accurate prediction of each agent's trajectory. To reduce the cost of consistency training for this joint model, CoFlow approximates the correction term using finite differences. Across 48 configurations on MPE and SMAC, CoFlow supports centralized and decentralized execution. Under centralized execution, our CoFlow outperforms our reproduced baseline by 11.2%, averaging relative gains equally across both benchmark suites, with a 12.93× speedup in model sampling. Its CVA improves normalized scores by 165.1% on average over the same model with cross-agent attention disabled. Compared with the implementation using exact derivatives, finite-difference training achieves a 1.78× speedup and reduces peak GPU memory by 41.1%. Code is provided in the supplementary material.