PRISM-JEPA for Franka FR3 Planar PushT — V2 (deploy candidate)

This is the V2 deploy candidate — a LeWM-style JEPA world model + PRISM prior head trained on 50 demos / 18k frames of Franka FR3 planar PushT, plus the inference scaffolding needed to run PRISM-MPPI on the robot.

Status: Lab-tested healthy by 3 sanity checks (see §"Diagnostic numbers" below). Cleared for real-robot deployment with the safety clamps in §"Safety".

⚠️ Predecessor: A v1 ckpt was previously published as a documented negative case (HF: YuhaiW/prism-jepa-franka-pusht). That one was broken because (i) the scene composition let the encoder latch onto the robot arm and ignore the T-block, and (ii) the LeWM-default predictor (10.79 M params) landed in a "near-identity attractor" on the small dataset. V2 fixes both issues (see §"V2 vs v1" below).

TL;DR


Source dataset	`Rongxuan-Zhou/pusht_lewm_fr3` (50 demos, 18,014 frames @ 10 Hz)
Action space	2-dim (dx, dy) — planar delta-EE in meters
World model	ViT-tiny encoder + small ARPredictor (1.14 M params, 10.5 % of LeWM default)
Prior head	PriorHead MLP — val MSE drop 19.8 % (passes 15 % gate)
Action conditionality	sens/base = 1.28 (vs sim-PushT ref 1.12, vs random 0) → MPPI cost is discriminative
Bundle size	~ 36 MB (small predictor → smaller ckpt than LeWM-default)

Bundle contents

File	Size	Role
`lewm_pusht_lewm_fr3_v2.ckpt`	34 MB	LeWM (JEPA encoder + small AR predictor), pickled
`prior_head_pusht_lewm_fr3_v2.pt`	2 MB	PRISM head state-dict + StandardScaler (action)
`franka_pusht_v2_inference_demo.py`	13 KB	Self-contained PRISM-MPPI / vanilla-MPPI inference loop
`jepa.py`, `module.py`	~ 11 KB	Model classes (required to unpickle the ckpt)
`prior_head.py`	2.4 KB	`PriorHead` class
`requirements.txt`	0.4 KB	Runtime dependencies
`README.md`	this file	Usage + deployment guide

Installation

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

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

Quick start

import numpy as np
from franka_pusht_v2_inference_demo import (
    PrismMPPIInferenceV2,
    pad_2d_to_6d_franka,
)

# 1. PRISM-MPPI (recommended)
planner_prism = PrismMPPIInferenceV2(
    lewm_ckpt  = "lewm_pusht_lewm_fr3_v2.ckpt",
    prior_ckpt = "prior_head_pusht_lewm_fr3_v2.pt",
    use_prism  = True,        # PoG-fuse prior into MPPI seed
    device     = "cuda",
)

# 2. Or vanilla LeWM-MPPI (for A/B comparison)
planner_vanilla = PrismMPPIInferenceV2(
    lewm_ckpt  = "lewm_pusht_lewm_fr3_v2.ckpt",
    prior_ckpt = "prior_head_pusht_lewm_fr3_v2.pt",
    use_prism  = False,
    device     = "cuda",
)

# Plan once, get 5 env-step actions (dx, dy)
obs_uint8  = camera.read_d455_agent_view()   # (224, 224, 3) uint8 RGB
goal_uint8 = goal_image                       # (224, 224, 3) uint8 RGB
actions_2d = planner_prism.plan(obs_uint8, goal_uint8)
# actions_2d.shape == (5, 2)

Franka FR3 deployment

Action format (V2 specific)

V2 was trained on a 2-dim action space (dx, dy only). The 4 other dims (dz, drx, dry, drz) recorded by the original 6D teleop pipeline contained only floating-point jitter and Quest controller drift (std ratios dz/dx = 0.04, drx/dx = 0.49, etc.) — they were dropped at the dataset level before training. On the Franka the planner outputs 2D and you pad to 6D before sending to the robot:

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

Training-distribution range (raw, per env-step at 10 Hz):

idx	meaning	training ± (m or rad)	unit
0	dx	± 0.025	m
1	dy	± 0.029	m

⚠️ Safety clamps (required)

Even a healthy V2 may occasionally output actions outside the training range. Clamp before sending to the robot:

ACTION_CLAMP_2D = np.array([0.030, 0.034])    # ≈ 1.2× max(|training|)
def clamp_safety(actions_2d):
    return np.clip(actions_2d, -ACTION_CLAMP_2D, +ACTION_CLAMP_2D)

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

Additionally on the hardware side:

