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`	The random values of one window, drawn before the forward pass.
`EncodedWindow`	One encoded training window (batch size 1).
`LossOutput`	Loss and detached metrics of one window.
`ObjectiveSettings`	Weights, timestep law and dropout rates of the objective.
`PraxisObjective`	The training objective around a `PraxisPolicy`.

Functions

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

API

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.

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

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.

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

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.

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]).

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.

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.

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

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])

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

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