A Surrogate Model is All You Need for Coordinated Spatial Reuse

We propose a simple yet effective approach to coordinated spatial reuse (Co-SR) scheduling in IEEE 802.11bn (Wi-Fi 8) networks: draw N random configurations, score each with a surrogate model, and select the best ones. No reinforcement learning, no generative model, no iterative optimization — just a single forward pass through a lightweight graph neural network.

How It Works

1. Random sampling — Draw N random Co-SR configurations (active APs, STA selection, MCS, transmit power) and encode each as a graph
2. Surrogate scoring — Score all candidates in a single batched forward pass using a GNN with a Mixture Density Network (MDN) head
3. Selection — Pick the best configurations (top-k or coverage-based) and apply them in round-robin

Installation

Requires Python >= 3.12. We recommend using uv:

git clone https://github.com/ml4wifi-devs/mapc-surrogate.git
cd mapc-surrogate
uv venv
source .venv/bin/activate
uv sync

For GPU support (CUDA 12):

uv sync --extra gpu

Usage

Dataset Generation

Generate training and validation datasets from simulated Wi-Fi scenarios:

cd mapc_surrogate
python dataset.py

This creates LZ4-compressed datasets in mapc_surrogate/datasets/.

Training

Train the surrogate model (uses Hydra for configuration):

cd mapc_surrogate
python training.py

Override defaults:

python training.py model.dim=128 optimizer.learning_rate=1e-3 train.n_steps=20000

Training logs are sent to Weights & Biases. Checkpoints are saved to runs/.

Project Structure

mapc_surrogate/
├── attributes.py        # Edge attribute enums (MCS, tx_power, selected) with JAX PyTree support
├── graphs.py            # Graph construction: WiFi topology → NetworkX → Jraph
├── model.py             # SurrogateModel: edge embedding → GNN + Transformer → MDN head
├── dataset.py           # Dataset generation from simulated scenarios
├── training.py          # Training loop with MDN loss, Orbax checkpointing, W&B logging
├── sim.py               # Surrogate inference, configuration selection strategies
├── baselines.py         # H-MAB baseline (hierarchical UCB)
├── datasets/            # Pre-generated training/validation data
└── configs/             # Hydra YAML configs (model, optimizer, training)

Citation

@article{wojnar2026surrogate,
  title={A Surrogate Model is All You Need for Coordinated Spatial Reuse},
  author={Wojnar, Maksymilian},
  year={2026}
}