Workspace bounding box (x_min..x_max, y_min..y_max, fixed z)
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

CONTROL_DT = 0.1            # 10 Hz (matches training)
N_EXEC = 5                   # = A_block; replan after each block
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 (lab-verified)

V2 was put through three pre-deploy sanity checks. Reference column = the official sim PushT ckpt from the LeWM paper, trained on 2.34 M frames.

Check	V2	sim PushT (ref)	V2 verdict
Encoder collapse (effective rank @ 90 % var)	35 / 192	80 / 192	✓ healthy, no collapse
Autoregressive rollout pred / id @ h = 5	0.960 (val)	0.854 (val)	acceptable — see "Action conditionality"
z-step size ‖z_t − z_{t+5}‖ vs sim	ratio 0.76	1.00	✓ active z-trajectory
Action sensitivity (sens / base)	1.28	1.12	✓ predictor is action-conditional

The sens / base number is the key one for MPPI deployment: it measures how much the predictor's output changes when the action input changes, vs the typical predictor displacement. ≥ 0.5 is enough for MPPI to discriminate candidates; V2 actually exceeds the sim reference baseline.

A note on the pred / id metric

V2's pred / id @ h = 5 is 0.96, which looks worse than the sim ckpt's 0.85. We initially read this as a problem, but the deciding metric for MPPI usability is sens / base, not pred / id. A model can have pred / id close to 1 (= predictions absolute-close to identity baseline) but still be highly action-conditional in the variance across candidates, which is what MPPI uses to discriminate. Confusing these two metrics led us to publish a different ckpt initially (see lewm_pusht_lewm_fr3_smallpred_2d_object.ckpt in the project repo) that turned out to be action-blind despite a lower pred / id — V2 is the correct trade-off.

V2 vs v1 — what changed

	v1 (broken negative case)	V2 (deploy candidate)
Dataset	36 demos, 8.8k frames	50 demos, 18k frames
Scene	Arm dominant, T small near frame edge	T centered, less arm bias
Action recorded	6-dim (dx, dy, dz, drx, dry, drz)	2-dim (dx, dy only)
Predictor	Default 10.79 M params	1.14 M params (10.5 % capacity)
Encoder is arm-proxy?	Yes (r = 0.51 with proprio)	No (r = 0.40)
sens / base	not measured for v1	1.28
Deploy status	Don't deploy (documented negative case)	Deploy-ready with safety clamps

Why a 10× smaller predictor?

LeWM's default predictor (depth = 6, heads = 16, mlp = 2048) was sized for the 2.34 M-frame sim PushT dataset. On Franka's 18k-frame real-world dataset, that capacity admits a "near-identity attractor" — the predictor can drop its loss to near-zero by outputting f(z_t) ≈ z_t + bias, ignoring the action entirely. Shrinking the predictor by 10× breaks that attractor and forces the model to actually use the action input.

This is a regularization-by-capacity-reduction story. Detailed analysis and the H2/H3/H4 ablation results are in docs/25_franka_pusht_v2_predictor_capacity.md of the source project.

Why drop the 4 action dims?

Inspection of the recorded 6D action showed:

dim	std	std / dx_std
dx	0.0040	1.00
dy	0.0047	1.17
dz	0.0002	0.04
drx	0.0020	0.49
dry	0.0026	0.65
drz	0.0010	0.25

dz is essentially floating-point jitter; drx / dry / drz are Quest controller drift. After per-dim StandardScaler they all become unit-std and look "equally important" to the predictor → 67 % of the input is noise. Dropping the 4 dims at the data level lets the predictor focus on the signal.

Caveats

Small dataset. 50 demos is enough for the predictor to learn action-conditional dynamics under the 10× capacity reduction, but more data would tighten the pred / id gap to sim and reduce variance. We recommend collecting another batch (target 100-150 demos) before publishing the model for general use.
Single seed. All numbers above come from one V2 training run (seed = 3072). Multi-seed variance has not been measured.
No success-rate evaluation. Action sensitivity ≠ task success. The model has cleared offline diagnostics but the real-robot SR (vs the v1 baseline's 0 % and sim PushT's 60-92 %) is what matters at the end. Deployment results will be published when available.

Cross-references

Source dataset: Rongxuan-Zhou/pusht_lewm_fr3
v1 (negative case): YuhaiW/prism-jepa-franka-pusht
Sim PushT reference ckpt: lewm-pusht (LeWM paper official)
Paper: PRISM-JEPA, §4 (sim experiments) and §5 (real-robot deployment)

License

apache-2.0

Downloads last month: -; Downloads are not tracked for this model. How to track

Video Preview

Robotics