> This page is for version 0.5.0 · 26.06.
> 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.model

Praxis world-action policy: the video backbone plus the action expert.

One call denoises a chunk-causal window of `K` chunks: the backbone predicts the video
velocity of the noised latents and hands its encoder taps to the action expert, which
predicts the action velocity. The frozen Wan2.2 VAE and the Qwen text encoder are not part
of this module; their outputs are the inputs here.

## Module Contents

### Classes

| Name                                                                                                  | Description                                                               |
| ----------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- |
| [`CheckpointConfigExporter`](#nemo_automodel-components-models-praxis-model-CheckpointConfigExporter) | Write the source checkpoint's `config.json` next to consolidated weights. |
| [`PraxisOutput`](#nemo_automodel-components-models-praxis-model-PraxisOutput)                         | Outputs of one denoising step.                                            |
| [`PraxisPolicy`](#nemo_automodel-components-models-praxis-model-PraxisPolicy)                         | Video backbone and action expert of a Praxis checkpoint.                  |

### Functions

| Name                                                                                                | Description                                                                                       |
| --------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------- |
| [`_per_layer_moe_training`](#nemo_automodel-components-models-praxis-model-_per_layer_moe_training) | Divide the balance-loss weight by the number of MoE layers.                                       |
| [`max_tokens_per_window`](#nemo_automodel-components-models-praxis-model-max_tokens_per_window)     | The longest token sequence one training window produces in either tower.                          |
| [`praxis_backend`](#nemo_automodel-components-models-praxis-model-praxis_backend)                   | The backend the Praxis checkpoint runs with: torch linears, flex attention, grouped-GEMM experts. |

### Data

[`ModelClass`](#nemo_automodel-components-models-praxis-model-ModelClass)

### API

```python
class nemo_automodel.components.models.praxis.model.CheckpointConfigExporter(
    config_path: pathlib.Path
)
```

Write the source checkpoint's `config.json` next to consolidated weights.

The checkpointer's metadata hook; a Praxis checkpoint carries no tokenizer or generation config, and its
`config.json` is copied verbatim so that the consolidated directory loads in the reference runtime.

```python
nemo_automodel.components.models.praxis.model.CheckpointConfigExporter.save(
    hf_metadata_dir: str,
    tokenizer: transformers.PreTrainedTokenizerBase | transformers.processing_utils.ProcessorMixin | None,
    original_model_path: str | None,
    v4_compatible: bool
) -> None
```

Copy the source `config.json` into `hf_metadata_dir`.

```python
nemo_automodel.components.models.praxis.model.CheckpointConfigExporter.validate(
    tokenizer: transformers.PreTrainedTokenizerBase | transformers.processing_utils.ProcessorMixin | None,
    original_model_path: str | None
) -> None
```

Check that the source config exists before any rank writes.

```python
class nemo_automodel.components.models.praxis.model.PraxisOutput(
    velocity: torch.Tensor | None,
    action_velocity: torch.Tensor,
    features: nemo_automodel.components.models.praxis.backbone.EncoderFeatures
)
```

Dataclass

Outputs of one denoising step.

**`action_velocity`** `Tensor`

Action velocity of shape \[batch, K \* F, action_patch_dim].

---

**`features`** `EncoderFeatures`

The encoder features the action expert read.

---

**`velocity`** `Tensor | None`

Video velocity of shape \[batch, channels, K \* F, h, w], `None` when not decoded.

---

```python
class nemo_automodel.components.models.praxis.model.PraxisPolicy(
    config: nemo_automodel.components.models.praxis.config.PraxisConfig,
    backend: nemo_automodel.components.models.common.utils.BackendConfig | None = None,
    moe_training: nemo_automodel.components.models.praxis.backbone.MoeTrainingSettings | None = None
)
```

**Bases:** [HFCheckpointingMixin](/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/common/hf_checkpointing_mixin#nemo_automodel-components-models-common-hf_checkpointing_mixin-HFCheckpointingMixin), `Module`, [MoEFSDPSyncMixin](/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/moe/fsdp_mixin#nemo_automodel-components-moe-fsdp_mixin-MoEFSDPSyncMixin)

Video backbone and action expert of a Praxis checkpoint.

**`action_expert`** `= ActionExpert(config.action_expert)`

---

**`backbone`**

---

**`device`** `device`

---

**`parallelizer`** `ModelParallelizer = PraxisModelParallelizer()`

---

**`pretrained_path`** `Path | None = None`

---

**`state_dict_adapter`**

---

```python
nemo_automodel.components.models.praxis.model.PraxisPolicy._get_consolidated_hf_metadata_exporter(
    tokenizer: transformers.PreTrainedTokenizerBase | transformers.processing_utils.ProcessorMixin | None,
    original_model_path: str | None
) -> nemo_automodel.components.models.praxis.model.CheckpointConfigExporter
```

The checkpointer's metadata hook: the consolidated directory carries the source checkpoint's config.

```python
nemo_automodel.components.models.praxis.model.PraxisPolicy.forward(
    text: torch.Tensor,
    anchor: torch.Tensor,
    clean: torch.Tensor | None,
    noised: torch.Tensor,
    timesteps: torch.Tensor,
    noised_actions: torch.Tensor,
    action_timesteps: torch.Tensor,
    proprio: torch.Tensor,
    rope_shifts: tuple[int, int, int, int] = (0, 0, 0, 0),
    decode: bool = True,
    backbone_grad: bool = True
) -> nemo_automodel.components.models.praxis.model.PraxisOutput
```

Denoise one window of `K` chunks of `F` latent frames.

**Parameters:**

**`text`** `Tensor`

Text-encoder states of shape \[1, text_tokens, text_dim].

---

**`anchor`** `Tensor`

Anchor latent of shape \[1, channels, 1, h, w].

---

**`clean`** `Tensor | None`

Clean copies of chunks `0 .. K-2` of shape \[1, channels, (K-1) \* F, h, w], or `None`.

---

**`noised`** `Tensor`

Noised chunks of shape \[1, channels, K \* F, h, w].

---

**`timesteps`** `Tensor`

Video flow time per chunk, shape \[1, K].

---

**`noised_actions`** `Tensor`

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

---

**`action_timesteps`** `Tensor`

Action flow time per chunk, shape \[1, K].

---

**`proprio`** `Tensor`

Normalized boundary pose per chunk, shape \[1, K, proprio_patch_dim].

---

**`rope_shifts`** `tuple[int, int, int, int]` — default: (0, 0, 0, 0)

`(encoder anchor, encoder text, encoder video, decoder video)` rotary offsets.

---

**`decode`** `bool` — default: True

Run the decoder tower for the video velocity.

---

**`backbone_grad`** `bool` — default: True

Build the autograd graph through the backbone. `False` runs it under
`torch.no_grad` for runs that train only the action expert.

---

**Returns:** `PraxisOutput`

class:`PraxisOutput`.

```python
nemo_automodel.components.models.praxis.model.PraxisPolicy.from_config(
    config: nemo_automodel.components.models.praxis.config.PraxisConfig,
    backend: nemo_automodel.components.models.common.utils.BackendConfig | None = None
) -> nemo_automodel.components.models.praxis.model.PraxisPolicy
```

classmethod

Build the model with uninitialized weights on the current default device.

```python
nemo_automodel.components.models.praxis.model.PraxisPolicy.from_pretrained(
    pretrained_model_name_or_path: str | pathlib.Path,
    backend: nemo_automodel.components.models.common.utils.BackendConfig | None = None,
    device: str | torch.device = 'cpu',
    moe_training: nemo_automodel.components.models.praxis.backbone.MoeTrainingSettings | None = None
) -> nemo_automodel.components.models.praxis.model.PraxisPolicy
```

classmethod

Build the model from a checkpoint directory or its Hugging Face repo.

With a real `device` the checkpoint tensors are assigned strictly and the model is returned in
eval mode with gradients disabled. With `device="meta"` only the architecture is built; a
training recipe shards it first and then fills it through the checkpointer
(`load_base_model`).

**Parameters:**

**`pretrained_model_name_or_path`** `str | Path`

Checkpoint directory with `config.json`, the safetensors
shards and `model.safetensors.index.json`, or the Hugging Face repo id of one
(`runwayml/praxis-1-droid`), downloaded to the Hub cache.

---

**`backend`** `BackendConfig | None` — default: None

Backend selection; defaults to `praxis_backend`.

---

**`device`** `str | torch.device` — default: 'cpu'

Device of the loaded tensors, or `"meta"` for the architecture alone.

---

**`moe_training`** `MoeTrainingSettings | None` — default: None

Router load-balancing terms for training (both zero at inference).

---

**Returns:** `PraxisPolicy`

The model; `pretrained_path` names the checkpoint directory.

```python
nemo_automodel.components.models.praxis.model.PraxisPolicy.moe_layers() -> list[nemo_automodel.components.moe.layers.MoE]
```

The MoE feed-forwards of both towers in module order.

```python
nemo_automodel.components.models.praxis.model.PraxisPolicy.place_rotary_tables(
    device: torch.device
) -> None
```

Move the rotary frequency tables to `device`.

The tables are plain attributes, not buffers, so module casts and sharding leave them alone;
they are built once on the CPU and only moved, never recomputed on the accelerator.

```python
nemo_automodel.components.models.praxis.model.PraxisPolicy.update_moe_gate_bias() -> None
```

Move every router's selection bias toward balanced expert load (call once per optimizer step).

```python
nemo_automodel.components.models.praxis.model._per_layer_moe_training(
    config: nemo_automodel.components.models.praxis.config.PraxisConfig,
    moe_training: nemo_automodel.components.models.praxis.backbone.MoeTrainingSettings | None
) -> nemo_automodel.components.models.praxis.backbone.MoeTrainingSettings | None
```

Divide the balance-loss weight by the number of MoE layers.

`MoeTrainingSettings.aux_loss_coeff` is the weight of the mean over the MoE layers, as in the
reference trainer; the shared Automodel router adds one term per layer, so each layer gets the
weight divided by the layer count.

```python
nemo_automodel.components.models.praxis.model.max_tokens_per_window(
    config: nemo_automodel.components.models.praxis.config.PraxisConfig
) -> int
```

The longest token sequence one training window produces in either tower.

The decoder runs on the noised chunks' latents one token per latent pixel, `patch_size[1] * patch_size[2]`
times the encoder's video tokens, so with its registers it is the longest sequence of a window and the
per-rank capacity a token dispatcher needs for one window per micro-step (`local_batch_size` 1; larger
local batches multiply it).

**Parameters:**

**`config`** `PraxisConfig`

The checkpoint configuration: window geometry, latent stride and register count.

---

**Returns:** `int`

Number of tokens.

```python
nemo_automodel.components.models.praxis.model.praxis_backend(
    dispatcher: typing.Literal['torch', 'deepep', 'hybridep', 'uccl_ep'] = 'torch'
) -> nemo_automodel.components.models.common.utils.BackendConfig
```

The backend the Praxis checkpoint runs with: torch linears, flex attention, grouped-GEMM experts.

**Parameters:**

**`dispatcher`** `Literal['torch', 'deepep', 'hybridep', 'uccl_ep']` — default: 'torch'

Token exchange of the routed experts under expert parallelism (`distributed.ep_size > 1`):
the DTensor all-gather of `"torch"`, or DeepEP, HybridEP or UCCL-EP. Without expert parallelism
the experts run locally whatever the value.

---

```python
nemo_automodel.components.models.praxis.model.ModelClass = PraxisPolicy
```