jepa-track-hidden-ball

Shell Game World Model Demo (~18M)

A small shell-game world model that predicts which cup holds a hidden ball from a single contact-sheet image.

Links

• GitHub repo: https://github.com/illegalcall/jepa-track-hidden-ball
• Upstream base repo: https://github.com/lucas-maes/le-wm
• License: MIT
• Local demo server: serve_demo_ui.py
• Single-sheet inference script: demo_jepawm_predict.py

What This Model Does

The bundled checkpoint is designed for a shell-game benchmark:

• three cups
• one hidden ball
• smooth cup swaps
• predict the final cup from the observed sequence

The repo slices a contact sheet into ordered frames, runs the checkpoint locally, and returns the final cup prediction.

Main Result

Hidden-ball balanced accuracy:

Random chance is 33.3%.

Important Caveat

This is not a plain next-step JEPA checkpoint.

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

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