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

Riccardo Barbano; Vincent Pauline; Runchang Li; George Webber; Alexander Denker; Željko Kereta; Stefan Bauer; Francisco Vargas

cs.ROOct 6, 2026

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

Authors: Riccardo Barbano, Vincent Pauline, Runchang Li, George Webber, Alexander Denker, Željko Kereta, Stefan Bauer, Francisco Vargas, +1 more

Organizations: UCL · TUM & MCML · CUHK · KCL · DESY · TUM & MCML & CIFAR · Proxima Bio · University of Edinburgh & CIFAR

Abstract

Deep generative models often produce structured outputs composed of interacting components. Modelling these outputs with a single model requires learning both the component distributions and their interactions. We pursue a modular alternative: reuse independently trained component generators and learn only how to coordinate them to produce coherent structured outputs. Our framework, Coordinated Multi-Agent Diffusion Steering (CMDS), treats frozen pretrained diffusion models as reusable generative primitives and coordinates their reverse processes through a learned control. We formulate coordination as a stochastic optimal control problem, balancing an assembly-level reward that specifies the desired properties of the combined output against deviations from the pretrained dynamics. The learned control amortises this optimisation, allowing reuse across new task instances. Experiments show that CMDS can recover a known target distribution, satisfy different spatial constraints with the same trained control, and recover individual sources from degraded mixtures. Across multi-agent maze navigation, articulated robot planning, and text-conditioned human motion, CMDS turns frozen models into coordinated multi-agent generators.

Figures & tables

Figure 1: Overview of CMDS . Frozen component diffusion models are sampled concurrently, while a joint control conditions on the multi-agent state and applies component-wise controls throughout the reverse process, yielding a coordinated assembled output.

Figure 2: Two-agent OU problem. Terminal samples from the independent OU reference process and the rotation-equivariant CMDS control for a prescribed separation $d_{T}=4$ . Yellow and blue points denote the two agents, and each line joins the terminal states belonging to the same sample.

Figure 4: Source separation from a blurred mixture ( $\tau=1$ ). Left: circle–square reconstructions for two held-out mixtures; red and blue hatching behind the CMDS reconstructions shows same-seed matched reference sources. Right: comparison of the reference, DPS, and CMDS over $1{,}024$ held-out mixtures, measuring clean-mixture recovery (RMSE, $\downarrow$ ), labelled-source overlap (IoU, $\uparrow$ ), and shape distortion ( $\downarrow$ ). Full metric definitions are in Section D.3 ; aggregate results and representative reconstructions under milder blur ( $\tau=0.5$ ) are shown in figs. 15 and 16 .

Figure 5: PointMaze (Medium) coordination. Dashed ( ) and solid ( ) paths use the same initial noise. Reference (uncontrolled) process trajectories collide, whereas CMDS steers to avoid collisions while preserving their start–end pairs.

	Medium				Large
	$N=2$		$N=4$		$N=2$		$N=4$
Method	Pool 1 $\uparrow$	Pool 2 $\uparrow$	Pool 1 $\uparrow$	Pool 2 $\uparrow$	Pool 1 $\uparrow$	Pool 2 $\uparrow$	Pool 1 $\uparrow$	Pool 2 $\uparrow$
Reference	0.174 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.005}$	0.195 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.008}$	0.068 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.005}$	0.010 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.005}$	0.195 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.016}$	0.219 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.000}$	0.039 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.000}$	0.049 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.005}$
Best-of-256	0.216 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.005}$	0.263 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.016}$	0.076 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.005}$	0.044 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.005}$	0.232 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.009}$	0.328 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.021}$	0.076 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.005}$	0.070 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.000}$
FK-Steering	0.797 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.028}$	0.820 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.034}$	0.398 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.036}$	0.331 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.016}$	0.924 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.030}$	0.953 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.008}$	0.622 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.009}$	0.729 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.023}$
DPS	0.964 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.024}$	0.956 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.016}$	0.958 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.025}$	0.919 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.033}$	0.984 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.016}$	0.977 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.000}$	0.880 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.030}$	0.927 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.009}$

Table 1: PointMaze joint success. Mean

\pm

standard deviation (std) across sampling seeds. Training-seed variability is reported separately in Section D.4.8 . Each pool contains

128

instances per layout and agent count. For CMDS and CMDS +TRG, we train five independent seeds and report the seed with median Pool 1 success for each setting, breaking ties using Pool 2.

