README.md

---
license: mit
tags:
- minecraft
- terrain-generation
- vae
- onnx
- voxel
- film
- transfer-learning
- games
datasets:
- ghosteau/minecraft-chunks
pipeline_tag: other
---

# STEVE-1 — Style-conditioned Terrain Encoder VAE

STEVE-1 generates full Minecraft chunks (16 x 384 x 16 voxels, bedrock to build limit) from **latent noise + a biome ID + a style ID** in a **single ONNX forward pass** — fast enough to run inside a live PaperMC server. It is also a **pretrained base for transfer learning**: anyone can fine-tune a new terrain *style* onto it from a folder of their own world's chunks, in minutes, training ~1.8% of the weights, without touching the base.

- **Training pipeline, plugin, and fine-tuning notebook:** [github.com/ghosteau/generative-terrain](https://github.com/ghosteau/generative-terrain)
- **Training data:** [ghosteau/minecraft-chunks](https://huggingface.co/datasets/ghosteau/minecraft-chunks) (2,040 chunks, ~200M voxels, 22 biomes)

## Architecture

A conditional VAE built around one hard constraint: in-game inference must be a single fast graph call (no diffusion, no autoregression).

| Component | What it does |
|---|---|
| Spatial latent `z` `[8, 4, 16, 4]` | A coarse 3D latent grid instead of one broadcast vector, so terrain varies from place to place (hills, valleys, caves) |
| Heightmap stage | The decoder first predicts a per-column surface height, then feeds a signed-distance-to-surface channel to the voxel head — the model commits to a surface instead of hedging |
| 3-level 3D U-Net (~4.6M params) | Chunk-scale receptive field: at the bottleneck the chunk is 2 x 48 x 2, so horizontal context is global |
| FiLM style conditioning | Every U-Net block and the heightmap head is modulated (per-channel scale/shift) by a style embedding. Style 0 = vanilla; 8 slots reserved for custom styles. This is the fine-tuning mechanism |
| 27 block groups | The model predicts semantic groups (`STONE`, `GRASS`, `WATER`, ores, ...); the plugin expands them to concrete blocks with local context |

Conditioning uses **only what exists before terrain does** (position, biome, noise, style). It deliberately never sees neighbour blocks, light, or surface flags — the data-leakage mistake that sank earlier versions of this project.

Training: EMA weights (exported), mixed precision, warmup + cosine LR, KL annealing with free bits, biome-balanced sampling, and biome dropout so a reserved `UNKNOWN` biome row learns generic terrain (that is what makes fine-tuning work on custom worlds with unrecognised biomes).

## Files

| File | Purpose |
|---|---|
| `terrain_vae_decoder.onnx` | The generator: `(z, biome_id, style_id) -> logits [1, 27, 16, 384, 16]` |
| `block_group_mapping.json` | Class id -> block group name |
| `biome_id_mapping.json` | Biome name -> id (includes `UNKNOWN`) |
| `style_mapping.json` | Style name -> id (`BASE` = 0) |
| `base_v1/` | The versioned base checkpoint for fine-tuning: PyTorch weights + a manifest with the exact architecture and mappings |

## Use it in a Minecraft server

Drop `terrain_vae_decoder.onnx` and the three JSON mappings into `<server>/plugins/GenerativeTerrain/` (plugin jar from the [GitHub repo](https://github.com/ghosteau/generative-terrain)), then in game:

```
/generateterrain                  base style, current chunk
/generateterrain style <name>     a fine-tuned custom style
```

## Use it from Python

```python
import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download

sess = ort.InferenceSession(hf_hub_download("ghosteau/STEVE-1", "terrain_vae_decoder.onnx"))

logits = sess.run(["logits"], {
    "z":        np.random.randn(1, 8, 4, 16, 4).astype(np.float32),  # the variety
    "biome_id": np.array([15], dtype=np.int64),                      # PLAINS (see biome_id_mapping.json)
    "style_id": np.zeros(1, dtype=np.int64),                         # 0 = BASE
})[0]

chunk = logits.argmax(1)[0]        # [16, 384, 16] block-group ids, Y index 0 = world Y -64
```

Sample a fresh `z` per chunk; scale it (`z * 0.8`) for tamer terrain, (`z * 1.2`) for wilder.

## Fine-tune it on your own terrain

This is the headline feature. The `base_v1/` directory is a self-contained transfer-learning checkpoint: weights plus a manifest that the pipeline's `fine_tune()` uses to rebuild the exact architecture, so version mismatches are impossible.

1. Export 50+ chunks of your world with the plugin: `/grabchunkarea 5`
2. Clone the [GitHub repo](https://github.com/ghosteau/generative-terrain), open `ml/notebooks/FineTune.ipynb`
3. Point it at `base_v1/`, your CSV folder, and a style name; run it

Fine-tuning trains the new style vector and the FiLM projections (~84k of 4.6M parameters) while everything else stays frozen — a unit test in the repo asserts the base weights come out bit-identical. The exported ONNX then carries both the vanilla style and yours, selectable per command. Up to 8 custom styles fit in one model.

## Training results

Vertical block-distribution profile, real vs generated (the terrain-plausibility check — bedrock floor, deepslate/stone body, thin surface band, air above):

![Real vs generated vertical profiles](assets/profiles.png)

Training curves (EMA-validated loss):

![Training curves](assets/vae_curves.png)

*Note: this was trained on my personal NVIDIA GeForce RTX 4070.*

## Limitations

- **Chunks generate independently** — edges between adjacent generated chunks do not line up yet. Cross-chunk edge conditioning is the next planned lever.
- **Rare biomes are undertrained**: the dataset has 1-5 chunks of TAIGA, FLOWER_FOREST, and SWAMP; those styles lean on the balanced sampler and will look generic.
- Trees and vegetation are coarse (single log/leaf groups expanded with heuristics by the plugin).
- Ores are placed procedurally after generation (as vanilla does), not modelled.

## License and attribution

MIT. Not affiliated with or endorsed by Mojang or Microsoft; "Minecraft" is a trademark of Mojang Synergies AB. The name STEVE is, of course, for Steve.