Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsPraxisnemo_automodel.components.models.praxis.action_expert

nemo_automodel.components.models.praxis.action_expert

View as MarkdownOpen in Claude

Action expert: a small transformer that denoises the action tokens.

The expert reads the encoder’s input sequence (its global tokens) and the encoder block outputs at the tap blocks; it never writes back into the video stream. Its layers alternate cross attention (even layers, one tap each) and self attention (odd layers).

Expert query sequence::

[ proprio (1 per chunk) | future tokens (F per chunk) | action tokens (1 per noised frame) ]

Cross-attention key/value sequence::

[ global tokens (register, anchor, text) | clean frame tokens | noised frame tokens ]

Queries and keys turn by the trunk rotary positions before their RMS norm; the rotary position of a frame is the position of its first token.

Module Contents

Classes

NameDescription
ActionExpertThe action transformer (see the module docstring).
ActionExpertBlockOne expert layer: cross or self attention, then a GEGLU MLP, each behind a scale-only AdaLN.
ActionTauEncoderAction token embedding with the flow time at the input: w3(silu(w2([w1(a); tau(t)]))).
ScaleAdaLNScale-only adaptive RMS norm: (1 + proj(silu(cond))) * rms_norm(x).
SinusoidalTimestep[cos(t f), sin(t f)] features of the flow time followed by a two-layer MLP.

Functions

NameDescription
_query_chunksChunk index of every expert query in the [proprio | future | actions] layout.
gegluGated GELU: gelu_tanh(first half) * second half over the last axis.
rms_normWeightless RMS norm over the last axis; the mean and scale run in float32.

API

class nemo_automodel.components.models.praxis.action_expert.ActionExpert(
    cfg: nemo_automodel.components.models.praxis.config.ActionExpertConfig
)

Bases: Module

The action transformer (see the module docstring).

_masks
= MaskCache()
action_encoder
= ActionTauEncoder(cfg.action_patch_dim, cfg.dim)
blocks
dim
= cfg.dim
embed_proprio
= nn.Linear(cfg.proprio_patch_dim, cfg.dim)
future_chunk_embed
= nn.Embedding(cfg.max_chunks, cfg.dim)
future_tokens
head
= nn.Linear(cfg.dim, cfg.action_patch_dim)
norm_out
num_future
= cfg.num_future_per_chunk
proprio_token
= nn.Parameter(torch.zeros(1, cfg.dim))
rope
tap_norm_scale
tap_proj
= nn.Linear(cfg.trunk_dim, cfg.dim, bias=False)
text_proj
= nn.Linear(cfg.trunk_dim, cfg.dim, bias=False)
timestep_embedder
= SinusoidalTimestep(cfg.dim)
nemo_automodel.components.models.praxis.action_expert.ActionExpert._masks_for(
    q_tags: nemo_automodel.components.models.praxis.masks.TokenTags,
    kv_tags: nemo_automodel.components.models.praxis.masks.TokenTags,
    window: int,
    key: tuple[int, int, int, int, str]
) -> tuple[torch.Tensor, torch.Tensor]
nemo_automodel.components.models.praxis.action_expert.ActionExpert.forward(
    features: nemo_automodel.components.models.praxis.backbone.EncoderFeatures,
    noised_actions: torch.Tensor,
    action_timesteps: torch.Tensor,
    proprio: torch.Tensor
) -> torch.Tensor

Predict the action velocity from the encoder features of one window.

Parameters:

features
EncoderFeatures

Encoder taps and sequence layout of the window (batch size 1).

noised_actions
Tensor

Noised action tokens of shape [1, K * F, action_patch_dim], chunk-major.

action_timesteps
Tensor

Flow time of each chunk, shape [1, K].

proprio
Tensor

Normalized boundary pose of each chunk, shape [1, K, proprio_patch_dim].

Returns: Tensor

Action velocity of shape [1, K * F, action_patch_dim] in the taps dtype.

class nemo_automodel.components.models.praxis.action_expert.ActionExpertBlock(
    dim: int,
    num_heads: int,
    head_dim: int,
    is_cross: bool
)

Bases: Module

One expert layer: cross or self attention, then a GEGLU MLP, each behind a scale-only AdaLN.