Figure 6: A dual-arm coordination example. Poses share normalised path progress $s=0.135$ ; crosses mark arm–arm contact, and red intervals mark strict mesh collisions. PBDM uses $n_{\mathrm{cand}}=6$ with up to five repair attempts per candidate, versus $n_{\mathrm{cand}}=1$ for the other methods; its modified, native-accepted path fails strict checking between knots. This outcome-selected illustration is not a matched-budget comparison. Additional poses and clearance traces are provided in Section D.5 .

Method	Strict success (%, $\uparrow$ )	Planning time (s, $\downarrow$ )
PBDM	60.9 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 8.0}$	0.427 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.003}$
Dual KUKA + guidance	74.3 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 5.7}$	0.356 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.002}$
Dual KUKA prior	36.3 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 8.7}$	0.274 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.033}$
Single-arm priors + guidance	51.0 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.6}$	0.355 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.001}$
Reference single-arm priors	11.4 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 1.5}$	0.193 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.002}$
CMDS +TRG	77.3 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.7}$	0.368 $\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.002}$

Table 2: Dual KUKA planning with one initial joint candidate. Mean

\pm

sample std. strict successes on 413 held-out test tasks. Planning time is the median value and includes scene setup, generation, verification and optional refinement.

Table 9

Figure 7: Four-person motion coordination. We compare reference MDM, PCD++, and CMDS on the four-person coordination task. Solid curves show left- and right-foot trajectories, with intermediate poses faded; the top-down insets summarise the trajectories and target completion. Here, the reference produces collisions and misses targets, whereas PCD++ and CMDS reach all targets without collisions. Figure 22 in Section D.6 compares all methods.

Appendix figures & tables26 assets

Supplementary material from the paper’s appendix.

Appendix

Variables and processes
Quantity	Notation
Abstract component & assembly	$X^{i},\;X,\;Y=\varphi(X)$
(Uncontrolled) reference reverse process	$\hat{X}_{s}^{i},\;\hat{X}_{s},\;\hat{Y}_{s}$
Generic controlled reverse process	$\hat{X}_{s}^{u,i},\;\hat{X}_{s}^{u},\;\hat{Y}_{s}^{u}$
Parametrised controlled reverse process	$\hat{X}_{s}^{\theta,i},\;\hat{X}_{s}^{\theta},\;\hat{Y}_{s}^{\theta}$
Stop-gradient controlled rollout	$\hat{X}_{s}^{\bar{\theta},i},\;\hat{X}_{s}^{\bar{\theta}},\;\hat{Y}_{s}^{\bar{\theta}}$

Appendix

Table 5: Notation convention. Hatted variables denote reverse-time quantities. The superscript

u

denotes a generic control, while

\theta

denotes the process induced by the parametrised control

U_{\theta}

. For the continuous-time process notation, quantities without a control superscript refer to the uncontrolled reference process, unless stated otherwise. In the discrete-time algorithms and implementation descriptions,

X_{\ell}

denotes the rollout state, with the control superscript suppressed when the generating method is clear from the context.

Figure 8: Two-agent SOC toy: discrete adjoint. Two agents start from an OU reference and are rewarded for reaching terminal configurations with a prescribed separation. We compare the (uncontrolled) OU reference, a value-function oracle, an importance-sampled reference (IS), and the learned control.

Figure 9: Two-agent SOC toy: adjoint matching. Two agents start from an OU reference and are rewarded for reaching terminal configurations with a prescribed separation. We compare the (uncontrolled) OU reference, a value-function oracle, an importance-sampled reference (IS), and the learned control.

Method	Optim. scheme	ELBO $(\uparrow)$	$\log Z_{\mathrm{IS}}$	Valid $(\uparrow)$
$N=1$ (1-square)	Adj. Match.	-1303.1	-7.2	0.0%
$N=1$ (1-square)	Disc. Adj.	-1309.5	-7.2	0.0%
$N=1$ (2-square)	Adj. Match.	-289.1	-0.7	18.9%
$N=1$ (2-square)	Disc. Adj.	-173.4	-0.7	25.4%
CMDS	Adj. Match.	-74.3	-1.6	99.2%
CMDS	Disc. Adj.	-61.2	-1.6	98.5%

Appendix

Table 6: Classifier-only steering diagnostics. We report

\log Z_{\mathrm{IS}}

