--- 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