README.md

---
license: apache-2.0
tags:
- robotics
- franka
- jepa
- lewm
- prism-mppi
- world-model
- visuomotor
- pusht
- multi-step-supervision
library_name: pytorch
pipeline_tag: robotics
---

# PRISM-LeWM for Franka FR3 Planar PushT — v4 (multi-step supervision)

**Research preview** of a LeWM-style JEPA world model trained on **200
real-robot demonstrations** (51 k frames, Franka FR3, 10 Hz) using
**multi-step autoregressive predictor supervision**. This is the first
public release of LeWM trained with this scheme.

**Status: deploy-marginal.** Cost-surface discriminativity improved **2.17×**
over LeWM-default training (CV 0.13 → 0.29), and the Test A in-distribution
deploy gate is the closest to PASS we've achieved (3.93× vs PASS ≤ 2×).
Real-robot SR is not yet measured. **Use as a baseline / starting point
for further work, not as a production-ready policy.**

## TL;DR

| | |
|---|---|
| **Source dataset** | [`Rongxuan-Zhou/pusht_lewm_fr3`](https://huggingface.co/datasets/Rongxuan-Zhou/pusht_lewm_fr3) (200 demos, 51 492 frames @ 10 Hz, Franka FR3) |
| **Action space** | **2-dim (dx, dy)** — planar delta-EE in meters |
| **Encoder** | ViT-tiny from scratch (5.5 M params) |
| **Predictor** | LeWM-default ARPredictor (10.79 M params, depth 6 / heads 16 / mlp 2048) |
| **Training novelty** | **3-step autoregressive supervision** (predictor sees own predictions during training; TD-MPC-style stop-gradient between steps). LeWM-default is 1-step teacher-forced. |
| **Prior head** | PriorHead MLP, val MSE drop 14.7 % (marginal vs 15 % gate) |
| **Cost-surface CV @ H=5** | **0.29** vs sim PushT 0.13 vs LeWM-default-on-our-data 0.13 |
| **Bundle size** | ~ 80 MB |

## What's in here

| File | Size | Role |
|---|---|---|
| `lewm_pusht_lewm_fr3_v4b.ckpt` | 72 MB | LeWM world model (JEPA encoder + multi-step-trained AR predictor) |
| `prior_head_pusht_lewm_fr3_v4b.pt` | 2 MB | PRISM prior head + StandardScaler |
| `franka_pusht_lewm_inference.py` | 13 KB | Self-contained PRISM-MPPI / vanilla-MPPI inference loop |
| `jepa.py`, `module.py` | ~ 14 KB | Model classes (required to unpickle) |
| `prior_head.py` | 2.4 KB | `PriorHead` class |
| `requirements.txt` | 0.4 KB | Runtime deps |
| `README.md` | this file | Usage + deploy guidance |

## Installation

```bash
pip install huggingface_hub
python -c "from huggingface_hub import snapshot_download; \
    snapshot_download(repo_id='YuhaiW/prism-jepa-franka-pusht-v4', \
                      local_dir='./franka_pusht_v4_bundle')"
cd franka_pusht_v4_bundle/
pip install -r requirements.txt
```

PyTorch ≥ 2.1 + a CUDA GPU recommended. CPU works but `plan()` is
~ 10× slower.

## Quick start — running PRISM-MPPI

```python
import numpy as np
from franka_pusht_lewm_inference import (
    PrismMPPIInference,
    pad_2d_to_6d_franka,
)

# PRISM-MPPI (uses the prior head to bias MPPI's initial action distribution)
planner = PrismMPPIInference(
    lewm_ckpt  = "lewm_pusht_lewm_fr3_v4b.ckpt",
    prior_ckpt = "prior_head_pusht_lewm_fr3_v4b.pt",
    use_prism  = True,
    device     = "cuda",
)

# In your control loop:
obs_uint8  = camera.read_d455_agent_view()    # (224, 224, 3) uint8 RGB
goal_uint8 = goal_image                        # (224, 224, 3) uint8 RGB
actions_2d = planner.plan(obs_uint8, goal_uint8)
# → (5, 2) array of (dx, dy) deltas in meters

for a2d in actions_2d:
    a6d = pad_2d_to_6d_franka(a2d)             # (6,)  dx, dy, 0, 0, 0, 0
    robot.send_delta_ee(a6d)
```

For vanilla MPPI (no prior, useful for A/B comparison):

```python
planner = PrismMPPIInference(..., use_prism=False, ...)
```

## How MPPI works in this code (~ 4 sentences)

For each `plan()` call:

1. **Encode** the current `obs` and `goal` images via the JEPA encoder
   → latent vectors `z_t`, `z_g` (each 192-dim).
2. **Init MPPI distribution** as N(0, var_scale²) in StandardScaler-normalized
   action space. If `use_prism=True`, **PoG-fuse** this with the prior
   head's (μ_p, σ_p) → biased starting distribution.
3. **MPPI loop** (30 iterations): sample K=128 candidate action sequences
   from the current N(μ, σ²); for each, autoregressively roll out the
   predictor for H=5 plan-steps to get a predicted `z_5`; compute
   `cost = ‖z_5 - z_g‖²`; reweight candidates by `softmax(-β · cost)`;
   update `μ` to the importance-weighted mean. **σ is frozen across
   iterations** — this is the PRISM-MPPI signature.
4. **Return** the first `A_block = 5` env-step actions of the converged μ,
   denormalised back to raw env units (meters).

If you'd like to tweak: `H`, `K`, `n_iters`, `var_scale`, `temperature`,
`prior_sigma_scale`, `history_size` are all constructor args of
`PrismMPPIInference`. Defaults match the LeWM-sim paper conventions.

## Franka FR3 deployment notes

### Action format

The model outputs **2-dim** actions `(dx, dy)` in meters per 10 Hz env-step.
The 4 other dims (`dz, drx, dry, drz`) that some Franka teleop pipelines
record are **floating-point jitter** in our dataset and were dropped at
the data level — see §"Why drop 4 dims" below for evidence.

For Franka FR3 controllers that expect 6-dim `(dx, dy, dz, drx, dry, drz)`,
pad with zeros:

```python
a2d = planner.plan(obs, goal)[0]
a6d = pad_2d_to_6d_franka(a2d)              # (6,) = [dx, dy, 0, 0, 0, 0]
robot.send_delta_ee(a6d)
```

Training distribution per env-step:
- `dx` range ±0.025 m, std 0.005 m
- `dy` range ±0.030 m, std 0.006 m
- Frame rate: 10 Hz → average speed ~ 5 cm/s

### Suggested safety clamps

Even a stable model occasionally outputs actions a few × the training
range. Clamp before sending:

```python
ACTION_CLAMP_2D = np.array([0.030, 0.035])   # ~1.2× max(|training|)
def clamp_safety(a):
    return np.clip(a, -ACTION_CLAMP_2D, +ACTION_CLAMP_2D)

actions_2d = clamp_safety(planner.plan(obs, goal))
```

Also on the hardware side:
- Workspace bounding box: roughly `x ∈ [0.35, 0.81]`, `y ∈ [-0.30, 0.37]`,
  `z = 0.099 ± 0.005` (planar-lock height) in robot base frame
- Operator e-stop physically reachable
- First N trials at 0.5× velocity scaling
- Time-based MAX_STEPS cap (≈ 5 s at 10 Hz)

### Receding-horizon control loop

```python
CONTROL_DT = 0.1           # 10 Hz (matches training)
N_EXEC = 5                  # = A_block; execute then replan
MAX_STEPS = 50              # ≈ 5 s safety cap

step = 0
while step < MAX_STEPS:
    obs = preprocess_to_224(camera.read())
    if task_complete(obs, goal_uint8):
        break
    actions_2d = clamp_safety(planner.plan(obs, goal_uint8))
    for a2d in actions_2d[:N_EXEC]:
        a6d = pad_2d_to_6d_franka(a2d)
        robot.send_delta_ee(a6d)
        time.sleep(CONTROL_DT)
        step += 1
        if step >= MAX_STEPS: break
robot.move_to_home()
```

## Diagnostic numbers — full honesty

We evaluate three concentric tests:

```
                  encoder         predictor       MPPI usability
                  health          quality         (the deploy gate)
                  ─────────       ──────────      ───────────────
  Test           LOEO R²          pred/id ratio   cost-surface CV
  Sim ref        n/a              0.79 (h=1)      0.13
  v3 baseline    +0.03            0.97 (h=1)      0.13
  v4b (this)     (not measured)   (not measured)  ★ 0.29 ★
```

### Cost-surface CV in detail

Tested on 12 random `(z_t, z_g)` training pairs, K=256 random action
candidates each. The CV (`std(cost) / mean(cost)`) measures how
discriminative the cost surface is across action candidates:

```
                CV mean   CV max    GT rank pct
  v3 baseline   0.132     0.419     33 %
  v4b (this)    0.286     0.712     44 %
  sim PushT     0.133     0.367     59 %
  PASS gate     ≥ 0.30    —         ≤ 25 %
```

v4b's CV (0.29) is **2.17× the LeWM-default training** (0.13) on the same
data — the largest movement we've seen from any single technique. It
exceeds the sim PushT reference (0.13). **But it falls 0.014 short of
our 0.30 PASS threshold**, and `GT_rank_pct` worsened (33 → 44 %),
meaning the predictor's cost ranks ground-truth-like actions only
slightly better than random.

### Test A — in-distribution sanity (deploy gate)

Feed training-distribution `(obs, goal)` pairs to the planner; compare
output to ground-truth teleop action.

```
                       pred |a|     ratio       cos(pred, truth)
  truth (reference)    6.5 mm       1.00×       +1.000
  v3 baseline          33 mm        5.08×       −0.05
  v4b PRISM H=5        26 mm        3.93×       +0.10
  v4b vanilla H=5      32 mm        4.92×       +0.10
```

v4b is the closest to truth across all configs we've tried — 22 %
better ratio than v3 baseline, direction sign **flipped from -0.05 to
+0.10**. Still well above the PASS gate (≤ 2× ratio, ≥ +0.3 cos), but
clear directional progress.

### Honest summary

| Axis | v3 → v4b | PASS? |
|---|---|---|
| Cost-surface CV | 0.13 → 0.29 (+118 %) | ✗ marginal (need 0.30) |
| Test A ratio | 5.08× → 3.93× (-23 %) | ✗ (need ≤ 2×) |
| Test A cos | −0.05 → +0.10 | ✗ (need ≥ +0.3) |
| Encoder collapse | healthy | ✓ |
| Predictor stability | epochs 1–2 unstable, recovered by 18 | ✓ (with our hyperparams) |

## What this experiment shows

1. **Multi-step autoregressive supervision genuinely improves
   LeWM-MPPI's cost-surface discriminativity on small real-robot data.**
   The 2.17× CV improvement is the largest single-technique improvement
   we measured across V2, v3, and v4b. It's the technique LeWM-default
   skipped (and didn't need on 2.34 M sim frames).

2. **AR supervision is fragile on a JEPA + SIGReg + small-data setup.**
   Our first attempt (`num_preds=5`, `lr=5e-5`, `dropout=0.1`) silently
   crashed at epoch 3 with diverging val/sigreg. The conservative recipe
   that works in this release: `num_preds=3`, `lr=2e-5`,
   `predictor.dropout=0.0`, full BN running stats, stop-gradient between
   AR rollout steps.

3. **The gap from "discriminative cost surface" to "deploy-ready" is
   non-trivial.** Even with CV 2× sim's, our `GT_rank_pct` regressed
   (33 → 44 %): the cost surface gets sharper but doesn't point more
   strongly at the ground-truth action. Likely fixes (not in this
   release): stronger / deeper prior head, more demos, scheduled-sampling
   schedule on the autoregressive rollout.

## v4 vs v3 vs v2

| | v2 (negative case) | v3 (LeWM-default scale-up) | **v4 (multi-step)** |
|---|---|---|---|
| Source repo | `YuhaiW/prism-jepa-franka-pusht-v2` (private) | (not separately released) | **this repo (public)** |
| Demos | 50 | 200 | 200 |
| Frames | 18 k | 51 k | 51 k |
| Predictor training | 1-step teacher-forced + small predictor + manual normalization bug | LeWM-default 1-step | **multi-step AR, num_preds=3** |
| Cost-surface CV | n/a | 0.13 | **0.29** |
| Deploy status | known-broken negative case | wandering on real robot | research preview, deploy-marginal |

## Why drop 4 action dims

Recorded 6-dim Franka teleop action has only 2 meaningful dims:

| dim | std | std / dx_std | interpretation |
|---|---|---|---|
| dx | 0.0049 | 1.00 | signal |
| dy | 0.0058 | 1.17 | signal |
| dz | 0.0002 | 0.04 | floating-point jitter |
| drx | 0.0026 | 0.49 | Quest controller drift |
| dry | 0.0032 | 0.65 | Quest controller drift |
| drz | 0.0012 | 0.25 | Quest controller drift |

`dz` is essentially zero-noise; `drx/dry/drz` are teleop tracking
artefacts. After per-dim StandardScaler they'd all become unit-std and
look "equally important" to the model. Dropping at the data level
lets the model focus on the actual signal.

## Caveats

- **Not validated on real-robot SR.** The release is offline-diagnostic-
  best, not closed-loop success.
- **Single seed.** Multi-seed variance unmeasured.
- **Encoder is ViT-tiny from scratch** on 51 k frames — known to be
  sensitive to deploy-scene distribution shift (lighting, table colour,
  camera intrinsics). For best results, match deploy scene to training
  scene as closely as possible.
- **PRISM prior is moderate-quality.** Val MSE drop 14.7 % (marginal vs
  15 % gate). A deeper prior head would likely help further.

## Citation / context

If you use this ckpt in research, please cite the LeWM paper for the
underlying architecture and the multi-step autoregressive training
trick from world-model literature (Dreamer / TD-MPC / PlaNet style):

- LeWM: *LeWorldModel: First Stable JEPA from Pixels*. HF:
  [`papers/2603.19312`](https://huggingface.co/papers/2603.19312)
- TD-MPC / TD-MPC2: Hansen et al., for the multi-horizon AR loss style
- Dreamer / DreamerV3: Hafner et al., for imagined rollouts in training
- This release is part of the **PRISM-JEPA** project (paper in
  preparation).

## License

apache-2.0