Fine-tuning OpenPI on your robot

You have fine-tuned language models: a dataset in a house format, a tokenizer that must match the checkpoint, LoRA because the full model does not fit, an eval harness you trust more than the loss. Fine-tuning openpi (opens in a new tab) on your WidowX AI demos is exactly that job with three treacherous differences. The tokenizer is now a stack of data transforms encoding physical conventions — radians, delta actions, camera slots — and nothing type-checks them. The eval cannot be perplexity, because the loss does not measure what you care about. And a wrong config does not produce a bad paragraph; it produces a metal arm lunging at full speed toward your RealSense tripod.

Conceptual Foundation: State, Action, and the Temporal Horizon

Before configuring the pipeline, we must rigorously define the physical quantities the model processes. In robotics, the state $q_t \in \mathbb{R}^D$ represents the robot's current kinematic configuration, derived from proprioceptive sensors (encoders). It is the 'where am I now?' vector. The action $a_t \in \mathbb{R}^D$ represents the target configuration the controller should drive the robot toward. In a delta-action space, the action is not an absolute position but a displacement: $a_t = q_{t+1} - q_t$. This distinction is critical: the model predicts changes in configuration, not the configuration itself, which allows the policy to generalize across different starting poses.

The model does not predict a single instantaneous command. Instead, it predicts an action chunk: a sequence of $H$ future state targets. For a control frequency of $f_c = 50$ Hz, a chunk of $H=50$ steps represents exactly one second of real-time motion. The model outputs a tensor of shape $(H, D)$, where each row $k$ specifies the desired joint positions for time step $t+k$. This temporal abstraction exists because robotic motion is continuous and smooth; predicting a single step at a time would ignore the physical constraints of inertia and velocity limits, leading to jittery, unstable control. By predicting a chunk, the model enforces temporal consistency, ensuring that the sequence of actions forms a physically plausible trajectory.

A critical failure mode in closed-loop execution is compounding error. In open-loop evaluation (teacher forcing), the model is fed the ground-truth observation at each step, isolating prediction accuracy. In closed-loop execution, the model's own predicted action is executed, changing the robot's state, which becomes the observation for the next step. If the model makes a small error at step $t$, the resulting state deviation at $t+1$ may push the observation out of the distribution the model was trained on, causing a larger error at $t+1$. This feedback loop can lead to exponential divergence, a phenomenon the training loss never captures because it is computed under teacher forcing.

Worked Example: Delta-Action Normalization Pipeline

Consider a WidowX arm with 6 joints and a gripper (7 dimensions total). The pretrained model expects delta actions normalized to a standard range. Suppose the current state is $q_t = [0.5, 1.2, -0.3, 0.8, 0.1, 0.0, 0.0]$ radians, and the target action for the next step is $a_{t+1} = [0.52, 1.21, -0.31, 0.82, 0.11, 0.0, 0.0]$ radians. The normalization statistics for joint 0 are $\mu = 0.0$ and $\sigma = 0.02$.

Now, consider the failure mode where the delta transform is omitted. The raw absolute action $a_{t+1,0} = 0.52$ is fed directly into the normalization step. The normalized value becomes $\hat{a}_0 = \frac{0.52 - 0.0}{0.02} = 26.0$. This value is far outside the distribution the model was trained on (typically $\pm 3\sigma$). The model, encountering an out-of-distribution input, may output a saturated or erratic value. If the model outputs a normalized value of 3.0 (saturated), the serving stack unnormalizes it: $a_{exec} = \mu + \sigma \cdot 3.0 = 0.0 + 0.02 \cdot 3.0 = 0.06$ rad. This is a 6 cm error in joint space, which is significant and could cause the arm to lunge. If the statistics were wrong (e.g., $\sigma = 0.15$), the error would be amplified to $0.45$ rad, causing a violent, potentially damaging motion.

The pipeline at a glance: five stages, five failure modes

Before touching any config, hold the whole pipeline in your head. Data flows through five stages, each producing one artifact. The stages are simple; the danger is that four of the five fail silently — the loss decreases beautifully while an artifact encodes a wrong physical convention, and a robot-side debug weeks later costs 10× a check at the stage boundary.

