Disclaimer: Fan/research project; not affiliated with Nintendo, Intelligent Systems, or AWBW.
This repo contains code for diffusion based map generators for the Advance Wars game. Advance Wars is a turn based strategy game where the player objective is to defeat the enemy army using military units to capture the opponent's headquarters. Although there are different versions of this game, this repo will be focusing on maps from some of the older games like advance wars days of ruin. Thank you to the dev team at the AWBW (https://awbw.amarriner.com/) discord channel for providing me a dataset to work with.
Example maps (left side are land maps, right side are naval maps)
The goal of this project is to build a map generator using diffusion models to create diverse and novel maps based on an existing dataset. The core models trained in this repo are map generators however there are few others that serve other purposes such as a terrain2vec model to create a learned representation for latent diffusion and a headquarters placement model that learns how to place headquarters with the long term goal of it learning to place any type of buildings like factories, ports, etc. in a way that is balanced.
The maps dataset is simply a 2d array of tile ids that map to some tile image (see data/tiles/full/ for example images). The array size can range anywhere from 8x8 to 50x50. To simplify the problem, in most of my experiments there is upsampling or downsampling to a target size such as 16x16 or 32x32 and also augmenting the maps such as flipping and rotations. Most of the experiments are trained on a small dataset of a couple maps to build a proof of concept.
- v4 Latent Diffusion: 10x faster training, 200x smaller model than v2, with better diversity
- Terrain2Vec: Learned embeddings capture terrain relationships without manual feature engineering
- Text Conditioning: Successfully generates novel map types not in training data
- Spatial Conditioning: Enables controlled generation with masks
- Comprehensive Evaluation: All models include metrics (tpkl-div, memorization checks, etc.) - see notebooks and
src/training/evaluation.py
The purpose of these models is to take a set of example maps and create new diverse and novel maps. I start with unconditioned diffusion before moving onto text conditioned diffusion and finally latent diffusion.
Note: The v1 diffusion model is not uploaded as it was mostly used for learning about diffusion models and there were no meaningful results to share.
Adapting the approach from the Lee paper [1] which uses unconditioned diffusion models to generate levels in Super Mario Bros. The core approach uses UNet based diffusion with one hot encoded tiles, taking a single level and augmenting the dataset by using a sliding window to upsample.
When applied to the advance wars map it essentially memorized the dataset and would generate identical maps to the training dataset. This is likely due to the sliding windows overlapping significantly meaning it is essentially training over a similar set of maps repeatedly leading to overfitting. A single training run would take approximately 2 and a half hours and the model size is approximately 1gb.
16x16 examples part way through training, will eventually overfit and will generate examples that are identical to the training dataset.
Key files:
experiments/v2_diffusion/
This approach builds on the previous version but experiments with text and spatial conditioning using a simple dataset of a few land maps and naval maps and trying to get it to generate new blended maps like "land+naval" despite not being trained on it.
It was able to successfully generate maps that looked blended however it appears too random and fragmented to be playable. Also the local patterns lack coherence. This problem can be partly solved using spatial conditioning with masks where part of the noise map is denoised using "land" and the rest "naval". Below are examples for an island mask where you can see clearer boundaries.
Blended examples using text conditioning
Spatial conditioning using island mask
Key files:
experiments/v3_text_diffusion/
This approach is based on the Dai paper [2] which uses latent diffusion and a single example to generate new variations of the map it was trained by random cropping, also significantly reduces memorization by removing self attention layer and keeping the UNet flat.
Significant improvements over v2/v3:
- Model size: ~5MB (vs ~1GB for v2)
- Training time: 15-20 minutes (vs 2.5 hours)
- Diversity: Better variation, less memorization (validated via tpkl-div and other evals)
- Quality: More coherent local patterns - roads connect, terrain transitions are natural
Naval maps show particularly strong results with clean water boundaries and natural island formations. Land maps have some local pattern coherence but there is no pattern at the global level.
Generated examples (32x32):
Key files:
experiments/v4_latent_diffusion/v4_latent_diffusion.ipynbexperiments/v4_latent_diffusion/train_v4_latent_diffusion.py
A model that converts tile ids to embeddings based on its neighbors. This is used for the latent diffusion model. This is done using a simple skipgram model.
Successfully learns terrain embeddings and places terrains that are similar in neighborhood closer to each other in embedding space. Model is evaluated by comparing the tile co-occurrence in the original grids vs the neighbors in the embedding space and also looking at top k similar tiles for a specific tile for example plain, mountain and tree tiles are close in the embedding space. It's also shown to work well with the v4 latent diffusion model.
Key files:
experiments/v1_terrain2vec/v1_terrain2vec.ipynb
A model that takes a 2D array of tile ids and places the HQ in a balanced way. It uses a basic UNet architecture with 3x downsampling conv layers, bottleneck and 3x upsampling conv layers. It's trained on the original maps split into a terrain map and another a HQ map.
Model shows that it has learned symmetry in HQ placement. Although it occasionally places HQs on invalid tiles (water, mountains) where they cannot be built in the actual game.
Key files:
experiments/v1_hq_placer/v1_hq_placer.ipynb
experiments/
v1_terrain2vec/ # Terrain embedding model (skip-gram)
v2_diffusion/ # Unconditioned diffusion baseline
v3_text_diffusion/ # Text + spatial conditioning
v4_latent_diffusion/ # Latent diffusion (current best)
v1_hq_placer/ # HQ placement model
src/
models/ # UNet, diffusion, embeddings
training/ # Training loops, evaluation
data/ # Data loading, preprocessing
visualization/ # Map rendering
Requirements:
- Python 3.12+
- PyTorch 2.9+
- CUDA-capable GPU (recommended)
Installation: pip install -r requirements.txt
Usage:
Run any of the notebooks under experiments/ with a Jupyter notebook server pointing to your virtualenv.
This code is available under CC BY-NC 4.0 - you can use and modify it for non-commercial purposes as long as you give credit.
[1] H. J. Lee and E. Simo-Serra, "Using Unconditional Diffusion Models in Level Generation for Super Mario Bros," 2023 18th International Conference on Machine Vision and Applications (MVA), Hamamatsu, Japan, 2023, pp. 1-5, doi: 10.23919/MVA57639.2023.10215856. keywords: {Video games;Machine vision;Games} https://ieeexplore.ieee.org/document/10215856
[2] Dai, S., Zhu, X., Li, N., Dai, T., & Wang, Z. (2024). Procedural Level Generation with Diffusion Models from a Single Example. Proceedings of the AAAI Conference on Artificial Intelligence, 38(9), 10021-10029. https://doi.org/10.1609/aaai.v38i9.28865