Training and evaluation

scripts/world_model/train_wm.py dispatches to the standard WorldModel or DinoWM trainer.

Standard WorldModel: CNN or DINOv2, MLP or RSSM

Example data loading options:

A local CNN run with generated NPZ data:

python scripts/world_model/make_toy_dataset.py --out_dir data/toy --episodes 200
python scripts/world_model/train_wm.py --method wm \
  --data_path data/toy --dataset_type npz \
  --backbone cnn --dynamics_type mlp \
  --epochs 20 --out_dir runs/cnn

A local CNN run with generated parquet data (frames.parquet and windows.parquet):

python scripts/world_model/train_wm.py --method wm \
  --data_path /path/to/parquet_dataset --dataset_type hf \
  --backbone cnn --dynamics_type rssm \
  --epochs 20 --out_dir runs/wm

The Parquet tables may contain image paths and array-valued sensor columns.

The same trainer supports CNN or DINOv2 with RSSM. Set --backbone cnn on that command for CNN + RSSM. The visual encoder can also be vjepa2 in the standard architecture. --dynamics_type mlp selects the non-RSSM dynamics path.

For a Hugging Face dataset and a frozen DINOv2 encoder:

python scripts/world_model/train_wm.py --method wm \
  --data_path lquan9/collect1 --dataset_type hf \
  --backbone dinov2 --backbone_name dinov2_vitb14 \
  --freeze_backbone --dynamics_type rssm \
  --epochs 20 --batch_size 64 --out_dir runs/dino_rssm

Hub rows may embed images and numbered sensor columns. hf accepts a Hub ID.

---

For the existing two-GPU DDP setup, the underlying trainer accepts --ddp:

torchrun --standalone --nproc_per_node=2 \
  scripts/world_model/train_world_model.py --ddp \
  --data_path lquan9/collect1 --backbone dinov2 \
  --backbone_name dinov2_vitb14 --freeze_backbone \
  --epochs 20 --batch_size 64 --lr 3e-4 \
  --horizon 8 --img_size 224 --wandb --out_dir runs/dinov2

---

JEPA-lite and state decoder

JEPA-lite is the wm trainer with a latent matching objective inspired by joint-embedding predictive architectures [ADM+23].

python scripts/world_model/train_wm.py --method wm \
  --wm_mode jepa-lite --data_path lquan9/collect1 \
  --backbone cnn --dynamics_type mlp \
  --epochs 20 --batch_size 64 --lr 3e-4 \
  --horizon 8 --img_size 224 --out_dir runs/jepa_lite_c1_512

The same objective accepts the DINOv2 backbone using --backbone dinov2 --backbone_name dinov2_vitb14 --freeze_backbone. The state-head trainer supports --latent_source observed (decode encoded sensor frames) and --latent_source rollout (decode predicted states used by MPC). A two-stage run initializes from observed frames, then fine-tunes on predicted rollouts:

python scripts/world_model/train_state_head_decoder.py \
  --checkpoint runs/jepa_lite_c1_512/latest.pt \
  --dataset /path/to/matching/dataset --dataset_type hf \
  --heads state --latent_source observed --epochs 5 --batch_size 8

python scripts/world_model/train_state_head_decoder.py \
  --checkpoint runs/jepa_lite_c1_512/latest_with_observed_heads.pt \
  --dataset /path/to/matching/dataset --dataset_type hf \
  --heads state --latent_source rollout --epochs 10 --batch_size 8

Each stage reports both observed and rollout validation position RMSE; select a planner checkpoint by held-out rollout error and planning performance. You can run rollout directly from latest.pt if you do not need the observed-latent probe. The dataset must match the checkpoint’s odometry, IMU, action dimensions and action convention. The trainer loads norm_stats.pt beside the input checkpoint when present, or accepts --norm_stats_path; it does not create new statistics. For a state-only head, set the ROS planner’s risk_weight to zero. Use --heads both only with valid collision and non-collision labels.

