mstar.model.pi05.components.action_expert

Action expert transformer with adaRMS timestep conditioning.

The action expert is a Gemma-style transformer that processes action suffix tokens during the flow-matching loop. It shares KV-cache dimensions with the PaliGemma expert (same num_kv_heads and head_dim) so it can attend to the prefix KV cache that PaliGemma wrote during the prefill walk.

Composed from mstar.model.components — AdaRMSNorm for the conditional norms, GatedDecoderLayer for the gated-residual block, and the parallel-aware ParallelAttention / ParallelGatedMLP for the inner blocks (with a trivial single-rank comm group for the non-TP case). The model-specific pieces left here are Pi05TimeMLP (the sincos → adaRMS conditioning MLP) and the stack assembly.

Functions

Pi05ActionExpertLayer(config)

One action expert decoder layer.

Classes

Pi05ActionExpert(config)	Stack of action expert layers plus a final adaRMS norm.
Pi05TimeMLP(hidden_size)	Two-layer SiLU MLP that maps the sincos timestep embedding to the adarms_cond vector consumed by every norm in the action expert.

class mstar.model.pi05.components.action_expert.Pi05ActionExpert(config)

Bases: Module

Stack of action expert layers plus a final adaRMS norm.

Operates entirely in config.action_hidden_size.

Parameters:

config (Pi05Config)

forward(query_sequence, adarms_cond, *, label)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

• query_sequence (Tensor)

• adarms_cond (Tensor)

• label (str)

Return type:

Tensor

mstar.model.pi05.components.action_expert.Pi05ActionExpertLayer(config)

One action expert decoder layer.

Operates in config.action_hidden_size (1024 for gemma_300m, 2048 for gemma_2b). The attention’s K/V dims still match PaliGemma’s so the action expert can attend to the prefix KV cache PaliGemma wrote.

Parameters:

config (Pi05Config)

Return type:

GatedDecoderLayer

class mstar.model.pi05.components.action_expert.Pi05TimeMLP(hidden_size)

Bases: Module

Two-layer SiLU MLP that maps the sincos timestep embedding to the adarms_cond vector consumed by every norm in the action expert.

Both layers operate in the action expert’s hidden dimension (which may differ from PaliGemma’s). Mirrors lerobot’s time_mlp_in / time_mlp_out chain: Linear → silu → Linear → silu.

Not a plain shared MLP because of the trailing SiLU — keeping it here.

Parameters:

hidden_size (int)

forward(time_emb)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

time_emb (Tensor)

Return type:

Tensor