import dataclasses import logging import einops import flax.nnx as nnx import flax.nnx.bridge as nnx_bridge import jax import jax.numpy as jnp from typing_extensions import override from typing import List, Optional from openpi.models import model as _model import openpi.models.gemma as _gemma import openpi.models.siglip as _siglip from openpi.shared import array_typing as at import openpi.shared.nnx_utils as nnx_utils from openpi.models.pi0moh_config import Pi0GatedConfig def make_attn_mask(input_mask, mask_ar): """Adapted from big_vision. Tokens can attend to valid inputs tokens which have a cumulative mask_ar smaller or equal to theirs. This way `mask_ar` bool[?B, N] can be used to setup several types of attention, for example: [[1 1 1 1 1 1]]: pure causal attention. [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between themselves and the last 3 tokens have a causal attention. The first entry could also be a 1 without changing behaviour. [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a block can attend all previous blocks and all tokens on the same block. Args: input_mask: bool[B, N] true if its part of the input, false if padding. mask_ar: bool[?B, N] mask that's true where previous tokens cannot depend on it and false where it shares the same attention mask as the previous token. """ mask_ar = jnp.broadcast_to(mask_ar, input_mask.shape) cumsum = jnp.cumsum(mask_ar, axis=1) attn_mask = cumsum[:, None, :] <= cumsum[:, :, None] valid_mask = input_mask[:, None, :] * input_mask[:, :, None] return jnp.logical_and(attn_mask, valid_mask) # Copied from pi0.py @at.typecheck def posemb_sincos( pos: at.Real[at.Array, " b"], embedding_dim: int, min_period: float, max_period: float ) -> at.Float[at.Array, "b {embedding_dim}"]: """Computes sine-cosine positional embedding vectors for scalar positions.""" if embedding_dim % 2 != 0: raise ValueError(f"embedding_dim ({embedding_dim}) must be divisible by 2") fraction = jnp.linspace(0.0, 1.0, embedding_dim // 2) period = min_period * (max_period / min_period) ** fraction sinusoid_input = jnp.einsum( "i,j->ij", pos, 1.0 / period * 2 * jnp.pi, precision=jax.lax.Precision.HIGHEST, ) return jnp.concatenate([jnp.sin(sinusoid_input), jnp.cos(sinusoid_input)], axis=-1) class Pi0Gated(_model.BaseModel): def __init__(self, config: Pi0GatedConfig, rngs: nnx.Rngs): # Initialize base model with max_horizon. super().__init__(config.action_dim, config.action_horizon, config.max_token_len) self.config = config self.pi05 = config.pi05 paligemma_config = _gemma.get_config(config.paligemma_variant) action_expert_config = _gemma.get_config(config.action_expert_variant) # TODO: rewrite gemma in NNX. For now, use bridge. llm = nnx_bridge.ToNNX( _gemma.Module( configs=[paligemma_config, action_expert_config], embed_dtype=config.dtype, adarms=config.pi05, ) ) llm.lazy_init( rngs=rngs, method="init", use_adarms=[False, True] if config.pi05 else [False, False], ) img = nnx_bridge.ToNNX( _siglip.Module( num_classes=paligemma_config.width, variant="So400m/14", pool_type="none", scan=True, dtype_mm=config.dtype, ) ) img.lazy_init(next(iter(config.fake_obs().images.values())), train=False, rngs=rngs) # Attribute names must match pi0.py for weight loading. self.PaliGemma = nnx.Dict(llm=llm, img=img) # Shared action input projection. self.action_in_proj = nnx.Linear(config.action_dim, action_expert_config.width, rngs=rngs) if config.pi05: # Pi0.5-style: adaRMS conditioning on timestep. self.time_mlp_in = nnx.Linear(action_expert_config.width, action_expert_config.width, rngs=rngs) self.time_mlp_out = nnx.Linear(action_expert_config.width, action_expert_config.width, rngs=rngs) else: # Pi0-style: state token + action-time MLP (no adaRMS). self.state_proj = nnx.Linear(config.action_dim, action_expert_config.width, rngs=rngs) self.action_time_mlp_in = nnx.Linear( 2 * action_expert_config.width, action_expert_config.width, rngs=rngs, ) self.action_time_mlp_out = nnx.Linear( action_expert_config.width, action_expert_config.width, rngs=rngs, ) self.action_out_proj = nnx.Linear(action_expert_config.width, config.action_dim, rngs=rngs) # Extra gating head for Mixture-of-Horizons. self.gate_out_proj = nnx.Linear(action_expert_config.width, 1, rngs=rngs) @at.typecheck def embed_prefix( self, obs: _model.Observation ) -> tuple[ at.Float[at.Array, "b s emb"], at.Bool[at.Array, "b s"], at.Bool[at.Array, " s"], ]: """Unchanged from pi0.py""" input_mask = [] ar_mask = [] tokens = [] # embed images for name in obs.images: image_tokens, _ = self.PaliGemma.img(obs.images[name], train=False) tokens.append(image_tokens) input_mask.append( einops.repeat( obs.image_masks[name], "b -> b s", s=image_tokens.shape[1], ) ) # image tokens attend to each other ar_mask += [False] * image_tokens.shape[1] # add language (aka tokenized inputs) if obs.tokenized_prompt is not None: tokenized_inputs = self.PaliGemma.llm(obs.tokenized_prompt, method="embed") tokens.append(tokenized_inputs) input_mask.append(obs.tokenized_prompt_mask) # full attention between image and language inputs ar_mask += [False] * tokenized_inputs.shape[1] tokens = jnp.concatenate(tokens, axis=1) input_mask = jnp.concatenate(input_mask, axis=1) ar_mask = jnp.array(ar_mask) return tokens, input_mask, ar_mask @at.typecheck def embed_suffix( self, state, noisy_actions: _model.Actions, timestep: at.Float[at.Array, " b"], action_pad_mask, ) -> tuple[ at.Float[at.Array, "b s emb"], at.Bool[at.Array, "b s"], at.Bool[at.Array, " s"], at.Float[at.Array, "b emb"] | None, ]: """ Pi0 / Pi0.5 compatible suffix embedding. Mirrors :class:`Pi0`'s ``embed_suffix`` (including adaRMS conditioning when ``pi05=True``) but takes ``state`` and ``action_pad_mask`` explicitly to support batched horizon processing used by MoH. """ input_mask = [] ar_mask: list[bool] = [] tokens = [] adarms_cond = None # Optional Pi0-style state token (no state token for Pi0.5 / adaRMS). if not self.pi05: state_token = self.state_proj(state)[:, None, :] tokens.append(state_token) input_mask.append(jnp.ones((state.shape[0], 1), dtype=jnp.bool_)) # image/language inputs do not attend to state or actions ar_mask += [True] # Timestep embedding. time_emb = posemb_sincos( timestep, self.action_in_proj.out_features, min_period=4e-3, max_period=4.0, ) # Project actions. action_tokens = self.action_in_proj(noisy_actions) if self.pi05: # Pi0.5: adaRMS on time embedding, actions unchanged. time_emb = self.time_mlp_in(time_emb) time_emb = nnx.swish(time_emb) time_emb = self.time_mlp_out(time_emb) time_emb = nnx.swish(time_emb) adarms_cond = time_emb action_expert_tokens = action_tokens else: # Pi0: concatenate time + actions and pass through MLP. time_tokens = einops.repeat(time_emb, "b emb -> b s emb", s=noisy_actions.shape[1]) action_time_tokens = jnp.concatenate([action_tokens, time_tokens], axis=-1) action_time_tokens = self.action_time_mlp_in(action_time_tokens) action_time_tokens = nnx.swish(action_time_tokens) action_time_tokens = self.action_time_mlp_out(action_time_tokens) action_expert_tokens = action_time_tokens if action_pad_mask is None: action_pad_mask = jnp.ones(action_expert_tokens.shape[:2], dtype=jnp.bool_) input_mask.append(action_pad_mask) tokens.append(action_expert_tokens) # image/language/state inputs do not attend to action tokens ar_mask += [True] + ([False] * (action_expert_tokens.shape[1] - 1)) tokens = jnp.concatenate(tokens, axis=1) input_mask_arr = jnp.concatenate(input_mask, axis=1) ar_mask_arr = jnp.array(ar_mask) return tokens, input_mask_arr, ar_mask_arr, adarms_cond def cv_squared(self, x: at.Array, eps: float = 1e-10) -> at.Array: """Computes the squared coefficient of variation. (From pi0_pytorch_moh.py)""" def compute_cv(): mean = jnp.mean(x, dtype=jnp.float32) var = jnp.var(x, dtype=jnp.float32) return var / (jnp.square(mean) + eps) # Handle num_experts = 1 case return jax.lax.cond( x.shape[0] == 1, lambda: jnp.array(0.0, dtype=jnp.float32), compute_cv ) @override def compute_loss( self, rng: at.KeyArrayLike, observation: _model.Observation, actions: _model.Actions, *, train: bool = False ) -> at.Float[at.Array, "*b ah"]: # def compute_loss( # self, # rng: at.KeyArrayLike, # observation: _model.Observation, # actions: at.Float[at.Array, "b s action_dim"], # ) -> tuple[at.Float[at.Array, ""], dict[str, at.Float[at.Array, ""]]]: preprocess_rng, noise_rng, time_rng = jax.random.split(rng, 3) observation = _model.preprocess_observation(preprocess_rng, observation, train=train) batch_size, max_horizon, action_dim = actions.shape num_horizons = len(self.config.horizons) horizons_arr = jnp.array(self.config.horizons) # Sample noise and time noise = jax.random.normal(noise_rng, actions.shape) time_scalar = jax.random.beta(time_rng, 1.5, 1, (batch_size,)) * 0.999 + 0.001 # Expand time and actions for each horizon # time shape: (H, B) time = einops.repeat(time_scalar, "b -> h b", h=num_horizons) # x_t shape: (H, B, max_H, D) x_t = time[..., None, None] * noise[None, ...] + (1 - time[..., None, None]) * actions[None, ...] # u_t (target) shape: (B, max_H, D) u_t = noise - actions # STAGE 1: VLM Prefix Pass (Compute KV cache once) prefix_tokens, prefix_mask, prefix_ar_mask = self.embed_prefix(observation) prefix_attn_mask = make_attn_mask(prefix_mask, prefix_ar_mask) prefix_positions = jnp.cumsum(prefix_mask, axis=1) - 1 (_, prefix_out), prefix_past_key_values = self.PaliGemma.llm( [prefix_tokens, None], mask=prefix_attn_mask, positions=prefix_positions ) # STAGE 2: Action Head Suffix Passes (Parallelized via batching) # Repeat prefix masks and KV cache for each horizon # New batch size is (B * H) batched_prefix_mask = jnp.repeat(prefix_mask, num_horizons, axis=0) batched_past_key_values = jax.tree_map( lambda x: jnp.repeat(x, num_horizons, axis=1), prefix_past_key_values ) batched_state = jnp.repeat(observation.state, num_horizons, axis=0) # Reshape x_t and time to align with the new batch dimension # (H, B, max_H, D) -> (B*H, max_H, D) batched_x_t = jnp.transpose(x_t, (1, 0, 2, 3)).reshape(batch_size * num_horizons, max_horizon, -1) # (H, B) -> (B*H,) batched_time = jnp.transpose(time, (1, 0)).reshape(-1) # Create a padding mask for actions based on valid horizon length # (H, max_H) action_pad_mask = jnp.arange(max_horizon)[None, :] < horizons_arr[:, None] # (B*H, max_H) action_pad_mask_expanded = jnp.broadcast_to( action_pad_mask[None, :, :], (batch_size, num_horizons, max_horizon) ) # (B, H, max_H) -> (B*H, max_H) batched_action_pad_mask = action_pad_mask_expanded.reshape(batch_size * num_horizons, max_horizon) # Embed the batched suffix inputs suffix_tokens, suffix_pad_masks, suffix_ar_mask, adarms_cond = self.embed_suffix( batched_state, batched_x_t, batched_time, action_pad_mask=batched_action_pad_mask ) # Combine prefix and suffix masks for cross-attention pad_masks = jnp.concatenate([batched_prefix_mask, suffix_pad_masks], axis=1) ar_masks = jnp.concatenate([prefix_ar_mask, suffix_ar_mask], axis=0) full_att_2d_masks = make_attn_mask(pad_masks, ar_masks) prefix_len = prefix_mask.shape[1] suffix_len = suffix_tokens.shape[1] # Create position IDs and attention mask for the suffix part only suffix_position_ids = jnp.arange(prefix_len, prefix_len + suffix_len)[None, :] suffix_att_2d_masks = full_att_2d_masks[:, -suffix_len:, :] b = suffix_att_2d_masks.shape[0] suffix_position_ids = jnp.broadcast_to(suffix_position_ids, (b, suffix_len)) adarms = [None, adarms_cond] if self.pi05 else [None, None] (_, suffix_out), _ = self.PaliGemma.llm( [None, suffix_tokens], mask=suffix_att_2d_masks, positions=suffix_position_ids, kv_cache=batched_past_key_values, adarms_cond=adarms, ) action_start_index = 0 if self.pi05 else 1 # pi0.5 has no state token v_t_batched = self.action_out_proj(suffix_out) v_t_actions_padded = v_t_batched[:, action_start_index: action_start_index + max_horizon, :] # (H, B, max_H, D_action) all_v_t_preds = v_t_actions_padded.reshape( batch_size, num_horizons, max_horizon, -1 ).transpose(1, 0, 2, 3) # 1. Primary Loss: Ensures each expert head is trained well. all_head_losses = [] for i, h in enumerate(self.config.horizons): v_t_head = all_v_t_preds[i, :, :h, :] target_v_t = u_t[:, :h, :] # Mean over batch, horizon, and action dim head_loss = jnp.mean(jnp.square(v_t_head - target_v_t)) all_head_losses.append(head_loss) individual_loss = jnp.sum(jnp.stack(all_head_losses)) # 2. Auxiliary Loss: Trains the gating network # (B*H, S_suffix, 1) gate_logits_batched = self.gate_out_proj(suffix_out) # (B*H, max_H, 1) gate_logits_padded = gate_logits_batched[:, action_start_index: action_start_index + max_horizon, :] # (B, max_H, H) # gate_logits = einops.rearrange(gate_logits_padded, "(b h) s 1 -> b s h", b=batch_size, h=num_horizons) gate_logits_reshaped = gate_logits_padded.reshape(batch_size, num_horizons, max_horizon, 1) gate_logits = jnp.transpose(gate_logits_reshaped, (0, 2, 1, 3)).squeeze(-1) # Create mask for softmax # (max_H, H) valid_heads_mask = jnp.arange(max_horizon)[:, None] < horizons_arr[None, :] # (B, max_H, H) - broadcast batch dim valid_heads_mask = jnp.broadcast_to(valid_heads_mask, gate_logits.shape) masked_gate_logits = jnp.where(valid_heads_mask, gate_logits, jnp.finfo(gate_logits.dtype).min) gate_weights = nnx.softmax(masked_gate_logits, axis=-1) # Combine predictions using the dynamic weights # all_v_t_preds: (H, B, max_H, D) -> (B, H, max_H, D) all_v_t_preds_permuted = einops.rearrange(all_v_t_preds, "h b s d -> b h s d") # gate_weights: (B, max_H, H) -> (B, H, max_H, 1) gate_weights_expanded = jnp.transpose(gate_weights, (0, 2, 1))[:, :, :, None] # (B, H, max_H, D) * (B, H, max_H, 1) -> sum over H -> (B, max_H, D) v_t_combined = jnp.sum(all_v_t_preds_permuted * gate_weights_expanded, axis=1) auxiliary_loss = jnp.mean(jnp.square(v_t_combined - u_t)) # Mean over B, H, D # 3. Balance Loss: Encourage the gate layer to output weights flexibly loss_components = [] boundaries = sorted(list(set([0] + self.config.horizons))) for i in range(len(boundaries) - 1): start_step, end_step = boundaries[i], boundaries[i + 1] active_expert_indices = [idx for idx, h in enumerate(self.config.horizons) if h > start_step] if len(active_expert_indices) > 1: # (B, S_segment, H_total) segment_gate_weights = gate_weights[:, start_step:end_step, :] # (B, S_segment, H_active) active_expert_weights = segment_gate_weights[:, :, jnp.array(active_expert_indices)] # (H_active,) avg_expert_prob_in_segment = jnp.mean(active_expert_weights, axis=(0, 1)) segment_loss = self.cv_squared(avg_expert_prob_in_segment) loss_components.append(segment_loss) load_balancing_loss = jnp.mean(jnp.stack(loss_components)) if loss_components else 0.0 total_loss = ( individual_loss + self.config.aux_weight * auxiliary_loss + self.config.balance_weight * load_balancing_loss ) return total_loss @override def sample_actions( self, rng: at.KeyArrayLike, observation: _model.Observation, *, num_steps: int | at.Int[at.Array, ""] = 10, noise: at.Float[at.Array, "b ah ad"] | None = None, ) -> _model.Actions: """ Samples actions using the gated fusion mechanism during denoising. """ observation = _model.preprocess_observation(None, observation, train=False) dt = -1.0 / num_steps batch_size = observation.state.shape[0] max_horizon = self.action_horizon num_horizons = len(self.config.horizons) horizons_arr = jnp.array(self.config.horizons) noise = jax.random.normal(rng, (batch_size, max_horizon, self.action_dim)) # First fill KV cache with a forward pass of the prefix prefix_tokens, prefix_mask, prefix_ar_mask = self.embed_prefix(observation) prefix_attn_mask = make_attn_mask(prefix_mask, prefix_ar_mask) positions = jnp.cumsum(prefix_mask, axis=1) - 1 (_, prefix_out), kv_cache = self.PaliGemma.llm( [prefix_tokens, None], mask=prefix_attn_mask, positions=positions ) # Prepare static batched inputs (these don't change in the loop) batched_prefix_mask = jnp.repeat(prefix_mask, num_horizons, axis=0) batched_kv_cache = jax.tree_map( lambda x: jnp.repeat(x, num_horizons, axis=1), kv_cache ) batched_state = jnp.repeat(observation.state, num_horizons, axis=0) # Create static masks for padding actions in the loop # (H, max_H) steps_arr = jnp.arange(max_horizon) action_pad_mask_per_h = steps_arr[None, :] < horizons_arr[:, None] # (B*H, max_H) batched_action_pad_mask = jnp.broadcast_to( action_pad_mask_per_h[None, :, :], (batch_size, num_horizons, max_horizon) ) batched_action_pad_mask = einops.rearrange(batched_action_pad_mask, "b h s -> (b h) s") # batched_action_pad_mask = einops.repeat(action_pad_mask_per_h, "h s -> (b h) s", b=batch_size) # (H, max_H, 1) action_mask_for_padding_x_t = (steps_arr[None, :, None] < horizons_arr[:, None, None]) # Create static mask for gate softmax # (max_H, H) valid_heads_mask = steps_arr[:, None] < horizons_arr[None, :] # (B, max_H, H) - for broadcasting valid_heads_mask = valid_heads_mask[None, :, :] action_start_index = 0 if self.pi05 else 1 # pi0.5 has no state token prefix_len = prefix_mask.shape[1] def step_fn(carry): x_t, time = carry # --- Prepare Batched Inputs for this step --- expanded_time = jnp.broadcast_to(time, (batch_size * num_horizons,)) # Pad x_t for each horizon # (1, B, max_H, D) x_t_expanded = x_t[None, ...] # (H, B, max_H, D) padded_x_t_batched = jnp.where(action_mask_for_padding_x_t, x_t_expanded, 0.0) # (B*H, max_H, D) batched_x_t = padded_x_t_batched.transpose(1, 0, 2, 3) batched_x_t = einops.rearrange(batched_x_t, "b h s d -> (b h) s d") # --- Run Batched Suffix Pass --- suffix_tokens, suffix_pad_masks, suffix_ar_mask, adarms_cond = self.embed_suffix( batched_state, batched_x_t, expanded_time, action_pad_mask=batched_action_pad_mask ) pad_masks = jnp.concatenate([batched_prefix_mask, suffix_pad_masks], axis=1) ar_masks = jnp.concatenate([prefix_ar_mask, suffix_ar_mask], axis=0) full_att_2d_masks = make_attn_mask(pad_masks, ar_masks) suffix_len = suffix_tokens.shape[1] suffix_position_ids = jnp.arange(prefix_len, prefix_len + suffix_len)[None, :] suffix_att_2d_masks = full_att_2d_masks[:, -suffix_len:, :] b = suffix_att_2d_masks.shape[0] suffix_position_ids = jnp.broadcast_to(suffix_position_ids, (b, suffix_len)) adarms = [None, adarms_cond] if self.pi05 else [None, None] (_, suffix_out), _ = self.PaliGemma.llm( [None, suffix_tokens], mask=suffix_att_2d_masks, positions=suffix_position_ids, kv_cache=batched_kv_cache, adarms_cond=adarms, ) # --- Gating and Fusion --- # (B*H, S_suffix, 1) gate_logits_batched = self.gate_out_proj(suffix_out) # (B*H, max_H, 1) gate_logits_padded = gate_logits_batched[:, action_start_index: action_start_index + max_horizon, :] # (B, max_H, H) gate_logits_reshaped = gate_logits_padded.reshape(batch_size, num_horizons, max_horizon, 1) gate_logits = jnp.transpose(gate_logits_reshaped, (0, 2, 1, 3)).squeeze(-1) masked_gate_logits = jnp.where(valid_heads_mask, gate_logits, jnp.finfo(gate_logits.dtype).min) gate_weights = nnx.softmax(masked_gate_logits, axis=-1) # Get all predictions # (B*H, S_suffix, D_action) v_t_batched = self.action_out_proj(suffix_out) # (B*H, max_H, D_action) v_t_actions_padded = v_t_batched[:, action_start_index: action_start_index + max_horizon, :] # (B, H, max_H, D_action) all_v_t_preds = v_t_actions_padded.reshape(batch_size, num_horizons, max_horizon, -1) # Combine predictions # gate_weights: (B, max_H, H) -> (B, H, max_H, 1) gate_weights_expanded = jnp.transpose(gate_weights, (0, 2, 1))[:, :, :, None] # (B, H, max_H, D) * (B, H, max_H, 1) -> sum over H -> (B, max_H, D) v_t = jnp.sum(all_v_t_preds * gate_weights_expanded, axis=1) # --- Euler Step --- x_t_new = x_t + dt * v_t time_new = time + dt return (x_t_new, time_new) def cond_fn(carry): x_t, time = carry # robust to floating-point error return time >= -dt / 2 x_0, _ = jax.lax.while_loop(cond_fn, step_fn, (noise, 1.0)) return x_0