lepong — KAZ JEPA world model
A ~13M-parameter Joint Embedding Predictive Architecture (JEPA) that watches
128×128 pixels of PettingZoo Knights-Archers-Zombies (knights_archers_zombies_v10)
and predicts a 28-dim game state (archer/knight positions + up to 10 zombies) from a
frozen embedding. Pixels in, state out — the model reads no RAM at inference time.
Checkpoints
Multi-step rollout sweep, 100 epochs each on 100K frames (kaz_ma_128x128.lance):
| File | Rollout steps | Warmup | Notes |
|---|---|---|---|
kaz_R1.pt | 1 | – | single-step baseline |
kaz_R3.pt | 3 | 10 | discount 0.9 |
kaz_R5.pt | 5 | 15 | discount 0.9 |
Each checkpoint is self-describing: it carries game, state_dim=28,
state_names, state_mean/std, num_actions=6, and embed_dim.
Usage
import torch
ckpt = torch.load("kaz_R5.pt", map_location="cpu", weights_only=False)
print(ckpt["game"], ckpt["state_dim"], ckpt["state_names"])
Load and drive it with the lepong repo's unified player:
python -m server.play --game kaz --checkpoint checkpoints/kaz_R5.pt
License
MIT