nemo_automodel.components.models.praxis.config

Typed view of the Praxis checkpoint config.json.

The checkpoint config is a nested JSON document. The sections the model consumes are parsed into frozen dataclasses with explicit fields; unknown keys raise so that a checkpoint of another model variant is refused instead of silently misread.

Module Contents

Classes

Name	Description
`ActionExpertConfig`	`action_expert`: the small transformer that denoises the actions.
`ActionSpaceConfig`	`action_space`: the per-frame action row layout.
`AttentionConfig`	`backbone.attention`.
`BackboneConfig`	`backbone`: the two-tower video transformer.
`ChunkingConfig`	`chunking`: the window geometry in raw frames and latents.
`ExpertRopeConfig`	`action_expert.rope`: the expert reuses the trunk rotary positions.
`FlexKernelOptions`	`action_expert.flex_kernel_options`: pinned flex-attention template parameters.
`MoeSectionConfig`	`backbone.moe`: a DeepSeek-style sigmoid router with a shared expert.
`NormalizationConfig`	`normalization`: the q01/q99 action statistics and the proprio mean/std.
`PraxisConfig`	The parsed `config.json` of a Praxis single-arm checkpoint.
`PraxisConfigError`	Raised when `config.json` is not a supported Praxis single-arm config.
`ProjectionConfig`	One entry of `input_projections` / `output_projections`.
`RopeConfig`	`backbone.rope`: integer-position rotary embedding of the trunk.
`TextEncoderConfig`	`text_encoder`: the frozen Qwen text encoder and the unconditional caption.
`TowerConfig`	`backbone.encoder` or `backbone.decoder`.
`VideoConfig`	`video`: the canvas of the camera views and the VAE geometry.

Functions

Name	Description
`_dataclass_from`	Build `cls` from `section` with its field names; unknown keys are an error.
`_section`	-

Data

API

class nemo_automodel.components.models.praxis.config.ActionExpertConfig(
    dim: int,
    num_layers: int,
    num_heads: int,
    head_dim: int,
    trunk_dim: int,
    action_patch_dim: int,
    proprio_patch_dim: int,
    num_future_per_chunk: int,
    max_chunks: int,
    max_action_tokens: int,
    rope: nemo_automodel.components.models.praxis.config.ExpertRopeConfig,
    global_order: tuple[str, ...] = ('register', 'anchor', 'text'),
    tap_norm: bool = True,
    clean_only_conditioning: bool = False,
    rope_cross: bool = False,
    rope_self: bool = False,
    rope_base: float = 100.0,
    absolute_positions: bool = True,
    action_pos_table: bool = False,
    future_chunk_offset: typing.Literal['zero'] = 'zero',
    attention: typing.Literal['sdpa', 'flex_dynamic'] = 'sdpa',
    flex_kernel_options: nemo_automodel.components.models.praxis.config.FlexKernelOptions | None = None,
    head_plus_zero: bool = False,
    layout_stamp: int | None = None
)

Dataclass

action_expert: the small transformer that denoises the actions.

absolute_positions

bool = True

action_patch_dim

int

action_pos_table

bool = False

attention

Literal['sdpa', 'flex_dynamic'] = 'sdpa'

clean_only_conditioning

bool = False

dim

int

flex_kernel_options

FlexKernelOptions | None = None

future_chunk_offset

Literal['zero'] = 'zero'

global_order

tuple[str, ...] = ('register', 'anchor', 'text')

head_dim

int

head_plus_zero

bool = False

layout_stamp

int | None = None

max_action_tokens

int

max_chunks

int

num_future_per_chunk

int

num_heads

int

num_layers

int

proprio_patch_dim

int

rope

ExpertRopeConfig

rope_base

float = 100.0

rope_cross

bool = False

rope_self

bool = False

tap_norm

bool = True

trunk_dim

int

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

classmethod

class nemo_automodel.components.models.praxis.config.ActionSpaceConfig(
    dim: int,
    layout: str,
    pose_layout: str
)

