Distributed Model-Based Diffusion (DMBD)

Overview

Distributed Model-Based Diffusion (DMBD) is a distributed robot-server framework for scalable multi-robot trajectory optimization based on Model-Based Diffusion (MBD).

Trajectory optimization for multi-robot systems is challenging due to the high-dimensional, non-convex, and often non-differentiable nature of multi-robot coordination problems. Existing sampling-based optimization methods typically operate over the joint trajectory space of all robots, causing the optimization dimension to grow with team size.

DMBD addresses this challenge by distributing the reverse diffusion process across individual robots. Instead of performing denoising over the full joint trajectory space, each robot independently performs local conditional denoising in its own control space while conditioning on the current trajectory estimates of other robots at each denoising step.

A central server aggregates and broadcasts the robots' current state trajectories at each denoising iteration.

For more information, please check our paper.

Key Idea

For a team of robots, centralized MBD operates in the joint trajectory space, whereas DMBD decomposes the inference problem into local trajectory spaces.

Thus, each robot performs trajectory optimization using its own:

• Dynamics
• Objective function
• Local constraints
• Control space

while conditioning on the current trajectory estimates of the other robots.

This decomposition improves scalability and sample efficiency while requiring only trajectory information to be exchanged between robots and the server.

---

Why DMBD?

Compared with centralized sampling-based trajectory optimization, DMBD:

• Decomposes high-dimensional inference into per-robot inference problems.
• Improves scalability as the number of robots increases.
• Supports heterogeneous robots with different dynamics and objectives.
• Requires only trajectory information to be communicated.
• Is zeroth-order and does not require differentiable objectives or constraints.
• Retains the multi-modal trajectory optimization capabilities of MBD.
• Does not require offline training or learned score networks.

---

Simulations

We evaluate DMBD on a diverse set of multi-robot trajectory optimization tasks:

• Circular Goal Swap (N=2-20)
• Rectangular Goal Swap (N=2-20)
• Multi-Floor Coverage (N=2-10)
• Parking (N=2)
• Rush Hour (N=4)

Experiments use:

• Double-integrator dynamics
• Kinematic bicycle dynamics and support both circular and rectangular robot geometries.

Representative Results

	Swap with 20 Circular Robots	Swap with 20 Rectangular Robots	Multi-Floor Coverage with 10 Rectangular Robots	Narrow Parking	Rush Hour
Denoising
Trajectory

Additional Scenarios

	Narrow-Space Parking (N=6)	Rush Hour (N=5)
Denoising
Trajectory

Installation

The code is implemented in Python and uses JAX v0.9 for accelerated numerical computation.

Install the requirements

python -m venv .venv
source .venv/bin/activate

pip install -r requirements.txt

python main.py