---
license: mit
tags:
  - jepa
  - world-model
  - reinforcement-learning
  - latent-planning
  - cem
library_name: pytorch
---

# relay — JEPA plans a turn-based contact-push puzzle in latent space

A 26M-parameter Joint Embedding Predictive Architecture (JEPA) trained
to play a 2-player turn-based T-block puzzle against humans. Plans
each turn via CEM in pure embedding space — no physics simulator on
the model's side at inference.

- **Code**: <https://github.com/SotoAlt/relay>
- **Live demo**: <https://sotoalt.dev/experiments/relay.html>
- **Paper recipe**: [LeWorldModel](https://arxiv.org/abs/2603.19312) (LeCun et al., 2026)
  + [DexWM](https://arxiv.org/abs/2512.13644) joint state head

## Architecture

- ViT-Tiny encoder (`vit_tiny_patch14_224`, trained from scratch, **unfrozen**, 192-dim)
- Projector MLP(192→2048→192) + BatchNorm
- AR causal Transformer predictor (AdaLN, 6 layers, 16 heads, mlp_dim=2048)
- DexWM-style joint state head MLP(192→256→256→8) trained alongside the predictor
- Loss: `MSE_pred + 0.09·SIGReg + 10·MSE_state`

## Training data

12,500 episodes × 40 steps × 5 regimes (contact_push 40%,
approach_no_push 25%, null_thrust 15%, near_t 10%, far_from_t 10%).
500K frames at 224×224 RGB with 5-frame action chunks (matches
gameplay turn structure).

## Quick start

```bash
git clone https://github.com/SotoAlt/relay.git
cd relay
pip install torch torchvision timm einops pymunk pygame opencv-python-headless \
    shapely fastapi 'uvicorn[standard]' numpy pillow h5py

huggingface-cli download sotoalt/relay relay_stage1_v9_trackE_ep02_uhead.pt \
    --local-dir checkpoints/

PYTHONPATH=. python -m world_model.infer_relay \
    --port 8800 --device cpu \
    --checkpoint-v9 checkpoints/relay_stage1_v9_trackE_ep02_uhead.pt \
    --model-execute-jepa
```

Open `http://localhost:8800/`.

## Results

Phase A joint training, λ_state=10, 5-epoch fine-tune validation:

| metric | pre-joint v9 | post-joint v9 |
|---|---|---|
| `val_pred` | 0.0071 | 0.036 |
| `val_state` | 0.032 (Phase B) | **0.007** (4.6× better) |
| pymunk calibration probe `agent_mae` | ~113 px | **~54 px (-52%)** |

Match-play (model-execute-jepa, 12 × 4 opponent policies):

| opponent | win % | net progress |
|---|---|---|
| random | 92% | +126 |
| passive | 100% | +149 |
| chase_t | 0% | -67 |
| adversarial | 25% | -10 |

## License

MIT