as an importance-sampling estimate of the log normalising constant of the reward-tilted target, providing a calibration diagnostic for the strength of the induced tilt.

Figure 10: Left-right steering overview. Two copies of a pretrained single-square diffusion generator are coordinated through a controlled reverse process and assembled by the fixed masking operator $\hat{Y}^{\theta}_{1}=\varphi(\hat{X}^{\theta,1}_{1},\hat{X}^{\theta,2}_{1}).$ Black filled squares show the controlled terminal assembly $\hat{Y}^{\theta}_{1}$ . Red and blue dashed outlines show the paired counterfactual uncontrolled terminal proposals $\hat{X}^{1}_{1}$ and $\hat{X}^{2}_{1}$ , obtained from the same initial noise draw but with the learned control removed. Thus, each dashed outline indicates where the corresponding generator copy would have ended up under the reference product law $\hat{\mathbb{P}}$ , before reward-based steering of the assembled canvas. The training curves report the terminal reward $r(\hat{Y}^{\theta}_{1})$ and the path-space objective $\mathrm{ELBO}(\theta)=\mathbb{E}_{\hat{\mathbb{P}}^{\theta}}[\lambda r(\hat{Y}^{\theta}_{1})-A_{\theta}],$ where $A_{\theta}$ denotes the accumulated control/action cost. The learned controls steer the two locally valid component generators into their assigned left/right regions while remaining close to the pretrained dynamics. The path-action cost remains low because the uncontrolled generator already places substantial mass near valid square configurations, so the control only needs a moderate path-space perturbation to steer samples into high-reward assemblies.

Figure 11: Left-right steering (Samples). Samples from the aggregated canvas after training. Black filled squares denote the CMDS assembled output. Red and blue dashed outlines denote the paired uncontrolled terminal proposals from the two generator copies, obtained from the same initial/noise draw but with the learned control removed. Thus, for each sample, the outlines show where that instance would have landed without reward-based steering of the assembled canvas. The learned controls steer each square into its assigned half of the canvas, yielding valid stitched assemblies under the fixed left/right masking operator.

Method	Optim. scheme	ELBO $(\uparrow)$	$\log Z_{\mathrm{IS}}$	Valid $(\uparrow)$
$N=1$ (1-square)	Adj. Match.	-1706.8	-22.8	0.0%
$N=1$ (1-square)	Disc. Adj.	-1565.0	-22.8	0.0%
$N=1$ (2-square)	Adj. Match.	-1561.7	-4.5	8.3%
$N=1$ (2-square)	Disc. Adj.	-381.4	-4.5	9.6%
CMDS	Adj. Match.	-156.7	-5.4	97.9%
CMDS	Disc. Adj.	-190.9	-5.4	96.8%

Appendix

Table 7: Coupled steering diagnostics. We report

\log Z_{\mathrm{IS}}

as an importance-sampling estimate of the log normalising constant of the reward-tilted target, providing a diagnostic for the strength of the induced tilt. Single-process baselines struggle to satisfy the coupled validity constraint, whereas CMDS coordinates two pretrained one-square generators through the assembled canvas and achieves near-perfect validity with improved ELBO.

Training scheme	ELBO ${}_{d_{\Delta x}=5}$ $(\uparrow)$	$\left.\log Z_{\mathrm{IS}}\right\|_{d_{\Delta x}=5}$	Valid $(\uparrow)$
			$d_{\Delta x}=3$	$d_{\Delta x}=4$	$d_{\Delta x}=5$	$d_{\Delta x}=6$	Avg.
Adj. Match.	-172.2	-5.4	95.8%	95.4%	97.3%	97.1%	96.4%
Disc. Adj.	-158.5	-5.4	95.2%	95.7%	97.2%	95.5%	95.9%

Appendix

Table 8: Amortized steering diagnostics. ELBO and

\log Z_{\mathrm{IS}}

are evaluated at

d_{\Delta x}=5

. Validity is evaluated at every supported target using 4096 samples per target with matched random seeds across target conditions. Avg. is the macro-average over

d_{\Delta x}\in\{3,4,5,6\}

Figure 12: Fixed-target coupled steering ( $d_{\Delta x}=5$ ). Black squares show the controlled assembly, red and blue hatched squares the matched reference terminal proposals, green arrows the induced corrections. Yellow arrows indicate the achieved horizontal separation $\Delta x$ .

