A DIAMOND-style pixel-space EDM (Elucidating the Design Space of Diffusion-Based Generative Models) implementation for world modeling in the Snake game environment. This model predicts future game frames conditioned on previous frames and player actions.
-
UNet Encoder-Decoder
- Encoder: 64×64 → 32×32 → 16×16 → 8×8 (3 downsampling blocks)
- Bottleneck: Self-attention at 8×8 resolution for global reasoning
- Decoder: 8×8 → 16×16 → 32×32 → 64×64 (3 upsampling blocks)
- Skip connections between encoder and decoder
-
Adaptive Group Normalization
- Conditions normalization on combined action + noise level embeddings
- Enables strong action conditioning throughout the network
-
EDM Preconditioning
- Preconditioned network output:
c_skip * x_noisy + c_out * network(x)
- Stable training with very few denoising steps (only 3 steps needed)
-
Frame Stacking
- 4 previous frames concatenated channel-wise (12 channels total)
- Provides temporal context for prediction
The simplest way is to directly download the model.pt and create and move it to the output/ folder.
Before playing, you will need to generate some data (1k is enough) for initialization. Follow the instructions in the GitHub repository's README.
1python scripts/play_pixel_edm.py \
2 --model_path model.pt \
3 --data_dir data/images \
4 --cfg_scale 2.0 \
5 --steps 3
