YuhaiW/prism-jepa-cube

PRISM — Cube

LeWM world model and PRISM prior head for OGBench cube-single: a robot arm moves one cube to a target pose (5-D actions). These are the weights behind the Cube results in the paper (Fig. 5).

Paper · Project page · Code

Files

File	Contents
lewm_object.ckpt	LeWM world model (frozen JEPA encoder + predictor), the upstream quentinll/lewm-cube weights converted with convert_hf_checkpoint.py
prior_head_cube.pt	PRISM prior head: MLP with two 512-unit GELU layers, 0.59M parameters, β-NLL (β = 0.5), σ floor 0.05, trained on frozen features with the last frame of each episode as goal
jepa.py, module.py, prior_head.py	classes needed to load the two checkpoints
requirements.txt	tested package versions

Results

Success rate (%), mean ± std over ten seeds, 50 episodes per seed; goal 25 environment steps ahead, H = 5, B = 5, J = 30, K_elite = 30, prior scale s = 1.

Planner	K = 32	K = 64	K = 128
Vanilla MPPI	47.4 ± 6.3	48.0 ± 8.3	48.2 ± 7.2
Warm-start MPPI	56.6 ± 10.1	57.2 ± 8.0	58.8 ± 9.4
PRISM-MPPI	82.8 ± 5.5	83.0 ± 6.1	83.0 ± 4.7

Usage

git clone https://github.com/YuhaiW/prism-jepa.git && cd prism-jepa
uv venv --python 3.10 && source .venv/bin/activate
uv pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cu128
uv pip install -r requirements.txt
export STABLEWM_HOME=$HOME/.stable-wm
bash scripts/download.sh cube     # dataset → $STABLEWM_HOME/ogbench/cube_single_expert.h5, this repo's files → $STABLEWM_HOME/cube/ and checkpoints/
bash scripts/test.sh cube         # vanilla, warm-start and PRISM MPPI at K = 32, 64, 128

Citation

@misc{wang2026prism,
  title         = {PRISM: PRior-guided Imagination Sampling in world Models},
  author        = {Yuhai Wang and Jiawei Xia and Bowen Sun and Rongxuan Zhou and Xiao Hu and Yongliang Shi and Jing Du and Yang Ye},
  year          = {2026},
  eprint        = {2606.07974},
  archivePrefix = {arXiv},
  primaryClass  = {cs.RO},
  url           = {https://arxiv.org/abs/2606.07974},
}

MIT License. The world model and its code come from LeWM.