Figure 13: Amortised coupled steering. Samples from a single pair of agent-specific controls trained jointly over $d_{\Delta x}\in\{3,4,5,6\}$ . Black squares show the controlled assembly, red and blue hatched squares the matched reference terminal proposals, and green arrows the induced corrections. Yellow arrows indicate the achieved horizontal separation $\Delta x$ .

Figure 14: Amortised assembly from copies of a single-square prior. Two frozen copies of a generator pretrained only on one-square canvases are coordinated by a control and combined through a fixed left/right mask. Left: for each requested separation $d_{\Delta x}\in\{3,4,5,6\}$ , red and blue dashed outlines show paired uncontrolled terminal proposals obtained from the same noise, black squares show the controlled assembly, and green arrows show the induced corrections. Dots indicate feasible barycentre locations within the two masked halves. Right: the separation closely matches the requested target $d_{\Delta x}$ , with points and error bars reporting the sample mean and standard deviation. Yellow arrows indicate the achieved horizontal separation $\Delta x$ .

Figure 15: Source separation from a blurred mixture ( $\tau=0.5$ ). Each task contains a ground-truth circle $X_{\mathrm{gt}}^{1}$ and square $X_{\mathrm{gt}}^{2}$ . We observe the blurred mixture $b={\rm A}_{0.5}(Y_{\mathrm{gt}})={\rm A}_{0.5}(X_{\mathrm{gt}}^{1}+X_{\mathrm{gt}}^{2})$ and recover the labelled sources. Left: two representative evaluation mixtures. The red and blue hatched shapes behind the CMDS reconstructions are the matched reference terminal sources. Right: results over $1024$ evaluation mixtures for the reference process, DPS, and CMDS . Reconstruction RMSE measures whether the two recovered sources add up to the correct clean mixture. Source IoU measures whether each recovered object occupies the same pixels as its corresponding ground-truth object ( $\uparrow$ ): it is $1$ for an exact match and $0$ when neither recovered source overlaps its corresponding ground-truth source. The bars show the evaluation-set mean with $95\%$ bootstrap confidence intervals computed from $10{,}000$ resamples. Shape distortion measures whether source $1$ still looks like a circle and source $2$ like a square, irrespective of their location ( $\downarrow$ ). It is zero for an exact valid shape and increases when an object is blurred, deformed, missing, or has the wrong shape.

Figure 16: Representative samples ( $\tau=0.5$ ). Each row represents a different instance of the problem. The columns are: (i) the ground-truth mixture; (ii) the observation $b$ ; (iii) the DPS mixture reconstruction; (iv) the CMDS mixture reconstruction, with the matched reference circle and square shown using red and blue hatching, respectively; (v) the blurred CMDS reconstruction overlaid with the observation $b$ ; and (vi–vii) the CMDS estimates of the circle and square sources, respectively. In the measurement overlay, green indicates agreement between the predicted and observed mixtures, red indicates structure present only in the observation, and blue indicates structure present only in the predicted mixture.

Figure 17: Representative samples ( $\tau=1.0$ ). Each row represents a different instance of the problem. The columns are: (i) the ground-truth mixture; (ii) the observation $b$ ; (iii) the DPS mixture reconstruction; (iv) the CMDS mixture reconstruction, with the matched reference circle and square shown using red and blue hatching, respectively; (v) the blurred CMDS reconstruction overlaid with the observation $b$ ; and (vi–vii) the CMDS estimates of the circle and square sources, respectively. In the measurement overlay, green indicates agreement between the predicted and observed mixtures, red indicates structure present only in the observation, and blue indicates structure present only in the predicted mixture.

Panel of Figure 18 — Figure 18: Fixed PointMaze cell partitions. Coordinates use zero-based $(\mathrm{row},\mathrm{column})$ matrix indexing. The colours show the exact partition of cells between Pool 1 and Pool 2.

	Medium ( $N=2$ )		Medium ( $N=4$ )		Large ( $N=2$ )		Large ( $N=4$ )
