Metadata-Version: 2.4
Name: bench-mfg-suite
Version: 0.1.0
Summary: A benchmark suite for Mean Field Game algorithms
License: MIT
Keywords: machine-learning,mean-field-games,optimization,reinforcement-learning
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.10
Requires-Dist: cma>=3.2.0
Requires-Dist: colorlog>=6.6.0
Requires-Dist: hydra-core>=1.3.0
Requires-Dist: jax>=0.4.13
Requires-Dist: matplotlib>=3.5.0
Requires-Dist: numpy>=1.21.0
Requires-Dist: pyyaml>=6.0.0
Requires-Dist: rich>=13.0.0
Requires-Dist: scipy>=1.7.0
Requires-Dist: tqdm>=4.62.0
Provides-Extra: cuda
Requires-Dist: jax[cuda12]>=0.4.13; (sys_platform == 'linux') and extra == 'cuda'
Provides-Extra: cuda12
Requires-Dist: jax[cuda12]>=0.4.13; (sys_platform == 'linux') and extra == 'cuda12'
Provides-Extra: cuda12-local
Requires-Dist: jax[cuda12-local]>=0.4.13; (sys_platform == 'linux') and extra == 'cuda12-local'
Provides-Extra: dev
Requires-Dist: black>=22.0.0; extra == 'dev'
Requires-Dist: pre-commit>=2.20.0; extra == 'dev'
Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
Requires-Dist: pytest>=7.0.0; extra == 'dev'
Requires-Dist: ruff>=0.1.0; extra == 'dev'
Provides-Extra: docs
Requires-Dist: sphinx-rtd-theme>=1.0.0; extra == 'docs'
Requires-Dist: sphinx>=5.0.0; extra == 'docs'
Provides-Extra: jupyter
Requires-Dist: ipykernel>=6.0.0; extra == 'jupyter'
Requires-Dist: ipywidgets>=8.0.0; extra == 'jupyter'
Requires-Dist: jupyter>=1.0.0; extra == 'jupyter'
Provides-Extra: wandb
Requires-Dist: wandb>=0.23.0; extra == 'wandb'
Description-Content-Type: text/markdown

<div align="center">

<img src="favicon_v3.svg" width="90" alt="BenchMFG icon"/>

# BenchMFG

Benchmark suite for Mean Field Game algorithms.

[![Python](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
[![JAX](https://img.shields.io/badge/JAX-accelerated-salmon.svg)](https://github.com/google/jax)
[![Hydra](https://img.shields.io/badge/Hydra-config-89b8cd.svg)](https://hydra.cc)
[![uv](https://img.shields.io/badge/uv-package%20manager-purple.svg)](https://github.com/astral-sh/uv)

[![License: MIT](https://img.shields.io/badge/License-MIT-red.svg)](https://opensource.org/licenses/MIT)
[![Unit tests](https://img.shields.io/badge/unit%20tests-passed-brightgreen.svg)](https://docs.pytest.org/)
[![ruff](https://img.shields.io/badge/ruff-%E2%9A%A1-gold.svg)](https://github.com/astral-sh/ruff)
[![pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen.svg)](https://pre-commit.com/)

</div>

## Contents

- [Install](#install)
- [Quick Start](#quick-start)
- [Registered Configs: envs and algos](#registered-configs)
- [Sweep](#sweep)
- [Python API](#python-api)
- [Outputs And Plots](#outputs-and-plots)
- [Extending BenchMFG: create new envs and algos](docs/EXTENDING.md)
- [MF-Garnet: scaling laws and robustness](docs/MFG_GARNET.md)

## Install

```bash
uv add bench-mfg-suite
# or
pip install bench-mfg-suite
```

For local development:

```bash
uv sync --extra dev
```

CUDA is optional. The default install uses CPU-compatible JAX.

```bash
# Linux/NVIDIA, pip-managed CUDA runtime components:
uv add "bench-mfg-suite[cuda12]"
pip install "bench-mfg-suite[cuda12]"

# Linux/NVIDIA, local CUDA installation:
pip install "bench-mfg-suite[cuda12-local]"
```

If GPU initialization fails, check `nvidia-smi` and the official JAX install matrix:
https://docs.jax.dev/en/latest/installation.html

## Quick Start

List registered configs:

```bash
benchmfg env list
benchmfg algo list
```

Run one experiment:

```bash
benchmfg train algorithm=pso environment=kinetic_congestion device=cpu
```

## Registered Configs

Environments:
`contraction_game`, `four_rooms_obstacles`, `kinetic_congestion`, `lasry_lions_chain`, `mf_garnet`, `multiple_equilibria`, `no_interaction_game`, `potential_game2d`, `rock_paper_scissors`, `sis_epidemic`.

Algorithms:
`damped_fixed_point`, `omd`, `pi`, `pso`.

Use `benchmfg env list` and `benchmfg algo list` for the installed package’s authoritative list.

## Sweep

Run a sweep:

```bash
benchmfg sweep \
  algorithm=omd \
  environment=lasry_lions_chain \
  experiment.name=omd_sweep \
  experiment.random_seed=42,10,111,1032 \
  algorithm.omd.learning_rate=0.5,0.05,0.005 \
  algorithm.omd.temperature=0.2,0.5,0.8
```

## Python API

```python
import benchmfg

cfg = benchmfg.load_config(["algorithm=omd", "environment=lasry_lions_chain"])
environment, initial_policy = benchmfg.make_environment(cfg)
solver = benchmfg.make_solver(
    cfg,
    environment=environment,
    initial_policy=initial_policy,
)
```

## Outputs And Plots

Runs write artifacts under:

```text
outputs/<Env>/<Algorithm>/seed_<seed>/<Experiment>/<run_id>/
```

Important files: `exploitabilities.npz`, `final_mean_field.npz`, `final_policy.npz`,
`metrics.npz`, `config.yaml`.

Plot commands:

```bash
benchmfg plot single-run <run_dir>
benchmfg plot sweep <environment> <algorithm>
benchmfg plot compare <environment>
```

Plot discovery defaults:

- `single-run <run_dir>` plots exactly that timestamped run.
- `sweep <environment> <algorithm>` scans `outputs/` by default. For each seed and
  hyperparameter version, it selects the latest timestamped run containing
  `exploitabilities.npz`.
- `compare <environment>` reads the `results/<environment>/<algorithm>/best_model.yaml`
  files written by `plot sweep`; rerun `plot sweep` first if new runs were added.
- Use `--outputs-dir <path>` on sweep/compare commands when artifacts are not under
  `outputs/`.

## Repository Layout

```text
src/benchmfg/
├── config/      # packaged Hydra configs
├── envs/        # MFG environments
├── learner/     # solvers
├── utility/     # training, saving, plotting helpers
├── cli.py       # benchmfg command
└── train.py     # Hydra train entrypoint
```

See [EXPERIMENTS.md](EXPERIMENTS.md) for batch-run scripts.
