prism-jepa-franka-pusht-v2

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

Bundle contents

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

⚠️ 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.

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

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:

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