> This page is for version Nightly (default).
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.

# nemo_automodel.components.models.praxis.action_expert

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

| Name                                                                                              | Description                                                                                  |
| ------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------- |
| [`ActionExpert`](#nemo_automodel-components-models-praxis-action_expert-ActionExpert)             | The action transformer (see the module docstring).                                           |
| [`ActionExpertBlock`](#nemo_automodel-components-models-praxis-action_expert-ActionExpertBlock)   | One expert layer: cross or self attention, then a GEGLU MLP, each behind a scale-only AdaLN. |
| [`ActionTauEncoder`](#nemo_automodel-components-models-praxis-action_expert-ActionTauEncoder)     | Action token embedding with the flow time at the input: `w3(silu(w2([w1(a); tau(t)])))`.     |
| [`ScaleAdaLN`](#nemo_automodel-components-models-praxis-action_expert-ScaleAdaLN)                 | Scale-only adaptive RMS norm: `(1 + proj(silu(cond))) * rms_norm(x)`.                        |
| [`SinusoidalTimestep`](#nemo_automodel-components-models-praxis-action_expert-SinusoidalTimestep) | `[cos(t f), sin(t f)]` features of the flow time followed by a two-layer MLP.                |

### Functions

| Name                                                                                    | Description                                                                       |
| --------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------- |
| [`_query_chunks`](#nemo_automodel-components-models-praxis-action_expert-_query_chunks) | Chunk index of every expert query in the `[proprio \| future \| actions]` layout. |
| [`geglu`](#nemo_automodel-components-models-praxis-action_expert-geglu)                 | Gated GELU: `gelu_tanh(first half) * second half` over the last axis.             |
| [`rms_norm`](#nemo_automodel-components-models-praxis-action_expert-rms_norm)           | Weightless RMS norm over the last axis; the mean and scale run in float32.        |

### API

```python
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)`

---

```python
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]
```

```python
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.

```python
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)`

---

```python
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].

```python
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)`

---

```python
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].

```python
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)`

---

```python
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`.

```python
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`**

---

```python
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.

```python
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.

```python
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.

```python
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`.