Dataclass

action_space: the per-frame action row layout.

dim

int

layout

str

pose_layout

str

class nemo_automodel.components.models.praxis.config.AttentionConfig(
    num_heads: int,
    num_kv_heads: int,
    head_dim: int,
    gated: bool = True,
    qk_norm_eps: float = 1e-06,
    qk_norm_compute_dtype: str = 'float32',
    rope_after_norm: bool = False,
    flex_block_size: int = 128,
    backend: typing.Literal['triton', 'flash'] = 'triton'
)

Dataclass

backbone.attention.

backend

Literal['triton', 'flash'] = 'triton'

flex_block_size

int = 128

gated

bool = True

head_dim

int

num_heads

int

num_kv_heads

int

qk_norm_compute_dtype

str = 'float32'

qk_norm_eps

float = 1e-06

rope_after_norm

bool = False

class nemo_automodel.components.models.praxis.config.BackboneConfig(
    dim: int,
    timestep_dim: int,
    num_registers: int,
    attention: nemo_automodel.components.models.praxis.config.AttentionConfig,
    rope: nemo_automodel.components.models.praxis.config.RopeConfig,
    moe: nemo_automodel.components.models.praxis.config.MoeSectionConfig,
    dense_ffn_hidden: int,
    encoder: nemo_automodel.components.models.praxis.config.TowerConfig,
    decoder: nemo_automodel.components.models.praxis.config.TowerConfig,
    modulation_dtype: typing.Literal['bfloat16', 'float32', 'mixed'] = 'mixed',
    norm_eps: float = 1e-06,
    norm_compute_dtype: str = 'float32',
    layout: typing.Literal['chunk_causal'] = 'chunk_causal'
)

Dataclass

backbone: the two-tower video transformer.

attention

AttentionConfig

decoder

TowerConfig

dense_ffn_hidden

int

dim

int

encoder

TowerConfig

layout

Literal['chunk_causal'] = 'chunk_causal'

modulation_dtype

Literal['bfloat16', 'float32', 'mixed'] = 'mixed'

moe

MoeSectionConfig

norm_compute_dtype

str = 'float32'

norm_eps

float = 1e-06

num_registers

int

rope

RopeConfig

timestep_dim

int

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

classmethod

class nemo_automodel.components.models.praxis.config.ChunkingConfig(
    n_chunks: int,
    action_chunk: int,
    action_frames_per_token: int,
    latents_per_chunk: int,
    past_frames: int,
    future_frames: int,
    window_frames: int,
    chunk_geometry: tuple[tuple[int, int], ...],
    grid_hz: float,
    global_frame_anchor: bool = True,
    contiguous_chunks: bool = True,
    context_format: str = 'vxt',
    proprio_frames: int = 1
)

Dataclass

chunking: the window geometry in raw frames and latents.

action_chunk

int

action_frames_per_token

int

chunk_geometry

tuple[tuple[int, int], ...]

context_format

str = 'vxt'

contiguous_chunks

bool = True

future_frames

int

global_frame_anchor

bool = True

grid_hz

float

latents_per_chunk

int

n_chunks

int

past_frames

int

proprio_frames

int = 1

window_frames

int

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

classmethod

class nemo_automodel.components.models.praxis.config.ExpertRopeConfig(
    kind: typing.Literal['trunk'],
    dim: int,
    effective_dim: int,
    base: int,
    table_size: int,
    order: typing.Literal['before_norm']
)

Dataclass

action_expert.rope: the expert reuses the trunk rotary positions.

base

int

dim

int

effective_dim

int

kind

Literal['trunk']

order

Literal['before_norm']

table_size

int

class nemo_automodel.components.models.praxis.config.FlexKernelOptions(
    fwd_BLOCK_M: int,
    fwd_BLOCK_N: int,
    num_warps: int,
    num_stages: int
)

Dataclass

action_expert.flex_kernel_options: pinned flex-attention template parameters.

