> 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.objective

Flow-matching training objective of the Praxis policy.

A training window holds `K` clean chunk latents, the anchor latent, the text states, the
normalized action tokens and the boundary poses. One step noises the window and asks the
policy for the video and action velocities:

* Video: the whole window first gets one timestep `t0` and one noise `eps0`; the target
  is `eps0 - x0`. Then every chunk `k` is re-noised with its own timestep `t_k` from the
  bf16-rounded `x_t`, so that the K noised chunks share one velocity target. The clean copies
  of chunks `0 .. K-2` are the recovered clean latents. Timesteps follow the shifted law
  `t = s u / (1 + (s - 1) u)`.
* Actions: every chunk gets an independent uniform timestep; the target is `eps - x0`.
* Text dropout swaps the text for the encoded blank caption (the unconditional branch of
  text guidance); proprio dropout zeroes the boundary poses.
* Training also moves the rotary positions of the anchor, the text and the video by random
  offsets inside one frame block so the model does not memorize absolute positions.

The loss is `video_weight * mse(video) + action_weight * mse(action)`; a window whose action
loss is not finite or above `action_guard` contributes an action loss of exactly zero.

## Module Contents

### Classes

| Name                                                                                        | Description                                                     |
| ------------------------------------------------------------------------------------------- | --------------------------------------------------------------- |
| [`Draws`](#nemo_automodel-components-models-praxis-objective-Draws)                         | The random values of one window, drawn before the forward pass. |
| [`EncodedWindow`](#nemo_automodel-components-models-praxis-objective-EncodedWindow)         | One encoded training window (batch size 1).                     |
| [`LossOutput`](#nemo_automodel-components-models-praxis-objective-LossOutput)               | Loss and detached metrics of one window.                        |
| [`ObjectiveSettings`](#nemo_automodel-components-models-praxis-objective-ObjectiveSettings) | Weights, timestep law and dropout rates of the objective.       |
| [`PraxisObjective`](#nemo_automodel-components-models-praxis-objective-PraxisObjective)     | The training objective around a `PraxisPolicy`.                 |

### Functions

| Name                                                                                        | Description                                      |
| ------------------------------------------------------------------------------------------- | ------------------------------------------------ |
| [`shifted_timesteps`](#nemo_automodel-components-models-praxis-objective-shifted_timesteps) | `shift * u / (1 + (shift - 1) * u)` elementwise. |

### API

```python
class nemo_automodel.components.models.praxis.objective.Draws(
    text_u: torch.Tensor,
    video_u0: torch.Tensor,
    video_noise: torch.Tensor,
    video_u: torch.Tensor,
    rope_shifts: tuple[int, int, int, int],
    action_u: torch.Tensor,
    action_noise: torch.Tensor,
    proprio_u: torch.Tensor
)
```

Dataclass

The random values of one window, drawn before the forward pass.

**`action_noise`** `Tensor`

Action noise of shape \[1, K \* F, action_patch_dim] (float32).

---

**`action_u`** `Tensor`

Uniform values of shape \[K] for the action timesteps.

---

**`proprio_u`** `Tensor`

Uniform value of shape \[1] for the proprio dropout.

---

**`rope_shifts`** `tuple[int, int, int, int]`

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

---

**`text_u`** `Tensor`

Uniform value of shape \[1] for the text dropout.

---

**`video_noise`** `Tensor`

Window noise of shape \[1, channels, K \* F, h, w] (bf16).

---

**`video_u`** `Tensor`

Uniform values of shape \[K] for the per-chunk timesteps.

---

**`video_u0`** `Tensor`

Uniform value of shape \[1] for the window timestep.

---

```python
class nemo_automodel.components.models.praxis.objective.EncodedWindow(
    text: torch.Tensor,
    anchor: torch.Tensor,
    chunks: list[torch.Tensor],
    actions: torch.Tensor,
    proprio: torch.Tensor
)
```

Dataclass

One encoded training window (batch size 1).

**`actions`** `Tensor`

Normalized action tokens of shape \[1, K \* F, action_patch_dim] (float32), chunk-major.

---

**`anchor`** `Tensor`

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

---

**`chunks`** `list[Tensor]`

`K` clean chunk latents, each of shape \[1, channels, F, h, w].

---

**`proprio`** `Tensor`

Normalized boundary poses of shape \[1, K, proprio_patch_dim] (float32).

---

**`text`** `Tensor`

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

---

```python
class nemo_automodel.components.models.praxis.objective.LossOutput(
    loss: torch.Tensor,
    loss_video: torch.Tensor,
    loss_action: torch.Tensor,
    action_dropped: torch.Tensor
)
```

Dataclass

Loss and detached metrics of one window.

**`action_dropped`** `Tensor`

1.0 when the guard zeroed the action loss, else 0.0.

---

**`loss`** `Tensor`

Scalar float32 loss with gradients.

---

**`loss_action`** `Tensor`

Detached action loss before the guard.

---

**`loss_video`** `Tensor`

Detached video loss.

---

```python
class nemo_automodel.components.models.praxis.objective.ObjectiveSettings(
    video_weight: float = 1.0,
    action_weight: float = 1.0,
    video_timestep_shift: float = 5.0,
    text_dropout: float = 0.1,
    proprio_dropout: float = 0.1,
    action_guard: float = 5.0,
    train_backbone: bool = True
)
```

Dataclass

Weights, timestep law and dropout rates of the objective.

**`action_guard`** `float = 5.0`

---

**`action_weight`** `float = 1.0`

---

**`proprio_dropout`** `float = 0.1`

---

**`text_dropout`** `float = 0.1`

---

**`train_backbone`** `bool = True`

---

**`video_timestep_shift`** `float = 5.0`

---

**`video_weight`** `float = 1.0`

---

```python
class nemo_automodel.components.models.praxis.objective.PraxisObjective(
    policy: nemo_automodel.components.models.praxis.model.PraxisPolicy,
    settings: nemo_automodel.components.models.praxis.objective.ObjectiveSettings,
    text_uncond: torch.Tensor
)
```

**Bases:** `Module`

The training objective around a `PraxisPolicy`.

```python
nemo_automodel.components.models.praxis.objective.PraxisObjective.action_inputs(
    batch: nemo_automodel.components.models.praxis.objective.EncodedWindow,
    draws: nemo_automodel.components.models.praxis.objective.Draws
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
```

Noise the actions on the CPU in float32.

**Returns:** `tuple[Tensor, Tensor, Tensor, Tensor]`

`(noised [1, K*F, A], timesteps [1, K], target [1, K*F, A], proprio [1, K, P])`.

```python
nemo_automodel.components.models.praxis.objective.PraxisObjective.draw(
    batch: nemo_automodel.components.models.praxis.objective.EncodedWindow,
    generator: torch.Generator
) -> nemo_automodel.components.models.praxis.objective.Draws
```

Draw the random values of `batch` from a CPU `generator` in a fixed order.

```python
nemo_automodel.components.models.praxis.objective.PraxisObjective.forward(
    batch: nemo_automodel.components.models.praxis.objective.EncodedWindow,
    generator: torch.Generator
) -> nemo_automodel.components.models.praxis.objective.LossOutput
```

Draw the random values and return the loss of `batch`.

```python
nemo_automodel.components.models.praxis.objective.PraxisObjective.loss(
    batch: nemo_automodel.components.models.praxis.objective.EncodedWindow,
    draws: nemo_automodel.components.models.praxis.objective.Draws
) -> nemo_automodel.components.models.praxis.objective.LossOutput
```

Loss of one window for the given random values (no dropout, no rotary offset in eval mode).

```python
nemo_automodel.components.models.praxis.objective.PraxisObjective.video_inputs(
    batch: nemo_automodel.components.models.praxis.objective.EncodedWindow,
    draws: nemo_automodel.components.models.praxis.objective.Draws,
    device: torch.device,
    dtype: torch.dtype
) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor]
```

Noise the window.

**Returns:** `Tensor`

`(noised [1, C, K*F, h, w], clean [1, C, (K-1)*F, h, w] or None, timesteps [1, K], target [1, C, K*F, h, w])`

```python
nemo_automodel.components.models.praxis.objective.shifted_timesteps(
    u: torch.Tensor,
    shift: float
) -> torch.Tensor
```

`shift * u / (1 + (shift - 1) * u)` elementwise.