Inference-only π0.5 VLA math kernel for vllm-omni.
Only the math that turns a robot observation into an action chunk; no serving or request glue. Deliberately shaped like models/pi0/modeling_pi0.py so the two can later be factored into a shared Pi-family module (RFC step 2) on a "behaviour unchanged" review.
π0.5 = PaliGemma (SigLIP vision + Gemma 2B LM) prefix + Gemma 300M action expert suffix + flow-matching head. Where π0 projects the robot state through state_proj, π0.5 discretizes it into prompt tokens — so there is no state_proj layer, sample_actions takes no state argument, and the suffix is action tokens only, which drops π0's leading state-token boundary from the suffix attention mask and leaves [1] + [0] * (horizon - 1).
π0.5 is nonetheless the larger model: the 37 AdaRMS dense projections add ~116M parameters against the ~8K that state_proj saves. (LeRobot's README says otherwise; the checkpoint disagrees.)
Reference implementations
- OpenPI: openpi/src/openpi/models_pytorch/pi0_pytorch.py, gemma_pytorch.py
- LeRobot: lerobot/src/lerobot/policies/pi05/modeling_pi05.py
Bases: Module
Dual-backbone transformer: PaliGemma (Gemma 2B) + AdaRMS expert (300M).
Same two-mode dispatch as π0 (prefix_only / suffix_only), with one structural change: after building a stock GemmaForCausalLM expert, every norm in it is swapped for a :class:Pi05AdaRMSNorm carrying a dense conditioning projection.
Swapping in place — rather than subclassing GemmaModel as #4419 does — keeps the module tree, and therefore the checkpoint key layout, identical to the expert's stock layout apart from the norms themselves.
instance-attribute ¶
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.
forward(
attention_mask: Tensor | None = None,
position_ids: LongTensor | None = None,
past_key_values=None,
inputs_embeds: list[Tensor | None] | None = None,
use_cache: bool = False,
adarms_cond: Tensor | None = None,
)
Dispatch to prefix_only / suffix_only and return ([prefix_out, suffix_out], past_key_values_or_None).
Bases: Module
Adaptive RMSNorm conditioned on the flow-matching timestep.
π0 conditions on time by concatenating a time embedding onto each action embedding. π0.5 instead feeds the time embedding into every action-expert norm, which produces a per-layer (scale, shift, gate) triple::
y = norm(x) * (1 + scale) + shift
out = residual + gate * sublayer(y)
dense is zero-initialized, so an untrained model starts as the identity modulation with a closed gate — matching OpenPI's parameterization.
Note the shape of the unconditioned branch: normed * (1 + weight) with weight zero-initialized, which is exactly transformers' GemmaRMSNorm. That equivalence is why only the expert norms need replacing here and the PaliGemma prefix can keep stock Gemma layers.
Bases: Module
π0.5 VLA model for robot action prediction via flow matching.
Inference flow
- Embed prefix (images + language, where the language already carries the discretized state) → prefix tokens.
- Forward prefix through PaliGemma → layer-wise KV cache.
- For each denoising step
t = 1.0, 1-dt, ..., 0: a. Embed the timestep → an AdaRMS conditioning vector. b. Embed the suffix (action tokens only). c. Forward the suffix through the AdaRMS action expert. d.x_t = x_t + dt * v_t(Euler integration). - Return
x_0as the predicted action chunk.
instance-attribute ¶
action_dim = getattr(
config, "max_action_dim", DEFAULT_ACTION_DIM
)
instance-attribute ¶
action_horizon = getattr(
config, "chunk_size", DEFAULT_ACTION_HORIZON
)
instance-attribute ¶
instance-attribute ¶
instance-attribute ¶
max_state_dim = getattr(
config, "max_state_dim", self.action_dim
)
instance-attribute ¶
num_inference_steps = getattr(
config,
"num_inference_steps",
DEFAULT_NUM_INFERENCE_STEPS,
)
instance-attribute ¶
paligemma_with_expert = PaliGemmaWithActionExpertPi05(
vlm_config, expert_config
)
Apply one flow-matching denoising step: predict v_t from x_t.
Signature differs from π0's by exactly one argument: no state.
embed_prefix(
images: list[Tensor],
image_masks: list[Tensor],
lang_tokens: Tensor,
lang_masks: Tensor,
) -> tuple[Tensor, Tensor, Tensor]
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.
embed_suffix(
noisy_actions: Tensor, timestep: Tensor
) -> tuple[Tensor, Tensor, Tensor, Tensor]
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.
sample_actions(
images: list[Tensor],
image_masks: list[Tensor],
lang_tokens: Tensor,
lang_masks: Tensor,
noise: Tensor | None = None,
num_steps: int | None = None,
generator: Generator | list[Generator] | None = None,
) -> Tensor
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.
create_sinusoidal_pos_embedding(
time: Tensor,
dimension: int,
min_period: float = 0.004,
max_period: float = 4.0,
device: device = None,
) -> Tensor
Compute a sine/cosine positional embedding for scalar timesteps.
Ref: openpi/models_pytorch/pi0_pytorch.py create_sinusoidal_pos_embedding
Build a 2D attention mask from a padding mask and an autoregressive mask.
Ref: openpi/models_pytorch/pi0_pytorch.py make_att_2d_masks
Convert (B, S, S) bool masks to (B, 1, S, S) float masks.
True → 0.0 (attend), False → OPENPI_ATTENTION_MASK_VALUE.