.venv/bin/torchrun --master_port 29502 --nproc_per_node=2 --nnodes=1 vla_foundry/main.py \
  --model "include vla_foundry/config_presets/models/diffusion_policy.yaml" \
  --model.vision_language_backbone.freeze_text_encoder True \
  --model.transformer.is_causal True \
  --data "include vla_foundry/config_presets/data/lbm/lbm_data_params.yaml" \
  --data.dataset_manifest ["s3://tri-ml-datasets-uw2/vla_foundry_datasets/v0.4.1/BimanualPutRedBellPepperInBin/shards/manifest.jsonl"] \
  --data.dataset_statistics ["s3://tri-ml-datasets-uw2/vla_foundry_datasets/v0.4.1/BimanualPutRedBellPepperInBin/shards/stats.json"] \
  --data.dataset_modality ["robotics"] \
  --data.augmentation.enabled True \
  --data.augmentation.image "include vla_foundry/config_presets/data/lbm/lbm_image_augmentation_params.yaml" \
  --data.image_size 224 \
  --data.action_dim 20 \
  --data.processor openai/clip-vit-base-patch32 \
  --data.img_num_tokens 64 \
  --data.allow_multiple_epochs True \
  --data.num_workers 4 \
  --data.seq_len 2048 \
  --distributed.fsdp True \
  --num_checkpoints 3 \
  --hparams.loss_function mse \
  --hparams.per_gpu_batch_size 128 \
  --hparams.global_batch_size 384 \
  --hparams.grad_clip_norm 1.0 \
  --hparams.lr 5e-4 \
  --hparams.lr_cooldown_end 1e-5 \
  --remote_sync s3://tri-ml-datasets-uw2/vla_foundry/model_checkpoints/diffusion_policy \
  --total_train_samples 100000 \
  --wandb True \
  "$@"