LC-MAPF: Learning to Communicate Locally for Large-Scale Multi-Agent Pathfinding

---

The repository contains the following parts:

• example.py — run a quick inference example
• benchmark.py — evaluate on the full POGEMA benchmark
• train.py — train or fine-tune the model
• dataset/ — dataloader scripts
• configs/ — training config files (config_default.py lists all parameters with defaults)
• eval_configs/ — POGEMA benchmark scenario configs, required by benchmark.py

🔧 Local Installation

Use uv for local installation:

uv venv --python 3.10
source .venv/bin/activate
uv pip install -r docker/requirements.txt

🐋 Server-side Installation

For training, use Docker:

cd docker && sh build.sh

---

🔍 Running an example

Run example.py from the repo root — weights are downloaded from Hugging Face automatically on first run:

python3 example.py --map_name validation-mazes-seed-000 --num_agents 32

Alternatively, download weights manually from the Hugging Face model page and place them in a weights/ folder at the repo root:

weights/
└── LC-MAPF-3M.pt

Then pass --repo_id '' to disable auto-download and use the local file.

Optional arguments: --map_name, --device, --num_agents, --seed, --max_episode_steps, --ckpt, --ckpt_file, --repo_id, --show_map_names. The --map_name argument selects a map from those available in the eval_configs folder; use --show_map_names to list all options. Examples from each set: validation-random-seed-000, validation-mazes-seed-000, wfi_warehouse, Berlin_1_256_00, puzzle-00.

GPU is recommended; smaller models run on CPU. For Apple Silicon use --device mps.

An SVG animation of the solution is saved to the svg/ folder alongside the run metrics.

📚 Dataset

The training dataset is hosted on Hugging Face. Download a minimal subset (1 file per domain) for quick local testing:

python dataset/download_dataset.py --minimal

This creates the following structure:

dataset/train
├── house
│   └── part_0_0.arrow
├── mazes
│   └── part_0_0.arrow
└── random
    └── part_0_0.arrow

Download the full dataset:

python dataset/download_dataset.py

---

📊 Running evaluation

Run benchmark.py to evaluate on all POGEMA benchmark scenarios:

python3 benchmark.py

Results are stored in eval_results/. Average success rate tables are printed to the console and can optionally be logged to wandb.

🚀 Training

All training commands must be run from the repo root. A config file sets model and training hyperparameters — see configs/config_default.py for all available parameters. The script supports multi-GPU training via DDP; adjust nproc_per_node to match the number of available GPUs.

Single GPU:

torchrun --standalone --nproc_per_node=1 train.py configs/config_3M.py

Multi-GPU (e.g. 2 GPUs):

torchrun --standalone --nproc_per_node=2 train.py configs/config_3M.py