vllm_omni.diffusion.models.pi05 ¶

Encode images with SigLIP vision tower + PaliGemma projector.

The two steps are run explicitly rather than via PaliGemmaModel.get_image_features because that helper divides the projector output by sqrt(text_hidden_size). Being explicit keeps the scale unambiguous and matches π0 exactly (SigLIP is unchanged in π0.5).

Embed language tokens, returning the * sqrt(hidden)-scaled embedding.

The scaling location moved across transformers releases: at ≤5.3 the normalizer lives inside GemmaModel.forward (which we bypass), and at ≥5.4 GemmaTextScaledWordEmbedding self-applies it. Detect and avoid double-scaling.

Dispatch to prefix_only / suffix_only and return ([prefix_out, suffix_out], past_key_values_or_None).

Return (normed, gate); gate is None in the unconditioned case.

Build from a config dict (LeRobot config.json or deploy yaml).

Build from a checkpoint directory's config.json.

Normalization stats are not in config.json. LeRobot keeps them in the processor sidecar, so they are loaded separately and backfilled here — see :func:load_lerobot_norm_stats.

Apply one flow-matching denoising step: predict v_t from x_t.

Signature differs from π0's by exactly one argument: no state.

Build prefix embeddings, per-token padding mask, and AR mask.

Prefix tokens form [img_cam_0..., ..., lang_tokens...] with fully bidirectional attention (all-zero att_masks). Identical to π0 — the state is inside lang_tokens, so nothing here changes shape-wise.

Cameras are embedded one at a time; each call is (B, 3, 224, 224). The number of slots is fixed by config.max_cameras for the deployed model; missing cameras occupy their slot with a false image mask.

Build the suffix: action tokens only, plus the AdaRMS condition.

π0's suffix is [state_token, action_tokens×H] with an AR mask of [1, 1, 0...]. π0.5 has no state token, so the suffix is [action_tokens×H] and the mask is [1] + [0]*(H-1): the first action token opens a causal block and the rest attend bidirectionally within it.

Timestep → AdaRMS conditioning vector (B, expert_width).

silu(time_mlp_out(silu(time_mlp_in(sinusoid(t))))). The trailing SiLU is part of the reference implementation — dropping it is a silent numerical error, not a crash.

Load and audit a LeRobot π0.5 safetensors checkpoint.

Same remap rules as π0 (strip the model. prefix, flatten→nested PaliGemma submodules, tied lm_head → embed_tokens, version-robust SigLIP nesting), plus two π0.5-specific ones:

• action_time_mlp_{in,out} → time_mlp_{in,out}: some checkpoints were exported under the π0 parameter names.
• state_proj.* is reported, not silently dropped. A π0.5 checkpoint should not contain it; its presence usually means a π0 checkpoint was pointed at the π0.5 model class, which would otherwise run happily with a randomly-initialized action expert.

The action-expert norms are AdaRMS here, so they expose dense.weight / dense.bias and no plain weight. A checkpoint that carries a plain expert-norm weight is a π0-shaped checkpoint; that too is rejected rather than skipped. strict=False exists only for focused remapping unit tests that intentionally provide a partial state dict; the serving path always uses the strict default.

Generate an action chunk via iterative flow-matching denoising.

Convention: t=1 is noise, t=0 is the target — opposite of the published π0 paper but matching both OpenPI and LeRobot.

Takes no state: π0.5's state rides inside lang_tokens.

Fill tensor for an unused camera slot — pure -1, matches OpenPI/LeRobot.

Convert one LeRobot-domain image to a normalized model tensor.

PIL and uint8 inputs have domain [0, 255]. Floating ndarray/tensor inputs must be finite and in [0, 1]. Values are converted deterministically to [-1, 1]; the pixel content is never used to guess its input domain.

No-op for the diffusion loader: π0.5 self-loads its checkpoint.

Raw observation → sample_actions inputs.

No state tensor comes back: π0.5 carries the state inside lang_tokens.

Model output → the action chunk the robot receives.

AbsoluteActionsProcessorStep runs after unnormalization, because a relative-action checkpoint's statistics are computed in relative space. The raw state is the one the prompt encoded, before normalization.

relative → absolute. Output side (step 2 of the post-pipeline).

state must be the raw state — the same one the model was given before normalization — because relative actions live in raw action space.

absolute → relative. Input side (step 3).

Not used on the inference path — there are no input actions to convert at serving time — but it is what the transform means, and the parity test exercises it as the inverse of :meth:to_absolute.