Matching loss	Pool 1 $\uparrow$	Pool 2 $\uparrow$	Pool 1 $\uparrow$	Pool 2 $\uparrow$	Pool 1 $\uparrow$	Pool 2 $\uparrow$	Pool 1 $\uparrow$	Pool 2 $\uparrow$
Squared $\ell_{2}$	$0.940\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.020}$	$0.884\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.048}$	$\mathbf{0.806}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.155}$	$0.667\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.151}$	$0.833\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.041}$	$0.792\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.042}$	$0.723\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.157}$	$0.606\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.152}$
Pseudo-Huber	$\mathbf{0.988}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.012}$	$\mathbf{0.933}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.032}$	$0.805\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.098}$	$0.665\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.169}$	$\mathbf{0.907}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.037}$	$\mathbf{0.833}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.045}$	$\mathbf{0.873}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.027}$	$\mathbf{0.744}\,{\scriptstyle\pm 0.036}$

Appendix

Table 9: Adjoint matching training-seed variance and its mitigation. Held-out joint success, mean

\pm

population standard deviation across

10

independent training seeds per (maze,

N

) cell.

	Medium ( $N=2$ )		Medium ( $N=4$ )		Large ( $N=2$ )		Large ( $N=4$ )
Method	Pool 1 $\uparrow$	Pool 2 $\uparrow$	Pool 1 $\uparrow$	Pool 2 $\uparrow$	Pool 1 $\uparrow$	Pool 2 $\uparrow$	Pool 1 $\uparrow$	Pool 2 $\uparrow$
CMDS +TRG	$\mathbf{0.998}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.003}$	$\mathbf{0.983}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.015}$	$\mathbf{0.984}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.005}$	$\mathbf{0.947}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.006}$	$\mathbf{0.973}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.009}$	$\mathbf{0.906}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.015}$	$\mathbf{0.975}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.006}$	$\mathbf{0.969}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.005}$

Appendix

Table 10: CMDS +TRG training variance. Held-out joint success rate (Pool 1: train-cell instances, Pool 2: instances from entirely held-out maze cells), mean

\pm

population standard deviation across 5 independent training seeds, evaluated on 128 held-out instances per split.

Panel of Figure 19 — Figure 19: Qualitative maze-coordination examples. For each maze and agent count, the top row is the reference process trajectory (dashed) and the bottom row is the trajectory produced with the trained CMDS control (solid), for the same $6$ held-out instances and the same initial noise.

Population	Construction	Evaluation tasks
In-domain	Released dual-arm scene and endpoints	413
Shifted obstacles	Horizontal Gaussian shifts, standard deviation 0.10 m	191
Larger obstacles	Each box half-extent multiplied by 1.2	121
Wider bases	Base positions moved to $x=\pm 0.6$ m	208
Three arms	Third base at $(0,0.8,0)$ m, identity rotation	172

Appendix

Table 11: KUKA evaluation populations.

Setting	Final configuration
Trajectory / reverse steps	56 knots / 10 DDIM steps
Stochasticity / prior conditioning	$\eta=1$ / $\omega_{\mathrm{cfg}}=2$
Controller	Width 344, 4 heads, 2 axial blocks
Feed-forward / pair-bias hidden width	$4\times$ controller width / 32
Optimiser	Adam, $(\beta_{1},\beta_{2})=(0.9,0.999)$ , $\epsilon=10^{-8}$
Updates / batch / microbatch	20,000 / 256 joint queries / 128

Appendix

Table 12: KUKA controller-training configuration. Values are taken from the final saved training configuration; the priors remain frozen.

Figure 20: Synchronized KUKA configurations. Reference (two independent single-arm priors), Dual KUKA, PBDM and CMDS+TRG on the same task at common path progress $s\in\{0,0.135,0.3,1\}$ : shared start, arm-arm collision, obstacle collision, shared goal. Squares mark starts, circles goals, red crosses collisions. Only CMDS+TRG remains collision-free.

Figure 21: KUKA poses and inter-arm clearance. Reference (two independent single-arm priors), Dual KUKA, PBDM and CMDS+TRG on the same held-out task at normalized path progress $s=0.135$ (dotted line), with the inter-arm clearance along each path below. Wrist paths fade from start to goal; squares mark starts, circles goals, red crosses collisions at the shown pose. Shaded negative clearance indicates collision; only CMDS+TRG stays positive over the full path.

Primitive	Joint endpoints	Radius (m)
Pelvis	$(1,2)$	0.120
Torso	$(0,9)$	0.150
Shoulder span	$(16,17)$	0.100
Upper torso–head	$(9,15)$	0.100
Head	$(15,15)$	0.110
Upper arms	$(16,18),\ (17,19)$	0.055

Appendix

Table 13: Body-capsule geometry for collision costs and checks. Endpoints are zero-based HumanML joint indices; the repeated head endpoint defines a sphere.