adaln_attn
= ScaleAdaLN(dim)
adaln_mlp
= ScaleAdaLN(dim)
mlp_in
= nn.Linear(dim, 8 * dim)
mlp_out
= nn.Linear(4 * dim, dim)
to_kv
= nn.Linear(dim, 2 * inner, bias=False)
to_out
= nn.Linear(inner, dim, bias=False)
to_q
= nn.Linear(dim, inner, bias=False)
nemo_automodel.components.models.praxis.action_expert.ActionExpertBlock.forward(
    x: torch.Tensor,
    cond: torch.Tensor,
    kv_source: torch.Tensor | None,
    mask: torch.Tensor,
    q_table: torch.Tensor,
    k_table: torch.Tensor
) -> torch.Tensor

Run the layer.

Parameters:

x
Tensor

Expert tokens of shape [batch, queries, hidden].

cond
Tensor

Timestep conditioning of shape [batch, queries, hidden].

kv_source
Tensor | None

Cross-attention source of shape [batch, keys, hidden], or None for self attention.

mask
Tensor

Boolean tensor of shape [queries, keys], true where attention is allowed.

q_table
Tensor

Rotary table of the queries, shape [batch, queries, head_dim // 2, 2].

k_table
Tensor

Rotary table of the keys, shape [batch, keys, head_dim // 2, 2].

Returns: Tensor

Tensor of shape [batch, queries, hidden].

class nemo_automodel.components.models.praxis.action_expert.ActionTauEncoder(
    action_patch_dim: int,
    dim: int
)

Bases: Module

Action token embedding with the flow time at the input: w3(silu(w2([w1(a); tau(t)]))).

tau
= SinusoidalTimestep(dim)
w1
= nn.Linear(action_patch_dim, dim)
w2
= nn.Linear(2 * dim, dim)
w3
= nn.Linear(dim, dim)
nemo_automodel.components.models.praxis.action_expert.ActionTauEncoder.forward(
    actions: torch.Tensor,
    tau: torch.Tensor
) -> torch.Tensor

Embed noised action tokens.

Parameters:

actions
Tensor

Tensor of shape [batch, action_tokens, action_patch_dim].

tau
Tensor

Flow time per token, shape [batch, action_tokens].

Returns: Tensor

Tensor of shape [batch, action_tokens, hidden].

class nemo_automodel.components.models.praxis.action_expert.ScaleAdaLN(
    dim: int
)

Bases: Module

Scale-only adaptive RMS norm: (1 + proj(silu(cond))) * rms_norm(x).

proj
= nn.Linear(dim, dim)
nemo_automodel.components.models.praxis.action_expert.ScaleAdaLN.forward(
    tokens: torch.Tensor,
    cond: torch.Tensor
) -> torch.Tensor

Normalize and scale.

Parameters:

tokens
Tensor

Tensor of shape [batch, queries, hidden].

cond
Tensor

Tensor of shape [batch, queries, hidden].

Returns: Tensor

Tensor of shape [batch, queries, hidden] in the dtype of tokens.

class nemo_automodel.components.models.praxis.action_expert.SinusoidalTimestep(
    dim: int,
    frequency_dim: int = 256,
    max_period: float = 10000.0
)

Bases: Module

[cos(t f), sin(t f)] features of the flow time followed by a two-layer MLP.

mlp
nemo_automodel.components.models.praxis.action_expert.SinusoidalTimestep.forward(
    t: torch.Tensor
) -> torch.Tensor

Embed timesteps.

Parameters:

t
Tensor

Float tensor of shape […].

Returns: Tensor

Tensor of shape […, dim] in the MLP weight dtype.

nemo_automodel.components.models.praxis.action_expert._query_chunks(
    num_action_tokens: int,
    num_future: int,
    chunk_size: int,
    device: torch.device
) -> torch.Tensor

Chunk index of every expert query in the [proprio | future | actions] layout.

nemo_automodel.components.models.praxis.action_expert.geglu(
    x: torch.Tensor
) -> torch.Tensor

Gated GELU: gelu_tanh(first half) * second half over the last axis.

nemo_automodel.components.models.praxis.action_expert.rms_norm(
    x: torch.Tensor,
    eps: float = 0.0001
) -> torch.Tensor

Weightless RMS norm over the last axis; the mean and scale run in float32.

Parameters:

x
Tensor

Tensor of shape […, hidden].

Returns: Tensor

Tensor of shape […, hidden] in the dtype of x.