Scaling Multi-Agent Environment Co-Design with Diffusion Models

The agent-environment co-design paradigm jointly optimises agent policies and environment configurations in search of improved system performance, promising to fundamentally reshape how we deploy multi-agent systems in domains such as warehouse logistics and windfarm management. However, current co-design methods collapse under high dimensional environment design spaces and suffer from sample inefficiency when addressing moving targets inherent to joint optimisation. We address this by developing Diffusion Co-Design (DiCoDe), a scalable and sample-efficient co-design framework incorporating two core innovations. We introduce Projected Universal Guidance (PUG), enabling exploration of constraint-satisfying reward-maximising environments, and devise a critic distillation mechanism to transfer knowledge from the reinforcement learning loop to a guided diffuision model. Together, these improvements lead to superior environment-policy pairs when validated on challenging multi-agent co-design benchmarks, for example, exceeding state-of-the art in a warehouse setting with 39% higher rewards and 66% fewer simulation steps.

Engineers teach robots to pick-and-place boxes, and architects design the layout of shelves in a building. The success of an automated warehouse depends on both of these factors, but currently these decisions are often made separately. Can we teach computers to design smart robots and suitable environments for these robots to operate in at the same time? We took apart diffusion models — the technology powering generative AI — and applied them to this problem. We found that robots could learn which environments bring out their best performance, and we transfer that knowledge to a diffusion model to automatically construct better environments in return. Because diffusion models are so expressive, our method scales to real-world problems that were previously too large to tackle, like optimising the layout of an entire wind farm alongside the control policy of every turbine.