nemo_automodel.components.models.praxis

Submodules

Package Contents

Classes

Name	Description
`PraxisConfig`	The parsed `config.json` of a Praxis single-arm checkpoint.
`PraxisOutput`	Outputs of one denoising step.
`PraxisPolicy`	Video backbone and action expert of a Praxis checkpoint.

Functions

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

API

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

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.

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.

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

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.

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.

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

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.