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

# nemo_automodel.components.models.praxis

## Submodules

* **[`nemo_automodel.components.models.praxis.action_expert`](/nemo-automodel/nemo_automodel/components/models/praxis/action_expert)**
* **[`nemo_automodel.components.models.praxis.backbone`](/nemo-automodel/nemo_automodel/components/models/praxis/backbone)**
* **[`nemo_automodel.components.models.praxis.config`](/nemo-automodel/nemo_automodel/components/models/praxis/config)**
* **[`nemo_automodel.components.models.praxis.hub`](/nemo-automodel/nemo_automodel/components/models/praxis/hub)**
* **[`nemo_automodel.components.models.praxis.layers`](/nemo-automodel/nemo_automodel/components/models/praxis/layers)**
* **[`nemo_automodel.components.models.praxis.lora_export`](/nemo-automodel/nemo_automodel/components/models/praxis/lora_export)**
* **[`nemo_automodel.components.models.praxis.masks`](/nemo-automodel/nemo_automodel/components/models/praxis/masks)**
* **[`nemo_automodel.components.models.praxis.model`](/nemo-automodel/nemo_automodel/components/models/praxis/model)**
* **[`nemo_automodel.components.models.praxis.objective`](/nemo-automodel/nemo_automodel/components/models/praxis/objective)**
* **[`nemo_automodel.components.models.praxis.parallelize`](/nemo-automodel/nemo_automodel/components/models/praxis/parallelize)**
* **[`nemo_automodel.components.models.praxis.state_dict_adapter`](/nemo-automodel/nemo_automodel/components/models/praxis/state_dict_adapter)**

## Package Contents

### Classes

| Name                                                                           | Description                                                 |
| ------------------------------------------------------------------------------ | ----------------------------------------------------------- |
| [`PraxisConfig`](#nemo_automodel-components-models-praxis-config-PraxisConfig) | The parsed `config.json` of a Praxis single-arm checkpoint. |
| [`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                                                                                       |
| --------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------- |
| [`praxis_backend`](#nemo_automodel-components-models-praxis-model-praxis_backend) | The backend the Praxis checkpoint runs with: torch linears, flex attention, grouped-GEMM experts. |

### API

```python
class nemo_automodel.components.models.praxis.config.PraxisConfig(
    backbone: nemo_automodel.components.models.praxis.config.BackboneConfig,
    action_expert: nemo_automodel.components.models.praxis.config.ActionExpertConfig,
    chunking: nemo_automodel.components.models.praxis.config.ChunkingConfig,
    video: nemo_automodel.components.models.praxis.config.VideoConfig,
    text_encoder: nemo_automodel.components.models.praxis.config.TextEncoderConfig,
    action_space: nemo_automodel.components.models.praxis.config.ActionSpaceConfig,
    normalization: nemo_automodel.components.models.praxis.config.NormalizationConfig,
    format_version: int = FORMAT_VERSION,
    model_type: str = MODEL_TYPE,
    embodiment: typing.Literal['single_arm'] = 'single_arm',
    torch_dtype: str = 'bfloat16',
    fp32_tensors: tuple[str, ...] = ('*.ffn.expert_bias',),
    pose_frame: typing.Literal['flange'] = 'flange'
)
```

Dataclass

The parsed `config.json` of a Praxis single-arm checkpoint.

**`action_expert`** `ActionExpertConfig`

---

**`action_space`** `ActionSpaceConfig`

---

**`backbone`** `BackboneConfig`

---

**`chunking`** `ChunkingConfig`

---

**`embodiment`** `Literal['single_arm'] = 'single_arm'`

---

**`format_version`** `int = FORMAT_VERSION`

---

**`fp32_tensors`** `tuple[str, ...] = ('*.ffn.expert_bias',)`

---

**`model_type`** `str = MODEL_TYPE`

---

**`normalization`** `NormalizationConfig`

---

**`pose_frame`** `Literal['flange'] = 'flange'`

---

**`text_encoder`** `TextEncoderConfig`

---

**`torch_dtype`** `str = 'bfloat16'`

---

**`video`** `VideoConfig`

---

```python
nemo_automodel.components.models.praxis.config.PraxisConfig.from_json(
    config: collections.abc.Mapping[str, nemo_automodel.components.models.praxis.config.JsonValue]
) -> nemo_automodel.components.models.praxis.config.PraxisConfig
```

classmethod

Parse and validate a checkpoint config document.

```python
nemo_automodel.components.models.praxis.config.PraxisConfig.from_pretrained(
    path: str | pathlib.Path
) -> nemo_automodel.components.models.praxis.config.PraxisConfig
```

classmethod

Read and parse `<path>/config.json`.

```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/nemo_automodel/components/models/common/hf_checkpointing_mixin#nemo_automodel-components-models-common-hf_checkpointing_mixin-HFCheckpointingMixin), `Module`, [MoEFSDPSyncMixin](/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.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.

---

## Docs

- [nemo_automodel.components.models.praxis.action_expert](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/praxis/action_expert.md)
- [nemo_automodel.components.models.praxis.backbone](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/praxis/backbone.md)
- [nemo_automodel.components.models.praxis.config](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/praxis/config.md)
- [nemo_automodel.components.models.praxis.hub](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/praxis/hub.md)
- [nemo_automodel.components.models.praxis.layers](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/praxis/layers.md)
- [nemo_automodel.components.models.praxis.lora_export](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/praxis/lora_export.md)
- [nemo_automodel.components.models.praxis.masks](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/praxis/masks.md)
- [nemo_automodel.components.models.praxis.model](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/praxis/model.md)
- [nemo_automodel.components.models.praxis.objective](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/praxis/objective.md)
- [nemo_automodel.components.models.praxis.parallelize](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/praxis/parallelize.md)
- [nemo_automodel.components.models.praxis.state_dict_adapter](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/praxis/state_dict_adapter.md)