> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.

# 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`](#nemo_automodel-recipes-wam-train-Encoders)                             | The frozen encoders every window goes through, and the encoding of the blank caption. |
| [`TrainWAMFinetuneRecipe`](#nemo_automodel-recipes-wam-train-TrainWAMFinetuneRecipe) | Fine-tune Praxis-1 on training windows with FSDP2, all weights or LoRA adapters.      |

### Functions

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

### Data

[`logger`](#nemo_automodel-recipes-wam-train-logger)

### API

```python
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`

---

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

**Bases:** [BaseRecipe](/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/recipes/base_recipe#nemo_automodel-recipes-base_recipe-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`**

---

```python
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.

```python
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.

```python
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.

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

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

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

```python
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`.

---

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