1. Dataset conversion — your Phase 03 teleop recordings become a LeRobot dataset: camera frames, a 7-dim state vector (6 joints + gripper on a WidowX-class arm), actions, and a task string per episode.
2. Data config — a declared transform stack: repack renames your keys, a robot-specific transform pads and reshapes to the model's conventions, and model transforms resize images to 224×224 and tokenize the prompt.
3. Normalization statistics — a script iterates over the transformed dataset and writes per-dimension stats (mean/std and quantiles) into an assets directory.
4. Training config — a named entry in openpi's config registry: base checkpoint (π₀ or π₀-FAST), the LoRA freeze filter, batch size, step count, learning-rate schedule.
5. Training run — JAX compiles the step, checkpoints accumulate under an experiment directory, and the norm-stats assets are bundled with the weights so serving loads both as one unit.

One framing decision up front: which base checkpoint. The previous two lessons gave you the machinery — flow matching denoises continuous 50-step chunks; FAST decodes the same chunks as discrete tokens. For fine-tuning mechanics they are nearly identical in openpi; what matters here is that FAST's tokenizer needs bounded normalized actions and that both LoRA paths live on the JAX backend — both facts return below. Start with flow-based π₀ and treat π₀-FAST as the comparison arm of the experiment.

Data transforms: teaching your robot to speak π₀'s dialect

The pretrained model has fixed expectations: three image slots — base/exterior, left wrist, right wrist — each 224×224 RGB with a boolean validity mask; a state vector zero-padded to 32 dimensions; action chunks of $H = 50$ steps, also padded to 32 dims (one second of motion at 50 Hz). Your rig has one exterior RealSense, one wrist camera, a 7-dim state. The repack transform is pure renaming; the robot-specific transform does the physics-aware work: pad state 7 → 32, route your cameras into the base and left-wrist slots, and fill the right-wrist slot with a zero image whose mask is False — the model was pretrained with dropped views, so a masked slot is in-distribution while an unmasked black image is not.

Then comes action-space alignment, the highest-stakes decision in the file. Two axes. First, joint vs end-effector space: fine-tune in the space your demos were collected in and keep it consistent through serving — on a WidowX AI teleop rig, almost always joint positions. Second, absolute vs delta: the pretraining convention for arms like yours is delta joint actions — each action a displacement relative to the chunk-start state — while the gripper stays absolute, since open/close is positional. With $q_t$ the joint vector when the chunk is predicted and $g_t$ the gripper command:

Why deltas? They center the action distribution near zero wherever in the workspace the demo happened, making one set of normalization statistics valid everywhere and letting the pretrained prior transfer: a delta of $0.02$ rad means the same thing on your WidowX AI as on the pretraining arms. The transform is four lines of numpy; forgetting it trains a policy on values 100× outside the distribution the rest of the pipeline assumes.

import numpy as np

MODEL_DIM = 32   # pi0 padded state/action width
ARM_DIM = 7      # WidowX-class: 6 joints + gripper

def repack(sample: dict) -> dict:
    """Rename your LeRobot keys to what the robot transform expects."""
    return {
        "image": sample["observation.images.cam_base"],
        "wrist_image": sample["observation.images.cam_wrist"],
        "state": sample["observation.state"],   # (7,) float32, radians
        "actions": sample["action"],            # (H, 7) absolute joint targets
        "prompt": sample["task"],
    }

def widowx_inputs(sample: dict) -> dict:
    state = np.zeros(MODEL_DIM, dtype=np.float32)
    state[:ARM_DIM] = sample["state"]

    base = sample["image"]            # (224, 224, 3) after model resize
    wrist = sample["wrist_image"]
    missing = np.zeros_like(base)     # no right-wrist camera on this rig

    out = {
        "state": state,
        "image": {
            "base_0_rgb": base,
            "left_wrist_0_rgb": wrist,
            "right_wrist_0_rgb": missing,
        },
        "image_mask": {
            "base_0_rgb": np.True_,
            "left_wrist_0_rgb": np.True_,
            "right_wrist_0_rgb": np.False_,   # masked, not just black
        },
    }

    if "actions" in sample:           # training path only
        acts = np.asarray(sample["actions"], dtype=np.float32)  # (H, 7)
        delta = acts.copy()
        delta[:, :6] -= sample["state"][:6]   # joints become deltas
        # column 6 (gripper) stays absolute: pretraining convention
        padded = np.zeros((acts.shape[0], MODEL_DIM), dtype=np.float32)
        padded[:, :ARM_DIM] = delta
        out["actions"] = padded
    return out

Normalization statistics: the vocabulary file of continuous control

Why normalize per dimension? Raw scales differ by orders of magnitude: your WidowX AI's base yaw (J0) sweeps ±3.1 rad, the wrist pitch joints (J3, J4) only ±1.6 rad, delta actions live within ±0.05 rad, the gripper in $[0, 1]$. A shared scale would let the widest dimension dominate the loss and starve the narrow ones of gradient. So openpi computes statistics per dimension over the transformed training set and stores both flavors — flow-based π₀ consumes z-scores; π₀-FAST maps through the 1st/99th quantiles to $[-1, 1]$, because the FAST tokenizer needs bounded inputs to discretize:

