Initial release: LeWM + PRISM prior head for Cube

README.md

@@ -0,0 +1,78 @@
+---
+license: mit
+library_name: pytorch
+pipeline_tag: robotics
+tags:
+- robotics
+- jepa
+- world-model
+- visual-manipulation
+- prism
+- mppi
+- planning
+---
+
+# PRISM-JEPA · Cube (OGBench `cube-single`)
+
+A JEPA visual world model paired with a learned **action prior head** that
+biases an MPPI planner via closed-form Product-of-Gaussians fusion. Together
+they form the **PRISM-MPPI** pipeline from the PRISM paper, evaluated on
+OGBench's `cube-single` task.
+
+| File | Size | Role |
+|---|---|---|
+| `lewm_object.ckpt`   | 69 MB  | LeWM JEPA model (encoder + AR predictor), pickled `swm.World` object |
+| `prior_head_cube.pt` | 2.3 MB | PRISM action prior head — 3-layer MLP, trained with β-NLL (β = 0.5) on cube demos |
+
+## Headline result (K = 128, mean ± std over seeds {0, 1, 42})
+
+| Method                                  |  SR (%) |
+|-----------------------------------------|--------:|
+| Vanilla MPPI                            |    44.0 |
+| BC-only (prior mean → planner)          |    66.0 |
+| **PRISM-MPPI (s = 1)**                  | **79.3 ± 6.1** |
+
+PRISM-MPPI's only hyperparameter is the prior scale `s`; we use `s = 1` (no
+inflation). See the paper Sec. 4.4 for the full `s`-sweep.
+
+## Usage
+
+```bash
+# 1.  Clone the code repo and set up the env (see the repo README):
+git clone git@github.com:YuhaiW/prism-jepa.git
+cd prism-jepa && source .venv/bin/activate
+export STABLEWM_HOME=~/.stable-wm
+
+# 2.  Download these weights:
+hf download YuhaiW/prism-jepa-cube --local-dir ./hf_cube
+mkdir -p $STABLEWM_HOME/cube
+mv hf_cube/lewm_object.ckpt $STABLEWM_HOME/cube/
+mv hf_cube/prior_head_cube.pt .
+
+# 3.  Run PRISM-MPPI:
+python eval_prism_head.py --config-name=cube policy=cube/lewm solver=mppi \
+    +head.injection_mode=pog +head.sigma_scale=1.0 \
+    +head.ckpt=prior_head_cube.pt eval.num_eval=50
+
+# Vanilla MPPI baseline (no prior):
+python eval_prism_head.py --config-name=cube policy=cube/lewm solver=mppi \
+    +head.injection_mode=none eval.num_eval=50
+```
+
+## Training (summary)
+
+- **`lewm_object.ckpt`** — LeWM JEPA trained from scratch on OGBench
+  `cube_single_expert` (upstream recipe in `train.py`).
+- **`prior_head_cube.pt`** — 3-layer MLP mapping JEPA embedding `z_t` to a
+  per-coordinate Gaussian `(μ_p, σ_p)` over the next action block. β-NLL loss
+  (β = 0.5), `sigma_floor = 0.05`, Adam, 50 epochs. See `train_prior_head.py`.
+
+## Citation
+
+_BibTeX to be added._
+
+## License
+
+MIT for the PRISM head + integration code. The vendored LeWM code in
+`prism-jepa` inherits its upstream license — see
+[le-wm](https://github.com/lucas-maes/le-wm).

lewm_object.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82d37a9d9338d8c23005017ab5c1ff91c8b5e3fd51fafbd620af8457c381d125
+size 72344949

prior_head_cube.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0bbfacb047d7ea68370d07a56185099807cc1a9536034fbe53cdbfb3f6d78dec
+size 2358901