Setting	Recorded value
$(\lambda_{\mathrm{col}},\lambda_{\mathrm{goal}},\lambda_{\mathrm{prog}})$ , training	$(100,50,1)$
$(\lambda_{\mathrm{col}},\lambda_{\mathrm{goal}},\lambda_{\mathrm{prog}})$ , baselines	$(100,1,1)$
$(\lambda_{\mathrm{bone}},\lambda_{\mathrm{FK}},\lambda_{\mathrm{disp}})$	$(500,5,5)$
$(\lambda_{\mathrm{contact}},\lambda_{\mathrm{floor}},\lambda_{\mathrm{speed}})$	$(1,200,10)$
Outer SOC reward multiplier $\lambda$	$100$ , control training only
Optimiser / updates / learning rate	Adam / $15{,}000$ / constant $10^{-4}$

Appendix

Table 14: Reward weights and controller training settings for the reported

N=4

experiment. Baselines share the cost normalisations but use a smaller goal weight.

Figure 22: Extended qualitative comparison of four-person motion generation. We compare uncontrolled diffusion, Best-of-256, DPS, PCD, PCD++, and CMDS on the same four-person task shown in the main paper. The prompts are “A person moves forward at a jogging pace.” for Person 1 ( $\bullet$ ) and Person 4 ( $\bullet$ ), and “A person walks forward.” for Person 2 ( $\bullet$ ) and Person 3 ( $\bullet$ ). Top-down insets show starts and target regions. Solid curves show foot trajectories; intermediate poses are faded, while start and final poses are opaque. Red crosses mark inter-person collisions. In this example, uncontrolled diffusion, Best-of-256, DPS, and PCD do not satisfy the full joint goal-and-collision constraint, whereas PCD++ and CMDS reach all goals without inter-person collisions.

Method	Joint success $\uparrow$ (%)	Collision-free $\uparrow$ (%)	Peak body overlap $\downarrow$ (cm)	All-person goals $\uparrow$ (%)	Goal error $\downarrow$ (m)	Kinematic validity $\uparrow$ (%)	Time $\downarrow$ (s/assembly)
Best-of-256	$0.0\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$	$16.7\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 5.9}$	$8.860\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 1.086}$	$0.0\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$	$2.456\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.146}$	$83.3\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 6.6}$	$69.895\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.292}$
DPS	$0.0\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$	$61.7\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 10.8}$	$3.335\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 1.049}$	$0.0\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$	$3.751\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.284}$	$57.5\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 6.8}$	$8.566\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.018}$
PCD	$4.2\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 2.9}$	$8.3\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 2.9}$	$19.002\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 1.865}$	$\mathbf{100.0}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$	$0.500\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.000}$	$27.5\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 4.8}$	$9.985\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.064}$
PCD++	$75.0\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 6.6}$	$\mathbf{98.3}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 2.3}$	$\mathbf{0.000}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.000}$	$98.3\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 2.3}$	$0.487\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.001}$	$75.0\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 6.6}$	$145.335\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 4.460}$
CMDS	$\mathbf{90.0}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 6.3}$	$90.0\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 6.3}$	$0.740\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.351}$	$\mathbf{100.0}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$	$\mathbf{0.158}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.008}$	$\mathbf{100.0}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$	$\mathbf{0.541}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.002}$

Appendix

Table 15: Four-person locomotion coordination. All methods use the same

24

scenes with walking, jogging, and crouched walking, evaluated over five sampling seeds (

120

attempts per method). Entries are mean

\pm

sample standard deviation across the five seed summaries. Binary rates weight all

24

attempts per seed equally, including generation failures as zero successes; continuous geometry and time use available generated outputs. Timing excludes failed-attempt costs. Bold denotes the best mean, including ties.

Explore similar work

May 12, 2026cs.RO

Coordinated Diffusion: Generating Multi-Agent Behavior Without Multi-Agent Demonstrations