For long rollouts, train the dynamics at the deployment horizon before fitting the rollout head. A staged curriculum can resume the same run at increasing --horizon values. --epochs is the cumulative final epoch, not the number of additional epochs, so a four-stage run can use horizons 8, 16, 24, 36 with epoch limits 10, 20, 30, 45. Checkpoints include and restore the frozen EMA target model; legacy checkpoints without it initialize the target from the online model and print a warning. Do not fit a head beyond the checkpoint’s trained_horizon; the head trainer rejects that mismatch.

Standard training writes action_bounds.pt from the transformed training samples by default. Disable this only when bounds are supplied separately with --no-save_action_bounds. The model checkpoint records a training_contract containing the trained rollout horizon, history length, and dataset action semantics. Head training verifies this action contract against its dataset and preserves it in the output checkpoint.

---

DinoWM

DinoWM [ZPLP24] uses a frozen DINOv2 encoder and a ViT latent predictor. The trainer takes CLI flags, including a required data path:

WANDB_MODE=disabled python scripts/world_model/train_wm.py \
  --method dinowm --data_path lquan9/collect1 \
  --epochs 10 --out_dir runs/dinowm

For a local frames/windows Parquet dataset, use --data_path /path/to/dataset with the same command. A separate val split is used when available.

The options --frameskip, --num_hist, --num_pred, --encoder, and --encoder_name control DinoWM’s sequence and visual encoder. For image reconstruction, add --has_decoder --train_decoder and choose --decoder_type vqvae or transposed_conv. Image metrics work with the core dependencies; install .[decoder] to request optional LPIPS scores from world_model.metrics.eval_images(..., include_lpips=True). configs/dino-wm/trial.yaml records an earlier configuration; the current trainer uses argparse flags.

Checkpoints and evaluation

Use a separate --out_dir for each run. Checkpoint paths below are relative to that directory:

Trainer	Latest checkpoint	Numbered checkpoints	Best validation checkpoint
Standard WorldModel (including JEPA-lite)	latest.pt after every epoch	epoch_NNN.pt every --save_every epochs (default 10), and at the final epoch	best.pt
DinoWM	checkpoints/model_latest.pth at each save interval	checkpoints/model_N.pth every --save_every_x_epoch epochs (default 1)	checkpoints/model_best.pth

Both trainers select the best checkpoint by validation loss at their checkpoint intervals, when a separate validation split exists. A train-only dataset has no best validation checkpoint.

To evaluate the standard WorldModel trained on the toy dataset above, run the toy-environment MPC example:

python scripts/world_model/eval_mpc.py --ckpt runs/cnn/latest.pt --episodes 20

Option	Meaning	Default
--data_path	Hugging Face dataset ID or local path	required
--dataset_type	hf for Hub rows or local Parquet tables; npz for NPZ episodes. DinoWM currently supports hf.	hf
--schema_config	Optional YAML mapping for local Parquet table paths, columns, sensor policy, and split fallback.	canonical names
--precompute_windows	Cache local Parquet windows in RAM before standard training.	off
--max_windows	Limit cached Parquet windows per split; requires --precompute_windows.	unlimited
--backbone	Standard visual encoder: cnn, dinov2, or vjepa2	cnn
--dynamics_type	Standard dynamics: mlp, rssm, or td-jepa	mlp
--wm_mode	standard or jepa-lite objective	standard
--freeze_backbone	Keep pretrained visual weights fixed	off
--ddp	Distributed standard training	off
--epochs	Training epochs	20 for standard training
--out_dir	Checkpoint and run directory	runs
--save_every	Standard numbered checkpoint interval	10
--save_action_bounds	Write transformed per-dimension action statistics for CEM.	on
--resume	Restore model, optimizer, epoch, and EMA target state.	automatic out_dir/latest.pt
--save_every_x_epoch	DinoWM checkpoint interval	1

Read next

• ROS2 contains ROS2 integration details (e.g., launch arguments, checkpoint requirements, topics and node parameters).