The strategic question: reuse the pretrained asset's statistics, or recompute on your data? Reuse when your platform matches an embodiment the checkpoint already knows — the Trossen-arm assets for a WidowX AI rig with standard conventions — because reused stats keep the fine-tune inside the input distribution the pretrained weights expect, which matters most in the low-data LoRA regime. Recompute when your action convention, gripper hardware, or workspace differs enough that your data would occupy a thin, off-center sliver of the pretrained normalized space. Either can be right; mixing them never is. Train with statistics $(\mu_t, \sigma_t)$, serve with $(\mu_s, \sigma_s)$, and the server unnormalizes the policy's output with the wrong pair:

Put numbers on it. Your recomputed elbow delta-action std is $\sigma_t = 0.02$ rad; the pretrained asset, averaged over more dynamic motion, carries $\sigma_s = 0.15$ rad. Serve with the wrong file and every elbow command is amplified 7.5× — a demonstrated 3° adjustment becomes a 22° lunge into the joint limit within the first chunk. Flip the ratio ($\sigma_s/\sigma_t = 0.13$) and the arm traces a correctly shaped trajectory at one-eighth amplitude, stopping 20 cm short of the mug. Both failures are affine, which is your diagnostic gift: wrong stats scale and shift motion; they do not change its shape.

One quantile-specific trap: a dimension that barely moves — a gripper closed for 95% of every episode — has $q_{01} \approx q_{99}$, and the FAST normalization divides by that near-zero range, amplifying noise into full-scale swings. Eyeball the generated norm_stats.json once per dataset: a std or quantile range orders of magnitude smaller than its physical range is a thirty-second catch.

LoRA on a single RTX GPU: budget and schedule

π₀ is roughly 3.3B parameters: a PaliGemma-class 3B vision-language backbone plus a ~300M action expert. Price full fine-tuning from first principles, the same arithmetic you run for LLM training: bf16 weights and gradients, fp32 Adam moments, an fp32 master copy — $2 + 2 + 4 + 4 + 4 = 16$ bytes per parameter:

LoRA changes the ledger the way you expect: the shipped low-memory configs freeze the base weights (3.3B × 2 B ≈ 6.6 GB in bf16, no gradients, no optimizer state) and attach low-rank adapters to the transformer projections — tens of millions of trainable parameters, under 1 GB even with full Adam state. Activations dominate, and the total lands at openpi's stated 22.5 GB minimum — precisely why an RTX 4090's 24 GB is the canonical single-GPU fine-tuning card, with XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 set so JAX's allocator can claim enough of it. Your RTX workstation is not a compromise here; it is the reference platform.

Two JAX-backend realities to internalize before Week 11. First, the LoRA and FAST paths live in JAX: openpi's experimental PyTorch port supports neither LoRA, FAST, mixed precision, nor FSDP, so this phase runs on the JAX stack — run the repo's GPU smoke tests before you need them. Second, XLA compiles the whole training step on first execution: minutes of tracing and compilation, then steady-state steps orders of magnitude faster. The rhythm is torch.compile's, except it is not optional, and any shape change — a stray episode at a different camera resolution — retriggers compilation mid-run.

Now the schedule for your dataset. 150 episodes at 250–400 frames each is roughly 40–60k training frames. The shipped fine-tuning configs run on the order of 30k steps at batch 32, warmup into cosine decay, peak learning rate in the low $10^{-5}$ range — about 1M samples drawn, or 15–25 epochs over your data. That is memorization territory for 150 demos, and it is intentional: the frozen backbone plus low-rank capacity resists the worst of it. Practically:

• Checkpoint early and often — keep checkpoints every few thousand steps; the best on-robot checkpoint for a 150-demo LoRA run is frequently at 10–20k steps, not the final one.
• Expect an overnight run, not a weekend — LoRA at batch 32 on a 4090-class GPU processes this schedule in hours; if your projection is days, something is wrong (data loading off the network disk is the usual suspect).
• Change one variable per run — base checkpoint, stats choice, horizon, and LoRA schedule interact; the ablation discipline from your training-infra work applies verbatim.

Prove it off-robot: validation and environment discipline

