nemo_automodel.recipes.wam.train

Fine-tune a Praxis world-action policy on recorded episodes or DROID windows.

The model is PraxisPolicy under FSDP2, trained on all weights or through LoRA adapters (the peft section). The frozen Wan2.2 VAE, the Qwen text encoder and the DROID window reader come from the reference praxis package, which must be importable; each rank encodes its own windows on the fly. Checkpoints are written in the Praxis layout that the reference simulation and server load.

Run::

automodel examples/wam/praxis/praxis_1_libero.yaml —nproc-per-node 8

Module Contents

Classes

Name	Description
`Encoders`	The frozen encoders every window goes through, and the encoding of the blank caption.
`TrainWAMFinetuneRecipe`	Fine-tune Praxis-1 on training windows with FSDP2, all weights or LoRA adapters.

Functions

Name	Description
`draw_seed`	Seed of the random draws (timesteps, noise, dropout) of one window: per (step, micro-step, rank).
`lora_trains_backbone`	Whether any trainable parameter sits in the video backbone, so that its autograd graph must be kept.

Data

logger

API

class nemo_automodel.recipes.wam.train.Encoders(
    vae: praxis.vae.Autoencoder,
    text_encoder: praxis.text_encoder.TextEncoder,
    text_uncond: torch.Tensor
)

Dataclass

The frozen encoders every window goes through, and the encoding of the blank caption.

text_encoder

TextEncoder

text_uncond

Tensor

vae

Autoencoder

class nemo_automodel.recipes.wam.train.TrainWAMFinetuneRecipe(
    cfg
)

Bases: BaseRecipe

Fine-tune Praxis-1 on training windows with FSDP2, all weights or LoRA adapters.

YAML sections: model (PraxisPolicy.from_pretrained), optional peft, encoders (vae and text_encoder of the reference package), objective (ObjectiveSettings), dataset and dataloader, optimizer, lr_scheduler, step_scheduler, clip_grad_norm, distributed (FSDP2 with mp_policy) and checkpoint.

cfg

nemo_automodel.recipes.wam.train.TrainWAMFinetuneRecipe._encode(
    window: praxis.train.data.TrainingWindow
) -> nemo_automodel.components.models.praxis.objective.EncodedWindow

Encode one window with the frozen VAE and text encoder on this rank.

nemo_automodel.recipes.wam.train.TrainWAMFinetuneRecipe._optimizer_step(
    batches: list[praxis.train.data.TrainingWindow],
    step: int
) -> dict[str, float]

Accumulate the gradients of batches, clip, step the optimizer and the schedulers.

nemo_automodel.recipes.wam.train.TrainWAMFinetuneRecipe.run_train_validation_loop()

Train for the scheduled steps, logging every step and checkpointing on the scheduler’s steps.

nemo_automodel.recipes.wam.train.TrainWAMFinetuneRecipe.setup()

Build the sharded model, the encoders, the data, the optimizer and the checkpointer.

nemo_automodel.recipes.wam.train.draw_seed(
    seed: int,
    step: int,
    micro: int,
    rank: int,
    world: int
) -> int

Seed of the random draws (timesteps, noise, dropout) of one window: per (step, micro-step, rank).

nemo_automodel.recipes.wam.train.lora_trains_backbone(
    policy: nemo_automodel.components.models.praxis.model.PraxisPolicy
) -> bool

Whether any trainable parameter sits in the video backbone, so that its autograd graph must be kept.

Linear and grouped-expert adapters are the only trainable parameters after apply_lora_to_linear_modules; the base weights are frozen.

Parameters:

policy

PraxisPolicy

The model after apply_lora_to_linear_modules.

nemo_automodel.recipes.wam.train.logger = logging.getLogger(__name__)