fwd_BLOCK_M

int

fwd_BLOCK_N

int

num_stages

int

num_warps

int

class nemo_automodel.components.models.praxis.config.MoeSectionConfig(
    num_routed_experts: int,
    top_k_routed: int,
    expert_hidden: int,
    num_experts: int | None = None,
    shared_expert: bool = True,
    top_k: int | None = None,
    router_activation: typing.Literal['sigmoid'] = 'sigmoid',
    router_dtype: str = 'float32',
    routed_scaling_factor: float = 2.5,
    no_expert_slots: bool = False,
    routed_weight_layout: typing.Literal['expert_first'] = 'expert_first',
    dispatch_chunk_size: int | None = None
)

Dataclass

backbone.moe: a DeepSeek-style sigmoid router with a shared expert.

dispatch_chunk_size

int | None = None

expert_hidden

int

no_expert_slots

bool = False

num_experts

int | None = None

num_routed_experts

int

routed_scaling_factor

float = 2.5

routed_weight_layout

Literal['expert_first'] = 'expert_first'

router_activation

Literal['sigmoid'] = 'sigmoid'

router_dtype

str = 'float32'

shared_expert

bool = True

top_k

int | None = None

top_k_routed

int

nemo_automodel.components.models.praxis.config.MoeSectionConfig.__post_init__() -> None

class nemo_automodel.components.models.praxis.config.NormalizationConfig(
    domain: str,
    action_q01: tuple[float, ...],
    action_q99: tuple[float, ...],
    q99_normed_dims: tuple[int, ...],
    proprio_mean: tuple[float, ...],
    proprio_std: tuple[float, ...],
    action_formula: str,
    proprio_formula: str
)

Dataclass

normalization: the q01/q99 action statistics and the proprio mean/std.

action_formula

str

action_q01

tuple[float, ...]

action_q99

tuple[float, ...]

domain

str

proprio_formula

str

proprio_mean

tuple[float, ...]

proprio_std

tuple[float, ...]

q99_normed_dims

tuple[int, ...]

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.config.PraxisConfigError()

Bases: ValueError

Raised when config.json is not a supported Praxis single-arm config.

class nemo_automodel.components.models.praxis.config.ProjectionConfig(
    c_in: int | None = None,
    c_out: int | None = None,
    move_dim: bool = False,
    timestep_modulated: bool = False
)

Dataclass

One entry of input_projections / output_projections.

c_in

int | None = None

c_out

int | None = None

move_dim

bool = False

timestep_modulated

bool = False

class nemo_automodel.components.models.praxis.config.RopeConfig(
    base: int = 26000000,
    dim: int = 128,
    effective_dim: int = 64,
    max_seq_len: int = 90000000,
    efficient_cache_max: int = 1000000,
    max_intra: int = 360,
    max_inter: int = 129600,
    chunked_video_layout: bool = True,
    n_chunks: int | None = None,
    time_scale: float = 1.0,
    n_tiles: int | None = None,
    ref_placement: typing.Literal['group'] = 'group'
)

Dataclass

backbone.rope: integer-position rotary embedding of the trunk.

base

int = 26000000

chunked_video_layout

bool = True

dim

int = 128

effective_dim

int = 64

efficient_cache_max

int = 1000000

max_inter

int = 129600

max_intra

int = 360

max_seq_len

int = 90000000

n_chunks

int | None = None

n_tiles

int | None = None

ref_placement

Literal['group'] = 'group'

time_scale

float = 1.0

class nemo_automodel.components.models.praxis.config.TextEncoderConfig(
    model: str,
    hidden_layer: int,
    pad_to_multiple_of: int,
    trim_padding: bool,
    hidden_dim: int,
    uncond: typing.Literal['blank_caption'] = 'blank_caption',
    blank_caption: str = ' '
)

Dataclass

text_encoder: the frozen Qwen text encoder and the unconditional caption.

blank_caption

str = ' '

