JingwuLuo/ur3-optical-resample

π0.5 fine-tune: dual-arm UR3, optical bench rod transfer (curvature-resampled action chunks)

π0.5 (openpi) fine-tuned on the LeRobot dataset LGG100/optical-bench-lerobot-cleaned: a dual-arm UR3 doing one task, "pull the metal rod out of the post holder and place it in the white box". 40 episodes, 9588 frames, 15 Hz, three cameras (top, left wrist, right wrist) at 240×320, 14-dim joint state and action. The left arm stays parked in every episode; the right arm does the task with one gripper close and one open. The vision-language model is frozen; only the action expert was trained.

Before deploying: one answer from this model is not one target per frame. It is 10 points placed along the trajectory by geometric curvature, covering the next 30 recorded frames (2 s at the recording's pace). They must be interpolated back onto the 15 Hz frame grid before they go to the robot. deploy/serve_config.py does that on the server, so result["actions"] is already one target per frame. Sending the 10 points as 10 consecutive frames would play the motion 3× too fast, with the wrong timing.

Files

Path	Contents
jax/params/	JAX (orbax) weights, checkpoint at step 9999
jax/assets/optical-bench-images-curvature-s30-d2-train/norm_stats.json	normalization statistics (quantile; see Normalization statistics)
jax/assets/optical-bench-images-curvature-s30-d2-train/arc_resample.json	the resampling settings the labels were built with, the robot layout and per-episode label statistics; the interpolator reads it
pytorch/model.safetensors	PyTorch weights, bf16, converted from jax/; checked against it (see PyTorch weights)
pytorch/assets/optical-bench-images-curvature-s30-d2-train/	the same norm_stats.json and arc_resample.json as under jax/assets/
pytorch/config.json	written by openpi's converter (action dim, horizon, variants, precision); openpi does not read it
deploy/serve_config.py	openpi inference config for this model, with the interpolation as the output transform
deploy/arc_interp.py	points → per-frame targets; numpy + scipy only

jax/ (params/ + assets/) and pytorch/ (model.safetensors + assets/) are each a complete openpi checkpoint directory. Pass either one to serve_config.get_train_config and policy_config.create_trained_policy; openpi runs PyTorch when model.safetensors is present. The optimizer state is not included.

Inputs and outputs

Observation keys:

Key	Format
observation/image	top camera (cam_high), uint8 HWC RGB, 240×320 (resized with padding to 224×224 inside)
observation/left_wrist_image	left wrist camera, same format
observation/right_wrist_image	right wrist camera, same format
observation/state	float32 (14,)
prompt	"pull the metal rod out of the post holder and place it in the white box"

State and action layout: 0-5 left arm joints (shoulder_pan, shoulder_lift, elbow, wrist_1, wrist_2, wrist_3; rad), 6 left gripper (0 open, 1 closed), 7-12 right arm joints, 13 right gripper. Actions are absolute joint targets; action[t] is the pose for frame t+1.

The model receives no state (discrete_state_input=False): it sees only the images and the prompt. Inside the model, all 12 joints and the left gripper are deltas from the state and the right gripper is absolute; openpi converts back, so the policy returns absolute values. The left gripper is open in 11 recordings and closed in 29, constant within each; as a delta it stays wherever the robot reports it.

Serving

The model must be built with the training config. Anything else will either fail to load these weights or misread them:

Pi0Config(pi05=True, action_horizon=10, discrete_state_input=False)   # gemma_2b + gemma_300m, no LoRA

with quantile normalization and the delta mask make_bool_mask(13, -1). deploy/serve_config.py contains exactly that and needs nothing but upstream openpi:

import sys; sys.path.insert(0, "<download>/deploy")
from openpi.policies import policy_config
import serve_config

checkpoint = "<download>/jax"   # or "<download>/pytorch"
config = serve_config.get_train_config(checkpoint, rate=1.0)
policy = policy_config.create_trained_policy(config, checkpoint)
result = policy.infer(obs)
result["actions"]           # (30, 14) per-frame targets; result["actions"][j] is for frame t + j + 1
result["arc_points"]        # (10, 14) the model's own points, absolute
result["arc_point_frames"]  # (10,)    the frames they were placed at: 3, 6, ..., 30

Serve it over websocket with openpi.serving.websocket_policy_server.WebsocketPolicyServer(policy=policy, host="0.0.0.0", port=8000, metadata=policy.metadata).serve_forever(); the robot side uses openpi_client.WebsocketClientPolicy and gets the same dict. With jax/, the first request compiles the model (15 s on an L40S); later round trips took a median 98–101 ms (about 2,400 requests, local websocket, an L40S shared with other jobs).

Execute the first 8 targets, then request again. The evaluation below uses that execution length.

Check: deploy/serve_config.py against the training harness's serving path (vla_openpi.launch.serve), with one loaded model, 6 recorded frames and identical noise: per-frame targets, points and point frames identical (maximum difference 0). Loaded in two separate processes the two differ by up to 2.9e-3 rad, and each process repeats its own outputs exactly.

PyTorch weights

pytorch/model.safetensors was converted from jax/ with openpi's converter (examples/convert_jax_model_to_pytorch.py) and stored in bf16. The model has no LoRA adapters, so every trained weight has a PyTorch counterpart. The converted weights fill every parameter of openpi's PI0Pytorch except the two lm_head matrices, which action prediction does not use.

The check ran both directories through deploy/serve_config.py with identical observations and identical noise. It used 16 recorded frames, at 1/3 and 2/3 of episodes 8, 15, 24, 27, 0, 10, 20 and 30. The fp32 rows load the JAX weights as float32 and a float32 conversion made only for this check, both with full fp32 matmul precision. Each entry is the largest joint difference over the 10 points of one answer, median / max over the 16 answers, in rad:

Comparison	Joints	Right gripper commands that differ
JAX vs PyTorch, both fp32 (TF32 off)	0.0004 / 0.0029	0 of 160
JAX bf16 vs JAX fp32	0.0056 / 0.029	0 of 160
PyTorch bf16, compiled vs PyTorch fp32	0.0072 / 0.062	0 of 160
JAX bf16 vs PyTorch bf16, not compiled	0.0073 / 0.043	0 of 160
JAX bf16 vs PyTorch bf16, compiled (the default for each)	0.014 / 0.080	0 of 160
for scale: JAX bf16, two different noise draws	0.089 / 0.43	4 of 160

The left gripper command agrees in every comparison, and the per-frame targets differ by about the same amounts. In fp32 the two frameworks agree to 4e-4 rad, so the conversion is faithful. In bf16, each framework departs from its fp32 result by 0.006–0.007 rad (median), and torch.compile changes the rounding again. The two defaults therefore differ by 0.014 rad (median): a quarter of the median spacing between consecutive points (0.057 rad), and about a sixth of the change from drawing different noise. The open-loop evaluation below ran on jax/.

The PyTorch path needs openpi's patched transformers==4.53.2 (copy src/openpi/models_pytorch/transformers_replace/* over the installed package). By default openpi compiles the model with torch.compile (pytorch_compile_mode="max-autotune"). Measured on one L40S shared with other jobs, timing policy.infer only (no websocket):

• loading took 72 s;
• the first two calls took 25.6 s and 4.5 s;
• later calls took a median 62 ms, against 84 ms for jax/ measured the same way.

Without compilation (pytorch_compile_mode=None in the model config) a call took 357 ms.

Action resampling

How the training labels were built

A per-frame policy predicts action[t : t+10], 10 consecutive frames, about 0.67 s at 15 Hz. This model predicts 10 points spread over a 30-frame window instead, so one answer covers 2 s while keeping a 10-step action head:

• Window. For frame t, the window is the 30 recorded segments starting at the current state (anchor: "state"; span_frames: 30, action_horizon: 10).
• Spacing by curvature. The points are spaced uniformly in a progress coordinate that accumulates c(i) · ℓ(i) along the joint path, where ℓ is the segment's arc length over the 12 arm joints and c = κ^(1/3), κ the path's geometric curvature. κ is measured with chords of fixed arc length curvature_length = 0.1389 rad (4 median path steps), so still frames add neither length nor spurious curvature. Points crowd where the arm turns and thin out on straight moves, regardless of how fast the operator moved. The progress per point is set so that a window without keyframes holds its 10th point exactly at the 30th recorded frame.
• Keyframe windows. Around each gripper event, the label switches to uniform-in-time spacing with event_density: 2: 2 points every 3 recorded frames (twice the free part's average). The windows are close (right gripper command rising) −8/+5 frames, release (falling) −8/+5, lift (arm starts moving after the grasp, joint speed above 0.012 rad/frame for 3 frames) −2/+6. Each gripper switch therefore sits between closely spaced points taken from the recorded path.
• End padding. Each episode ends with 45 extra frames repeating the last recorded frame with the open command held, so the policy learns to stay still once the rod is in the box. Points that reach past the end of the recording fall on these rows and repeat the final pose.
• Loss weights. key_weight and pad_weight are both 1: every point counts the same.

Over the 36 training episodes: the 10th point reaches a median 30.0 recorded frames ahead; 17.6% of points fall on padding rows and 16.6% inside keyframe windows.

How the server rebuilds the trajectory

The model predicts the points but not their times. The interpolator (deploy/arc_interp.py, ArcInterpolator.from_report(arc_resample.json, rate)):

1. places point k (k = 1..10) at 3k / rate frames after the observation (3, 6, ..., 30 at rate 1);
2. interpolates the joints with a not-a-knot cubic spline through the current pose (at frame 0) and the 10 points;
3. sets each gripper to the value of the point nearest in time, re-binarised to 0/1. The current state is not used for the grippers: it echoes the previous command and may lag;
4. returns one target per control frame, floor(30 / rate) targets (30 at rate 1).

At rate 1 a window without keyframes is executed over 30 frames, the recording's average pace over that window; turns are slower and straight moves faster than the operator moved them. Windows containing a gripper event are executed more slowly (density 2 > 1), so the arm lingers near each switch.

What the reconstruction costs, measured on the labels

Feeding the ground-truth labels through the same reconstruction, closed loop over the 40 episodes (rate 1, executing 8 targets per chunk, the next chunk rebuilt from the pose reached):

	Median	p90	Max
Distance of the executed targets from the recorded path	0.0003 rad	0.0032 rad	0.092 rad
Arm pose at each gripper switch vs the recorded pose there	0.0011 rad	—	0.116 rad

No gripper switch is lost, and the task takes 10442 frames against 9548 recorded (1.09×).

Faster execution

rate > 1 traverses the same spline path faster. arc_interp.py also supports velocity timing (VelocityTiming). This keeps the same spline but times it so that every joint stays within a fraction of its velocity limit, with the rate as a cap, and holds the arm still at each gripper switch. It is applied on the robot side, from result["arc_points"] and the poses commanded so far, through ArcInterpolator.from_report(path, rate, timing).chunk_targets(...). In label replay with the UR3 limits (π rad/s base, shoulder and elbow; 2π rad/s wrists), a 100 ms inference latency, asynchronous execution, and 90% of 37 recordings completing within 2 path steps of the recorded path (3 recordings that themselves exceed the limits left out), the highest speed-up over the recording is:

Timing	Speed-up	Setting
uniform (rate only)	0.89×	rate 1, 14 targets per chunk
velocity timing, no switch hold	4.50×	cap 6, 2 targets per chunk
velocity timing, holds 0.533 s (close) / 0.667 s (open)	3.80×	cap 8, 2 targets per chunk

These are replays of the ground-truth labels. The model has been evaluated only at rate 1 with uniform timing, and nothing has run on the robot.

Normalization statistics

norm_stats.json holds openpi quantile statistics, computed over the 10191 training rows after the delta transform. The model sees x_norm = 2 (x − q01) / (q99 − q01 + 1e-6) − 1; its outputs are mapped back with the inverse.

Actions are the 10 label points, joints and left gripper as deltas from the state, right gripper absolute. They reach up to 30 frames ahead, so their ranges are wider than per-frame deltas:

Dim	q01	q99	mean	std
0–5 left joints	−0.025	0.025	0	0.0125
6 left gripper (delta)	−0.025	0.025	0	0.0125
7 right shoulder_pan	−0.3150	0.2138	−0.0067	0.0758
8 right shoulder_lift	−0.5038	0.7711	0.0695	0.2010
9 right elbow	−1.4411	0.5123	−0.0677	0.3043
10 right wrist_1	−1.2688	1.4868	0.0283	0.5077
11 right wrist_2	−0.5035	0.5378	−0.0183	0.1497
12 right wrist_3	−1.0752	0.8932	−0.0254	0.2412
13 right gripper (absolute)	0	0.9998	0.5433	0.4981

Dims 0–6 are widened. The parked left arm's deltas measured q99 − q01 = 1.0e-4 to 1.5e-4 rad (encoder noise) on the joints and exactly 0 on the left gripper. Quantile normalization divides by that range, which would train the model to predict noise stretched to [−1, 1] (and divide by 1e-6 on the gripper). The range is set to ±0.025 rad around the centre (std 0.0125), so the same noise maps to values near 0. The left-arm targets the model then commands stay within 1.2e-3 rad of the observed pose (see Open-loop evaluation).

State statistics are stored because openpi normalizes every key, but they have no effect: the model receives no state, and the state is un-normalized again before the deltas are added back. Recorded ranges (q01 / q99): left joints parked at −1.6553, −1.5411, −1.3606, −1.7022, 1.6544, 3.0440 (q99 − q01 ≤ 2e-4 rad); left gripper 0 / 0.9998; right joints 1.2876 / 2.0045, −1.5223 / −0.0127, −1.2492 / 2.0796, −4.2848 / −0.2946, −1.6316 / −0.8037, −1.7529 / 1.6304; right gripper 0 / 0.9998.

Data

Step	Result
Source	LGG100/optical-bench-lerobot-cleaned: 40 episodes, 9588 frames, 176–388 per episode (median 225), AV1 videos decoded once to PNG; every decoded frame checked against the stored image
Audit and cleaning	12 large joint steps reviewed (11 continuing a fast turn, 1 discontinuity at episode 35, frame 80): all kept as real fast motion; nothing repaired or excluded
Still frames	all kept (pauses under 1 s); 45 padding frames appended per episode → 11388 rows
Keyframes	grasp (close edge), slowest approach frame, lift, release (open edge); one of each per episode; 35 consistency flags reviewed on wrist-camera sheets and accepted
Held-out split	episodes 8, 15, 24, 27, stratified over the two large clusters (box position and angle, recording session)
Training set	36 episodes, 10191 rows, 1620 of them padding

Training

Base	pi05_base (openpi)
Trained	action expert (Gemma 300M) and the action and time projections, about 430M parameters; SigLIP and Gemma 2B frozen (about 2.9B)
Steps / batch	10,000 / 32; two L40S data parallel (16 per GPU), 1.7 s per step, 5 h 01 min
Optimizer	AdamW (β = 0.9, 0.95; ε 1e-8; weight decay 1e-4; grad clip 1.0), no EMA
Learning rate	cosine, 500 warmup steps, peak 2.5e-5, down to 2.5e-6 at 10k
Loss (flow matching, normalized space)	0.2236 at step 0, 0.1293 at 100, 0.0159 at 1000, 0.0061 at 5000, 0.0039 at 9900

Open-loop evaluation

Every recorded observation of an episode was sent to the served checkpoint (rate 1). The output was compared with the recording. Because each request uses the recorded observation, a wrong prediction never feeds back. These numbers are prediction errors, not a task success rate. The model has not been run on the robot. Errors are the worst of the 12 arm joints at each frame, in rad. One frame of motion (the median per-frame joint step) is 0.017 rad. Pooled median / p90 over the real (non-padding) frames:

	Held-out 8, 15, 24, 27	Training 0, 10, 20, 30
Model's 10 points vs the label	0.340 / 0.967	0.120 / 0.333
Per-frame targets, first 8 frames	0.159 / 0.414	0.112 / 0.306
Replanning every 8 frames, stitched trajectory	0.077 / 0.285	0.057 / 0.214
Drift while the recording holds still (45 padding frames)	0.005 / 0.125	0.004 / 0.007
Largest move commanded on the parked left arm	1.2e-3	7.4e-4

By phase, stitched / points median: held-out approach 0.071 / 0.215, carry (grasp to release) 0.081 / 0.457; training approach 0.064 / 0.113, carry 0.052 / 0.131. The gap between held-out and training episodes is largest in the far points during the carry to the box.

Gripper switches in the stitched trajectory, frames relative to the recorded switch (positive = later):

	Close	Open
Held-out	+8, +5, +5, +5	−1, −2, −4; episode 8 never opens
Training	+2, +1, +7, +3	−2, −8; episodes 10 and 30 do not open within the recorded rows

A second run over the three episodes without an open, this time including the 45 padding frames: episode 30 opened 4 frames after the recorded open (inside the padding), episode 10 one frame after, and episode 8 did not open at all. The model draws new noise on every request, so the individual switch frames vary between runs.

Limitations

• Release on held-out episodes. On held-out episode 8 the stitched trajectory never opens the gripper, and the arm keeps moving through the padding frames (the drift p90 of 0.125 rad comes mostly from that episode). On held-out episodes the gripper closes a median 5 frames (0.33 s) after the recorded close.
• Timing is assumed, not predicted. The server places points at fixed frames (3k / rate). The real time between points varies with the path's curvature and the operator's speed. For this static task that changes the speed along the path, not the path itself.
• No state input. The model knows whether the gripper is closed only from the images.
• Small dataset. 36 training episodes; the held-out evaluation covers 4.
• Faster execution (rate > 1, velocity timing) is measured only on replays of the labels.
• openpi downloads the PaliGemma tokenizer from gs://big_vision on first use.

License

Apache-2.0. The weights derive from π0.5 base (PaliGemma / Gemma), so the Gemma Terms of Use also apply.