The training loss measures how well the model fits demonstrated actions under teacher forcing. Closed-loop success depends on behavior under self-induced distribution shift, which the loss never sees; you met this as compounding error in the behavioral-cloning lesson, and a fancier architecture does not repeal it. Worse, the flow-matching objective regresses a velocity field across random noise levels, so even its absolute value is uninformative — there is no perplexity-style number where 2.1 is good and 3.4 is bad.

The cheapest eval that touches physical units is open-loop action prediction on held-out episodes. Hold out ~10 episodes before training (your Phase 03 frozen splits). Step through each at chunk boundaries: feed the demonstrated observation, predict a 50-step chunk, compare against the demonstrated actions — teacher-forced, so it measures fit-plus-generalization without compounding. Run the policy through the exact serving-side transform and unnormalization stack, so the numbers are radians headed over UDP to the arm's controller and the test covers your stats plumbing for free.

import numpy as np

H = 50        # action horizon: one chunk = 1 s at 50 Hz
ARM_DIM = 7   # 6 joints + gripper

def open_loop_eval(policy, episodes):
    """policy(obs) -> (H, ARM_DIM) actions in physical units (rad),
    run through the same transforms + unnormalization as serving.
    episodes: held-out list of dicts with per-step obs and actions."""
    errs = []
    for ep in episodes:
        T = len(ep["actions"])
        for t0 in range(0, T - H, H):        # jump chunk to chunk
            obs = ep["obs"][t0]              # demonstrated obs (no rollout)
            pred = policy(obs)               # (H, 7)
            ref = np.asarray(ep["actions"][t0:t0 + H])
            errs.append(np.abs(pred - ref))
    errs = np.concatenate(errs, axis=0)      # (N, 7)

    mae = errs.mean(axis=0)
    p95 = np.percentile(errs, 95, axis=0)
    names = ["joint_0", "joint_1", "joint_2", "joint_3",
             "joint_4", "joint_5", "gripper"]   # WidowX AI index convention
    for j in range(ARM_DIM):
        print(names[j].ljust(10),
              "MAE", round(float(mae[j]), 4),
              "p95", round(float(p95[j]), 4))
    return mae, p95

Then plot predicted vs demonstrated trajectories for all seven dimensions per held-out episode, chunk boundaries marked. The plots carry diagnoses the scalar table cannot: a constant vertical offset on one joint is a stats bias term; a clean amplitude ratio is the gain error derived above; discontinuities at every chunk boundary mean the state input is mis-scaled; a gripper closing 10 frames late misses a moving grasp despite a beautiful arm MAE. Rough calibration for a WidowX-class task: held-out arm-joint MAE of 0.02–0.05 rad is typically deployable; above ~0.15 rad the robot trial is a waste of an afternoon.

The final pre-robot gate is a sanity replay in a mock environment: run the actual policy server, point a fake client at it that replays recorded observations at the real 50 Hz cadence, and assert on shapes, physical ranges, chunk timing, and inference latency. An hour of work, and it catches the whole class of serving-boundary bugs — assets not found, quantile clipping engaged, an unmasked image slot — with zero grams of robot at risk. It is also the skeleton of the inference contract the next lesson formalizes.

One last off-robot gate remains, and it is the unglamorous discipline that makes everything above reproducible: collection, training, and inference are three separate environments with independent pins, because they have incompatible dependency graphs that all name the same fast-moving library. Collection pins the LeRobot version your teleop and camera drivers were validated against; openpi's training environment pins its own LeRobot plus JAX; inference splits again into the server environment and a lightweight client on the robot host. LeRobot's dataset format and API have churned across versions — that a dataset written by your collection pin is readable by openpi's training pin is something you verify once, then freeze, never assume.

The Trossen OpenPI tutorial (opens in a new tab) walks this exact stack end to end and is your closest published reference — with one caveat to tape to your monitor: its worked examples and reported results are bimanual, built on two-arm Trossen AI kits. Every single-arm adaptation — one arm's state and action dims, one wrist camera masked — is yours to make and yours to validate; the tutorial's numbers are an existence proof for the pipeline, not a baseline for your configuration. On a single-arm rig you are the first person running your exact config — budget the validation gates accordingly.

Where this goes next: the previous lesson, π₀-FAST: actions as tokens done right, gave you the second base checkpoint this pipeline fine-tunes; everything here — transforms, stats, LoRA, the open-loop gate — applies to both action heads. The next lesson, Serving a policy: the inference contract, picks up the artifact this one produces: checkpoint plus normalization assets become a policy server with a precise observation-in, chunk-out contract, whose latency behavior is exactly what your Phase 05 capstone on latency-aware action chunking will measure and optimize.