JingwuLuo/ur3_simple_pick

π0.5 LoRA fine-tune: dual-arm UR3, simple pick (progress-resampled action chunks)

π0.5 (openpi) fine-tuned with LoRA on the LeRobot dataset LGG100/ur3-test: a dual-arm UR3 doing one task, "pick up the red cube and place it on the white plate". 22 episodes, 3223 frames, 15 Hz, three cameras (top, left wrist, right wrist) at 240×320, 14-dim joint state and action.

Before deploying: one answer from this model is not one target per frame. It is 10 points spaced by progress along the trajectory, reaching 20 recorded frames (1.33 s) ahead. They must be interpolated back onto the 15 Hz frame grid before they go to the robot — deploy/ur3_arc_interp.py does that. Sending the 10 points as if they were 10 consecutive frames would play the motion about twice as fast and with the wrong timing.

Files

Path	Contents
jax/params/	JAX (orbax) weights as trained: bf16 base weights plus fp32 LoRA adapters
jax/assets/ur3-test-arc-s20/norm_stats.json	normalization statistics (quantile)
pytorch/model.safetensors	PyTorch weights, bf16, LoRA adapters folded into the base weights; checked against jax/ (see below)
pytorch/assets/ur3-test-arc-s20/norm_stats.json	same normalization statistics
deploy/ur3_arc_interp.py	points → per-frame targets; numpy + scipy only
deploy/ur3_arc_serve_config.py	openpi inference config for this model, using the interpolator

jax/ and pytorch/ are each a complete openpi checkpoint directory: pass either one to policy_config.create_trained_policy, which picks JAX or PyTorch by whether model.safetensors is present.

Inputs and outputs

Observation keys (the same as a per-frame UR3 policy):

Key	Format
observation/image	top camera, uint8 HWC RGB, 240×320 (resized to 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	"pick up the red cube and place it on the white plate"

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. Inside the model the joints are deltas from the state and the grippers are absolute; openpi converts back, so the policy returns absolute values.

Using it in an existing openpi server

The model must be built with this config — it is the training config, and anything else will not load these weights or will misread them:

Pi0Config(pi05=True, action_horizon=10, discrete_state_input=False,
          paligemma_variant="gemma_2b_lora", action_expert_variant="gemma_300m_lora")

with quantile normalization and the delta mask make_bool_mask(6, -1, 6, -1) (joints as deltas, grippers absolute). deploy/ur3_arc_serve_config.py contains exactly that, so the simplest route is:

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

policy = policy_config.create_trained_policy(ur3_arc_serve_config.get_train_config(), "<download>/jax")
result = policy.infer(obs)
result["actions"]           # (21, 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, 5, ..., 21

result["actions"] has the same meaning as a per-frame policy's actions, so an existing client works unchanged — it just receives 21 targets instead of 10. If your server has its own output path instead, the only step to add is the interpolation, applied to the policy's absolute points and the state the observation carried:

from ur3_arc_interp import ArcInterpolator

interp = ArcInterpolator(span_frames=20)   # must be 20 for this checkpoint
targets = interp(points, state)            # points (10, 14) or padded (10, 32); state (14,) or (32,)
# targets: (21, 14); targets[j] is the target for frame t + j + 1

Execute the first ~6 targets, then request again (see the limitations below for why).

How the action chunk was built, and what the interpolation assumes

For each frame the training label is 10 points placed uniformly in a progress coordinate s(t) = Σ |Δ²q|^(1/3) (a curvature measure over the 12 arm joints), scaled so that 10 points cover 20 recorded frames. Points therefore crowd where the arm bends and thin out on straight moves. Around each gripper command change (8 frames before, 5 after) the label instead spends 3 points spread evenly in time, and those steps get 3× loss weight.

The model predicts the points but not their times. The interpolator places point k at frame 1 + 2k (3, 5, ..., 21), the average spacing the construction produces (the measured medians are 3.1, 5.1, ..., 21.6). It then interpolates the joints with PCHIP, starting from the current pose, and holds each gripper at its nearest point, re-binarised to 0/1.

Training

Base	pi05_base (openpi)
Fine-tuning	LoRA, rank 16 on the VLM and rank 32 on the action expert, attention + MLP; alpha = rank
Steps / batch	10,000 / 32, one L40S, about 12 h
Optimizer	AdamW (β = 0.9, 0.95; 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
Final loss	≈ 0.05 (weighted by the per-step loss weights, so not comparable to a per-frame model's loss)

Open-loop check

Each of the 22 training episodes was replayed through the JAX checkpoint frame by frame, using the recorded observations, and the output was compared with the recorded targets. These numbers show the pipeline works and how closely the model fits the data it was trained on. They are not a task success rate, and the model has not yet been run on the robot. Errors are the worst left-arm joint, reported in rad and in frames of motion (the median per-frame joint movement in each episode's recording, 0.013–0.032 rad). Values are medians over episodes.

	Median	p90
Model's 10 points vs. the training label	0.027 rad (1.4 frames)	0.040 rad (2.1 frames)
Per-frame targets, first 6 frames	0.020 rad (1.1 frames)	0.037 rad (1.9 frames)
Per-frame targets, all 21 frames	0.040 rad (2.2 frames)	0.166 rad (7.4 frames)
Replanning every 6 frames, stitched trajectory	0.012 rad (0.7 frames)	0.024 rad (1.3 frames)

All 44 recorded gripper changes appear in the stitched trajectory, with no open/close flicker. Round-trip latency was about 100 ms per request on one L40S (JAX, local websocket).

PyTorch weights

pytorch/ was converted from jax/ with the LoRA adapters folded into the base weights first: openpi's stock converter (examples/convert_jax_model_to_pytorch.py) maps only base weight names and silently drops the adapters. The fold follows openpi's LoRA einsums exactly. In particular, openpi's attention output projection sums lora_b over heads, so its adapter is not a per-head lora_a @ lora_b.

Checked against jax/ on 16 recorded frames with identical observations and identical noise: joint outputs differ by 2.8e-3 rad (median) and 8.7e-3 rad (max), against a median step of 2.7e-2 rad between output points, which is bf16 rounding. All 160 gripper commands are identical.

The PyTorch path needs openpi's patched transformers==4.53.2 (src/openpi/models_pytorch/transformers_replace). With the default config it runs torch.compile (mode max-autotune) on the first call, which takes several minutes.

Limitations

• Timing is assumed, not predicted. Near gripper events the label spaces points about 4.3 frames apart, not 2, so placing them 2 frames apart compresses the motion there. Replanning every 6 frames, the gripper closes 1–4 frames early (median 3, about 0.2 s) and opens 0–3 frames early (median 1). Feeding the ground-truth labels through the same interpolation gives the same shifts, so this comes from the timing assumption, not the model. Error also grows over the tail of the 21 frames, which is why only the first ~6 should be executed.
• The right arm never moves in the training data. In every episode the right arm stays within 3e-4 rad of its starting pose and the right gripper is constant. The model's right-arm targets stay within about 2e-4 rad of the current pose, so in practice this is a left-arm policy.
• Small dataset, no held-out evaluation. All 22 episodes were used for training.
• 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.