YuhaiW/prism-jepa-pusht

PRISM — PushT

LeWM world model and PRISM prior head for PushT: a planar pusher aligns a T-block with a target pose (2-D actions). These are the weights behind the PushT 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-pusht weights converted with convert_hf_checkpoint.py
prior_head_pusht.pt	PRISM prior head: MLP with two 512-unit GELU layers, 0.51M 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	56.0 ± 5.7	58.6 ± 6.2	58.8 ± 6.9
Warm-start MPPI	70.6 ± 9.1	71.0 ± 5.9	68.8 ± 6.8
PRISM-MPPI	80.6 ± 5.6	82.8 ± 6.9	85.4 ± 7.1

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 pusht     # dataset → $STABLEWM_HOME/pusht_expert_train.h5, this repo's files → $STABLEWM_HOME/pusht/ and checkpoints/
bash scripts/test.sh pusht         # 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.