One for All, All for One: Coordinated Multi-Agent Diffusion Steering via Stochastic Optimal Control

Many structured outputs consist of interacting components, but a single generative model must learn both the individual components and the diversity of valid structured outputs they can jointly produce. This motivates a modular alternative: train component generators independently and learn only how to steer them jointly to produce coherent structured outputs. We introduce Coordinated Multi-Agent Diffusion Steering (CMDS), a framework that treats frozen pretrained diffusion models as reusable generative primitives and coordinates their reverse processes concurrently through a learned amortised control. An assembly map combines the generated components, while an assembly-level reward specifies the desired properties of the assembled output. We formulate this coordination through stochastic optimal control, balancing the assembly-level reward against deviation from the pretrained dynamics. Experiments validate amortised satisfaction of relational design specifications and posterior inference for source separation. Notably, in PointMaze, CMDS coordinates frozen single-agent diffusion planners into an amortised planner for up to four agents, enabling all agents to reach their goals while avoiding wall and inter-agent collisions.