LAM

This project contains the latent action model (LAM) training path as a standalone package. VLA model code, VLA training, deployment servers, benchmark examples, and unrelated scripts are intentionally absent.

Layout

• latent_action_model/cli.py: Lightning CLI entrypoint.
• latent_action_model/models/: LAM model, vision encoders, decoders, and quantizers.
• latent_action_model/training/: Lightning module and training callbacks.
• latent_action_model/data/: LAM datamodule, dataset wrapper, collate, video augmentation, and LeRobot readers.
• configs/: LAM training configs.
• scripts/: single-node and multi-node launch scripts.

Install

conda create -n lam-train python=3.10 -y
conda activate lam-train
pip install -r requirements.txt
pip install -e .

Training also needs local datasets and vision encoder weights matching the selected YAML config. The default config points to local paths under /mnt/project_rlinf/jlchen.

Single-Node Training

bash scripts/train.sh \
  --config configs/dino_base_ae.yaml

Equivalent direct CLI:

python -m latent_action_model.cli fit \
  --config configs/dino_base_ae.yaml

Multi-Node Training

Set the distributed environment variables and launch:

NNODES=2 NODE_RANK=0 MASTER_ADDR=<master-ip> MASTER_PORT=29500 \
  bash scripts/train_ddp.sh \
  --config configs/dino_base_ae.yaml

train_ddp.sh passes --trainer.num_nodes ${NNODES} to Lightning.

Smoke Checks

python -m compileall latent_action_model
python -m latent_action_model.cli fit \
  --config configs/dino_base_ae.yaml \
  --print_config

For a minimal real training check, override the trainer to run one GPU step on an available local mixture:

LAM_ENABLE_MANUAL_WANDB=0 python -m latent_action_model.cli fit \
  --config configs/dino_base_ae.yaml \
  --data.data_mix=libero \
  --data.batch_size=1 \
  --data.num_workers=0 \
  --trainer.max_steps=1 \
  --trainer.devices=1 \
  --trainer.strategy=auto \
  --trainer.limit_val_batches=0 \
  --trainer.num_sanity_val_steps=0