The working shell-game variant uses explicit hidden-state supervision during training. Pure next-step JEPA-style prediction did not solve this benchmark.

Quick Start

bash

git clone https://github.com/illegalcall/jepa-track-hidden-ball.git
cd jepa-track-hidden-ball
python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

PYTHONPATH=local_inference_assets python3 demo_jepawm_predict.py \
  --checkpoint /path/to/lewm_auxonly_123456_h12_epoch_12_object.ckpt \
  --sheet demo_cases/case_1/sheet.png \
  --history-size 12 \
  --device cpu \
  --output result.json

Local Demo UI

python3 serve_demo_ui.py --host 127.0.0.1 --port 8123

Then open:

http://127.0.0.1:8123/demo_ui/

Model Size

This checkpoint has 18,048,683 trainable parameters.

Files In This Release

lewm_auxonly_123456_h12_epoch_12_object.ckpt
local-ui-overview.png
local-ui-output-card.png

Limitations

this checkpoint is specialized to the shell-game benchmark in the repo
it is not a general-purpose vision model
retraining from scratch still depends on upstream LeWM / stable-worldmodel code

illegalcall/jepa-track-hidden-ball

Shell Game World Model Demo (~18M)

Links

What This Model Does

Main Result

Important Caveat

Quick Start

Local Demo UI

Model Size

Files In This Release

Limitations

illegalcall/
jepa-track-hidden-ball