hidden_dim

int

hidden_layer

int

model

str

pad_to_multiple_of

int

trim_padding

bool

uncond

Literal['blank_caption'] = 'blank_caption'

class nemo_automodel.components.models.praxis.config.TowerConfig(
    depth: int,
    dense_blocks: tuple[int, ...],
    text_in: nemo_automodel.components.models.praxis.config.ProjectionConfig | None,
    video_in: nemo_automodel.components.models.praxis.config.ProjectionConfig,
    video_out: nemo_automodel.components.models.praxis.config.ProjectionConfig,
    tap_blocks: tuple[int, ...] = (),
    modulation_dim: int | None = None,
    unpatch_factor: tuple[int, int, int] = (1, 1, 1),
    positive_zero_residual_add: bool = True
)

Dataclass

backbone.encoder or backbone.decoder.

dense_blocks

tuple[int, ...]

depth

int

modulation_dim

int | None = None

positive_zero_residual_add

bool = True

tap_blocks

tuple[int, ...] = ()

text_in

ProjectionConfig | None

unpatch_factor

tuple[int, int, int] = (1, 1, 1)

video_in

ProjectionConfig

video_out

ProjectionConfig

nemo_automodel.components.models.praxis.config.TowerConfig.from_json(
    section: collections.abc.Mapping[str, nemo_automodel.components.models.praxis.config.JsonValue],
    where: str
) -> nemo_automodel.components.models.praxis.config.TowerConfig

classmethod

class nemo_automodel.components.models.praxis.config.VideoConfig(
    height: int,
    width: int,
    views: tuple[str, ...],
    layout: str,
    tile_bounds_rows_cols: dict[str, tuple[int, int, int, int]],
    vae: str,
    vae_encode_chunk_frames: int,
    latent_channels: int,
    vae_stride: tuple[int, int, int],
    patch_size: tuple[int, int, int],
    tile_height: int = 256,
    tile_width: int = 416,
    canvas: str = 'tiles',
    keep_aspect: bool = True,
    antialias: bool = False,
    aspect_tol: float = 0.03,
    vae_encode_impl: str = 'megakernel',
    vae_megakernel_chunk_frames: int | None = None
)

Dataclass

video: the canvas of the camera views and the VAE geometry.

antialias

bool = False

aspect_tol

float = 0.03

canvas

str = 'tiles'

height

int

keep_aspect

bool = True

latent_channels

int

layout

str

patch_size

tuple[int, int, int]

tile_bounds_rows_cols

dict[str, tuple[int, int, int, int]]

tile_height

int = 256

tile_width

int = 416

vae

str

vae_encode_chunk_frames

int

vae_encode_impl

str = 'megakernel'

vae_megakernel_chunk_frames

int | None = None

vae_stride

tuple[int, int, int]

views

tuple[str, ...]

width

int

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

classmethod

nemo_automodel.components.models.praxis.config._dataclass_from(
    cls: type[nemo_automodel.components.models.praxis.config._T],
    section: collections.abc.Mapping[str, nemo_automodel.components.models.praxis.config.JsonValue],
    where: str
) -> nemo_automodel.components.models.praxis.config._T

Build cls from section with its field names; unknown keys are an error.

nemo_automodel.components.models.praxis.config._section(
    parent: collections.abc.Mapping[str, nemo_automodel.components.models.praxis.config.JsonValue],
    name: str
) -> collections.abc.Mapping[str, nemo_automodel.components.models.praxis.config.JsonValue]

nemo_automodel.components.models.praxis.config.FORMAT_VERSION = 2

nemo_automodel.components.models.praxis.config.JsonScalar = str | int | float | bool | None

nemo_automodel.components.models.praxis.config.JsonValue = JsonScalar | list['JsonValue'] | dict[str, 'JsonValue']

nemo_automodel.components.models.praxis.config.MODEL_TYPE = 'praxis'

nemo_automodel.components.models.praxis.config._T = TypeVar('_T')