OMAD

Online Multi-Agent Diffusion — an online multi-agent reinforcement learning implementation built on diffusion policies. The core algorithm is DIME (diffusion-based off-policy RL). Configuration is managed with Hydra; training uses JAX/Flax on top of the Stable-Baselines3 training loop and replay buffer.

Features

• Diffusion policy actor: continuous actions via DIS (Diffusion Importance Sampler) and a conditional score network
• Multi-agent support: PettingZooToGymnasiumWrapper maps PettingZoo multi-agent envs to a single-agent interface (concatenated observations/actions, summed rewards)
• Environments:
  • PettingZoo MPE: simple_spread, simple_adversary, simple_tag, simple_push, simple_reference
  • MaMuJoCo (gymnasium-robotics): Ant, HalfCheetah, Hopper, Walker2d, and more
• Experiment tracking: optional Weights & Biases and TensorBoard
• Checkpointing: save and resume training

Project structure

OMAD/
├── run_omad.py              # Training entry point (Hydra)
├── configs/                 # Hydra configs
│   ├── base.yaml            # Main config
│   ├── alg/dime.yaml        # DIME hyperparameters
│   └── sampler/             # Diffusion sampler config
├── diffusion/               # DIME algorithm and diffusion policy
│   ├── dime.py
│   ├── diffusion_policy.py
│   └── od/                  # OT / diffusion integrators
├── common/                  # Off-policy algorithm, policy base, buffers
├── models/                  # Critic, evaluation callbacks
└── README.md

Requirements

• Python ≥ 3.9
• CUDA recommended (JAX GPU acceleration)
• Linux (tested on Linux)

Installation

cd OMAD
pip install -r requirements.txt

Core dependencies include jax, flax, optax, stable-baselines3, gymnasium, pettingzoo, hydra-core, and wandb. See requirements.txt for the full list.

Optional environment packages

# PettingZoo MPE
pip install "pettingzoo[mpe]"

# Multi-agent MuJoCo
pip install gymnasium-robotics

Quick start

PettingZoo MPE (simple_spread)

python run_omad.py env_name=mpe/simple_spread seed=0

simple_adversary

python run_omad.py env_name=mpe/simple_adversary seed=0

simple_tag (configurable number of adversaries)

python run_omad.py env_name=mpe/simple_tag env.simple_tag.N=3 seed=0

MaMuJoCo (Ant, 2×4 agent configuration)

python run_omad.py env_name=mamujoco/Ant env.Ant.agent_conf=2x4 seed=0

Disable WandB

python run_omad.py env_name=mpe/simple_spread wandb.activate=false

Common hyperparameter overrides

python run_omad.py \
  env_name=mpe/simple_spread \
  seed=0 \
  tot_time_steps=1e6 \
  alg.batch_size=256 \
  alg.optimizer.lr_actor=3e-4 \
  eval_freq=10000

Configuration

Key	Description	Default
env_name	Environment ID	mpe/simple_spread
seed	Random seed	0
tot_time_steps	Total training steps	1e7
eval_freq	Evaluation interval (steps)	10000
alg.batch_size	Batch size	256
alg.learning_starts	Steps before learning starts	5000
alg.actor.diff_steps	Diffusion steps	16
checkpoint.enable	Enable checkpoint saving	false

Environment-specific options live under the env section in configs/base.yaml (e.g. env.simple_spread.N).

Multi-agent wrapper

PettingZooToGymnasiumWrapper exposes a Gymnasium single-agent API:

1. Observations: concatenate observations from trainable agents (and opponents in some scenarios)
2. Actions: only trainable agents are controlled; fixed opponents use a random policy
3. Rewards: sum rewards over trainable agents

Trainable vs. fixed agent roles per MPE scenario are defined in PettingZooToGymnasiumWrapper.__init__ in run_omad.py.

Checkpoints

Enable in config:

checkpoint:
  enable: true
  path: ./checkpoints
  freq: 20000
  load_path: null  # Set to a checkpoint directory to resume training

Or via the command line:

python run_omad.py checkpoint.enable=true checkpoint.path=./checkpoints

Citation

If you use this code, please cite our paper:

Diffusing to Coordinate: Efficient Online Multi-Agent Diffusion Policies

Zhuoran Li, Hai Zhong, Xun Wang, Qingxin Xia, Lihua Zhang, Longbo Huang

@inproceedings{li2026diffusing,
  title={Diffusing to Coordinate: Efficient Online Multi-Agent Diffusion Policies},
  author={Zhuoran Li and Hai Zhong and Xun Wang and Qingxin Xia and Lihua Zhang and Longbo Huang},
  booktitle={Forty-third International Conference on Machine Learning},
  year={2026},
  url={https://openreview.net/forum?id=tHY0G4PhCR}
}

Paper: OpenReview

Acknowledgements

This repository extends and adapts ideas from several open-source projects and research works. We are grateful to their authors and maintainers.

Core algorithm and implementation

• DIME (Diffusion-Based Maximum Entropy Reinforcement Learning, ICML 2025): our training pipeline, diffusion actor, DIS sampler, and much of the JAX/Flax off-policy learning code are built on top of the public DIME implementation. We thank the DIME authors for releasing their codebase and for the diffusion-based RL framework that OMAD extends to the multi-agent setting.
• Maximum-entropy RL: automatic entropy tuning follows the formulation in Soft Actor-Critic and related discussions in the softlearning project.

Libraries and tooling

• JAX, Flax, and Optax for differentiable, JIT-compiled training.
• Stable-Baselines3 for the off-policy training loop, replay buffer interface, and callbacks.
• Gymnasium and PettingZoo for standard RL and multi-agent environment APIs.
• gymnasium-robotics for MaMuJoCo multi-agent locomotion benchmarks.
• Hydra and OmegaConf for configuration management.
• Weights & Biases for experiment logging (optional).

Additional code references

• Critic ensemble design draws inspiration from CORL and XPAG.
• Batch renormalization follows Batch Renormalization and the Flax normalization implementation.
• Distribution utilities are adapted from walk_in_the_park.

Environments

• MPE scenarios are provided via PettingZoo’s Multi Particle Environment suite.
• MaMuJoCo tasks use the multi-agent MuJoCo environments in gymnasium-robotics.

If we have inadvertently omitted a credit, please open an issue and we will update this section promptly.

License

Add a LICENSE file before publishing and verify license compatibility of third-party dependencies.