Imitation learning powered by generative models has proven effective for modeling complex single-agent behaviors. However, teaching multi-agent systems, like multiple arms or vehicles, to coordinate through imitation learning is hindered by a fundamental data bottleneck: as the joint state-action space grows exponentially with the number of agents, collecting a sufficient amount of coordinated multi-agent demonstrations becomes extremely costly. In this work, we ask: how can we leverage single-agent demonstration data to learn multi-agent policies? We present Coordinated Diffusion (CoDi), a framework that couples independently trained single-agent diffusion policies through a user-defined multi-agent cost function, without requiring any coordinated demonstrations. We derive a new diffusion-based sampling scheme wherein the diffusion score function decomposes into independent, single-agent pre-trained base policies plus a cost-driven guidance term that coordinates these base policies into cohesive multi-agent behavior. We show that this guidance term can be estimated in a gradient-free manner, making CoDi applicable to black-box, non-differentiable cost functions without additional training. Theoretically and empirically, we analyze the conditions under which this composition can faithfully approximate a target multi-agent behavior. We find a complementary role for demonstration data versus the cost function: single-agent demonstrations must cover the support of the desired multi-agent behavior, while the cost function must promote desired behavior from this product of single-agent policies. Our results in simulation and hardware experiments of a two-arm manipulation task show that CoDi discovers robust coordinated behavior from single-agent data, is more data-efficient than multi-agent baselines, and highlights the importance of joint guidance, base policy support, and cost design.

Lasse Peters, Laura Ferranti, Andrea Bajcsy +1

University of California, Berkeley · Delft University of Technology · Carnegie Mellon University

Diffusion Policies Imitation Learning Diffusion Dynamics Agent Behavior Modeling

Ask PDF

May 2, 2026cs.AI

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

Generative models have emerged as a promising paradigm for offline multi-agent reinforcement learning (MARL), but existing approaches require many iterative sampling steps. Recent few-step acceleration methods either distill a joint teacher into independent students or apply averaged velocity fields independently to each agent. Unfortunately, these few-step approaches hurt inter-agent coordination. We show that the efficiency-coordination trade-off is not inherent: single-pass multi-agent generation can preserve coordination when the velocity field is natively joint-coupled. We propose Coordinated few-step Flow (CoFlow), an architecture that combines Coordinated Velocity Attention (CVA) with Adaptive Coordination Gating. A finite-difference consistency surrogate further replaces memory-prohibitive Jacobian-vector product backpropagation through the averaged velocity field with two stop-gradient forward passes. Across 60 configurations spanning MPE, MA-MuJoCo, and SMAC, CoFlow matches or surpasses Gaussian policies, value-based methods, transformer policies, diffusion models, and prior flow baselines on episodic return. Three independent coordination probes confirm that CoFlow's improvements arise from inter-agent coordination rather than per-agent capacity. A denoising-step sweep shows that single-pass inference suffices on every configuration. CoFlow reaches state-of-the-art coordination quality in 1-3 denoising steps under both centralized and decentralized execution. Project Page: https://guowei-zou.github.io/coflow/

Guowei Zou, Haitao Wang, Beiwen Zhang +2

Sun Yat-sen University

Flow Map

Ask PDF

Sep 29, 2026cs.LG

Multi-Agent Flow Matching with Decoupled Generative Guidance

Generative modeling is widely used for producing diverse objects from complex, multimodal distributions. However, its expressivity does not, in general, come with formal guarantees that the generated objects satisfy hard constraints or requirements. In multi-agent generation, this problem becomes more challenging because a hard requirement can depend on multiple agents, while each agent may need to determine its own guidance input without relying on the simultaneously computed guidance inputs of other agents. To this end, we introduce DeGG-Flow, a general framework for multi-agent flow matching with decoupled generative guidance. By representing the generative process as a control-affine dynamical system, we develop guidance conditions for two classes of coupled requirements: shared requirements whose satisfaction depends on multiple agents together, and private requirements associated with each individual agent dependent on its neighbors. For both classes, we establish feasibility conditions and finite-horizon convergence guarantees. We further derive a Wasserstein bound that characterizes the distributional deviation induced by the guidance. We demonstrate DeGG-Flow on multi-robot collaboration for crossing a spatial gap by reconfiguring the environment, and on multi-object scene generation with affordance requirements. Across both applications, DeGG-Flow directly generates objects that satisfy all corresponding hard requirements, including at team sizes unseen during training.

Ruoyu Lin, Magnus Egerstedt, Fabio Pasqualetti

Department of Electrical Engineering and Computer Science, University of California, Irvine, Irvine, CA 92697 USA · University of North Carolina at Chapel Hill, Chapel Hill, NC 27599 USA

Conditional Flow Matching Generative Processes Indoor Scene Generation Guidance

Ask PDF