Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsPraxisnemo_automodel.components.models.praxis.model

nemo_automodel.components.models.praxis.model

View as MarkdownOpen in Claude

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

NameDescription
CheckpointConfigExporterWrite the source checkpoint’s config.json next to consolidated weights.
PraxisOutputOutputs of one denoising step.
PraxisPolicyVideo backbone and action expert of a Praxis checkpoint.

Functions

NameDescription
_per_layer_moe_trainingDivide the balance-loss weight by the number of MoE layers.
max_tokens_per_windowThe longest token sequence one training window produces in either tower.
praxis_backendThe backend the Praxis checkpoint runs with: torch linears, flex attention, grouped-GEMM experts.

Data

ModelClass

API

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.

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.

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.

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.

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, Module, 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
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.

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]Defaults to (0, 0, 0, 0)

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

decode
boolDefaults to True

Run the decoder tower for the video velocity.

backbone_grad
boolDefaults to 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.

classmethod

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

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 | NoneDefaults to None

Backend selection; defaults to praxis_backend.

device
str | torch.deviceDefaults to 'cpu'

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

moe_training
MoeTrainingSettings | NoneDefaults to None

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

Returns: PraxisPolicy

The model; pretrained_path names the checkpoint directory.

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.

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.

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

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.

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.

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']Defaults to '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.

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