run_trainer.py

#!/usr/bin/env python3

# Copyright (c) 2025-2026 Robotics and AI Institute LLC dba RAI Institute. All rights reserved.


"""
Pythonic version of the train.sh script launcher for memory-visuomotor-policies.

This script provides a clean, object-oriented interface for launching training experiments
with proper parameter validation, environment setup, and command generation.
"""

import os
import sys
import random
import subprocess
import argparse
from pathlib import Path
from typing import Optional, List, Dict, Any
from dataclasses import dataclass, field
from enum import Enum
from termcolor import colored

def required_env_var(env_var_name: str) -> str:
    """Validate that an environment variable is set."""
    if env_var_name not in os.environ:
        raise ValueError(f"Error: {env_var_name} is required and cannot be empty")
    return os.environ[env_var_name]

class LaunchLocation(Enum):
    """Supported launch locations."""
    LOCAL = "local"

class VisionEncoder(Enum):
    """Supported vision encoders."""
    CLIP = "clip"
    DINOV3 = "dinov3"
    CROSSMAE = "crossmae"

@dataclass
class TrainingConfig:
    """Configuration for training experiments."""
    # Required parameters
    downsample_obs: int
    batch_size: int
    num_gpus: int
    model_config: str
    data_config: str
    seed: int = 1
    nocompress: bool = False  # If True, disable compression (default: False, compression enabled)
    gating_flag: str = "False"
    seq_length: int = 4
    block_finetune: Optional[str] = None
    vision_encoder: VisionEncoder = VisionEncoder.CROSSMAE
    launch_location: LaunchLocation = LaunchLocation.LOCAL
    dry_run: bool = False  # Add dry_run parameter
    basic_run: bool = False  # Add basic_run parameter for single-GPU training
    load_in_mem: bool = False  # Add load_in_mem parameter for loading dataset in memory
    compile_model: bool = False  # Add compile_model parameter for compiling the model
    lr: float = 5e-4  # Add lr parameter for learning rate
    attn_drop: float = 0.0  # Add attn_drop parameter for attention dropout probability
    # Mamba
    use_mamba: bool = False
    use_lstm: bool = False
    use_transformer_xl: bool = False
    use_linear_attention: bool = False
    xl_chunk_factor: int = 32

    # k_ptp parameter for multi-step prediction
    k_ptp: int = 0  # Add k_ptp parameter

    # Wandb configuration
    wandb_project_name: Optional[str] = None  # Will be set based on experiment type if None

    # Computed fields
    extra_flags: List[str] = field(default_factory=list)
    exp_base_dir: str = "memory_exps"
    exp_name: str = ""
    task_config: str = ""
    gbs: int = 0
    accum_iter: int = 0
    num_repeat_traj: int = 0
    port_num: int = 0

    # resume from checkpoint
    resume: Optional[str] = None

    full_attn_inds: Optional[list[int]] = None
    pool_block_inds: Optional[list[int]] = None
    # block attention indices
    block_attn_ind: Optional[list[int]] = None
    break_after_n_epochs: Optional[int] = None

    compressor_latent_len: int = 1

    def __post_init__(self):
        """Validate and compute derived fields after initialization."""
        self._validate_required_params()
        self._validate_environment()
        # Only install latest code if not in dry run mode
        self._compute_derived_fields()
        self._generate_experiment_name()
        self._generate_extra_flags()
        self._set_wandb_project_name()  # Add this new method call

    def _validate_required_params(self):
        """Validate that all required parameters are provided."""
        required_params = {
            'downsample_obs': self.downsample_obs,
            'batch_size': self.batch_size,
            'num_gpus': self.num_gpus,
            'model_config': self.model_config,
            'data_config': self.data_config,
        }

        for param_name, param_value in required_params.items():
            if param_value is None or param_value == "":
                raise ValueError(f"Error: {param_name} is required and cannot be empty")

        if self.k_ptp > 0:
            assert self.gating_flag == "block_nogate" or self.gating_flag == "nogate", "k_ptp is only supported with block_nogate or block_sigmoid_g5 gating flag"

    def _validate_environment(self):
        """Validate that all environment variables are set."""
        required_env_vars = [
            'PRISM_DATAROOT',
            'EXP_STORAGE_BASE_DIR'
        ]
        for env_var_name in required_env_vars:
            required_env_var(env_var_name)
        return

    def _compute_derived_fields(self):
        """Compute derived fields based on input parameters."""
        libero_full_datasets = ["task_libero_100.json", "task_libero_90.json", "task_libero_10.json", "task_libero_spatial.json", "task_libero_object.json", "task_libero_goal.json"]
        # Global batch size calculation
        if self.data_config in libero_full_datasets:
            token_batch_size = 8192 # with sl = 256, batch size of 32
            self.gbs = token_batch_size // self.seq_length
        else:
            token_batch_size = 65536
            self.gbs = max(256, token_batch_size // self.seq_length)
        print(f"token_batch_size: {token_batch_size}, gbs: {self.gbs}, num_gpus: {self.num_gpus}, batch_size: {self.batch_size}")
        self.gbs = ((self.gbs + self.num_gpus * self.batch_size - 1) // (self.num_gpus * self.batch_size)) * (self.num_gpus * self.batch_size)

        # Accumulation iterations
        self.accum_iter = max(1, self.gbs // (self.batch_size * self.num_gpus))

        # Number of repeat trajectories
        self.num_repeat_traj = max(256, int(256 * (256 / self.seq_length)))

        # Adjust for specific datasets
        if self.data_config in ["task_robocasa_atomic.json", "task_robocasa_atomic_all.json", "task_robocasa_mem_mix.json"]:
            self.num_repeat_traj = max(1, self.num_repeat_traj // 2)
        if self.data_config in libero_full_datasets:
            self.num_repeat_traj = max(1, self.num_repeat_traj // 16)
            if self.data_config in ["task_libero_100.json", "task_libero_90.json"]:
                self.num_repeat_traj = max(self.num_repeat_traj // 2, 1)
        if self.data_config in ["task_rw_mutex_mem_washandreturn.json"]:
            self.num_repeat_traj = self.num_repeat_traj * 3

        # Task config path
        self.task_config = f"config/{self.data_config}"

        # Random port number
        self.port_num = 2452 + random.randint(0, 99)

    def _generate_extra_flags(self):
        """Generate extra command line flags based on configuration."""
        self.extra_flags = []

        # Block attention indices
        if self.block_attn_ind is not None:
            print(colored(f"Using block attention in layers {self.block_attn_ind}", "green"))
            self.extra_flags.extend([
                "--model-cfg.policy-cfg.block_attn_ind", " ".join(map(str, self.block_attn_ind))
            ])

        if self.full_attn_inds is not None:
            print(colored(f"Using full attention in layers {self.full_attn_inds}", "green"))
            self.extra_flags.extend([
                "--model-cfg.policy-cfg.full_attn_inds", " ".join(map(str, self.full_attn_inds))
            ])

        if self.pool_block_inds is not None:
            print(colored(f"Using pool block attention in layers {self.pool_block_inds}", "green"))
            self.extra_flags.extend([
                "--model-cfg.policy-cfg.pool_block_inds", " ".join(map(str, self.pool_block_inds))
            ])

        # Gating flags
        gating_flag = self.gating_flag
        if 'block' in gating_flag:
            self.extra_flags.extend([
                "--model-cfg.policy-cfg.use_block_attention"
            ])
            gating_flag = gating_flag.replace("block_", "")
        if gating_flag == "True" or gating_flag == "False":
            pass  # No extra flags for boolean gating
        elif gating_flag == "nogate" or gating_flag == "none":
            pass
        elif gating_flag and gating_flag != "":
            # Extract gate name and add block attention with gating
            gate_name = gating_flag.replace("gate", "")
            self.extra_flags.extend([
                "--model-cfg.policy-cfg.gate_full_attn_layers",
                "--model-cfg.policy-cfg.gating_type", gate_name
            ])
        else:
            raise ValueError(f"Invalid gating flag: {gating_flag}")


        # Vision encoder flags
        if self.vision_encoder == VisionEncoder.CLIP:
            self.extra_flags.extend([
                "--model-cfg.vision-encoder-cfg.vision-encoder", "vit_base_patch32_clip_224.openai"
            ])
            if self.exp_base_dir is None:
                self.exp_base_dir = "clip_exps"
        elif self.vision_encoder == VisionEncoder.DINOV3:
            self.extra_flags.extend([
                "--model-cfg.vision-encoder-cfg.vision-encoder", "facebook/dinov3-vits16plus-pretrain-lvd1689m"
            ])
        elif self.vision_encoder == VisionEncoder.CROSSMAE:
            # get the PRISM_DATAROOT and if the file is not there, install it using: https://huggingface.co/mlfu7/ICRT/resolve/main/crossmae_rtx/cross-mae-rtx-vitb.pth
            import os
            PRISM_DATAROOT = os.environ.get("PRISM_DATAROOT")
            if PRISM_DATAROOT is None:
                raise ValueError("PRISM_DATAROOT is not set")
            VISION_ENCODER_PATH = os.path.join(PRISM_DATAROOT, "crossmae_rtx", "cross-mae-rtx-vitb.pth")
            if not os.path.exists(VISION_ENCODER_PATH):
                os.makedirs(os.path.dirname(VISION_ENCODER_PATH), exist_ok=True)
                subprocess.run(["wget", "https://huggingface.co/mlfu7/ICRT/resolve/main/crossmae_rtx/cross-mae-rtx-vitb.pth", "-O", VISION_ENCODER_PATH])
            self.extra_flags.extend([
                "--model-cfg.vision-encoder-cfg.vision-encoder", VISION_ENCODER_PATH
            ])

        # Dataset-specific flags
        if "robocasa" in self.data_config:
            self.extra_flags.extend([
                "--shared-cfg.has_base_action"
            ])

        if "libero" in self.data_config:
            if self.exp_base_dir is None:
                self.exp_base_dir = "libero_exps"

        # nopool is always assumed by default
        if self.exp_base_dir is None:
            self.exp_base_dir = "nopool_exps"
        # Add compression flag if compression is enabled (nocompress is False)
        if not self.nocompress:
            self.extra_flags.append("--model-cfg.policy-cfg.compress_full_attn")

        self.extra_flags.extend([
            "--model-cfg.policy-cfg.compressor_latent_len", str(self.compressor_latent_len)
        ])

        if self.load_in_mem:
            self.extra_flags.append("--dataset-cfg.load-in-mem")

        if self.compile_model:
            self.extra_flags.append("--trainer-cfg.compile_model")

        if self.use_lstm:
            if self.exp_base_dir is None:
                self.exp_base_dir = "lstm_exps"
            self.extra_flags.append("--model-cfg.policy-cfg.use_lstm")
        if self.use_mamba:
            if self.exp_base_dir is None:
                self.exp_base_dir = "mamba_exps"
            self.extra_flags.append("--model-cfg.policy-cfg.use_mamba")
        if self.use_transformer_xl:
            if self.exp_base_dir is None:
                self.exp_base_dir = "transformer_xl"
            self.extra_flags.append("--model-cfg.policy-cfg.use_transformer_xl")
            self.extra_flags.extend([
                "--model-cfg.policy-cfg.xl_chunk_factor", str(self.xl_chunk_factor)
            ])
        if self.use_linear_attention:
            if self.exp_base_dir is None:
                self.exp_base_dir = "linear_attention_exps"
            self.extra_flags.append("--model-cfg.policy-cfg.use_linear_attention")

        # Debug flags (if not in debug mode)
        if os.environ.get("DEBUG", "").lower() not in ["true", "1"]:
            print(colored("Adding logging flags", "green"))
            self.extra_flags.extend([
                "--logging-cfg.log-name", self.exp_name,
            ])

        if self.break_after_n_epochs is not None:
            self.extra_flags.extend([
                "--trainer-cfg.break_after_n_epochs", str(self.break_after_n_epochs)
            ])

        if self.attn_drop > 0:
            self.extra_flags.extend([
                "--model-cfg.policy-cfg.attn_drop", str(self.attn_drop)
            ])

    def _generate_experiment_name(self):
        """Generate the experiment name based on configuration."""
        # Regular experiment naming
        task_config_name = Path(self.task_config).stem

        # Base experiment name
        self.exp_name = (
            f"exp_ds{self.downsample_obs}_{Path(self.model_config).stem}_"
            f"{Path(self.data_config).stem}_GBS{self.gbs}_cfg{task_config_name}_"
            f"sl{self.seq_length}_s{self.seed}_{self.vision_encoder.value}"
        )

        # Add gating suffix
        if self.gating_flag == "False":
            self.exp_name += "_pool"
        elif self.gating_flag and self.gating_flag != "":
            self.exp_name += f"_{self.gating_flag}"

        # Add k_ptp suffix if > 0
        if self.k_ptp > 0:
            self.exp_name += f"_kptp{self.k_ptp}"

        # Add nopool suffix (always present by default)
        self.exp_name += "_nopool"
        # Add compress suffix if compression is enabled
        if not self.nocompress:
            self.exp_name += "_compress"

        if self.block_attn_ind is not None:
            self.exp_name += f"_block_attn_ind{'_'.join(map(str, self.block_attn_ind))}"

        if self.full_attn_inds is not None:
            self.exp_name += f"_full_attn_inds{len(self.full_attn_inds)}"

        if self.pool_block_inds is not None:
            self.exp_name += f"_pool_block_inds{len(self.pool_block_inds)}"

        if self.use_lstm:
            self.exp_name += "_lstm"
        if self.use_mamba:
            self.exp_name += "_mamba"
        if self.use_transformer_xl:
            self.exp_name += "_txl"

        if self.use_linear_attention:
            self.exp_name += "_linear_attn"

        if self.use_transformer_xl:
            self.exp_name += f"_txl_cf{self.xl_chunk_factor}"

        if self.attn_drop > 0:
            self.exp_name += f"_ad{self.attn_drop}"

        if self.compressor_latent_len != 1:
            self.exp_name += f"_clen{self.compressor_latent_len}"

        print(colored(self.exp_name, "green"))

    def _set_wandb_project_name(self):
        """Set wandb project name based on experiment type if not provided."""
        if self.wandb_project_name is None:
            self.wandb_project_name = "prism_litev2"

    def get_training_command(self) -> List[str]:
        """Generate the complete training command with all arguments."""
        exp_storage_base_dir = os.environ.get("EXP_STORAGE_BASE_DIR", "/tmp/experiments")

        # Set num_workers based on basic_run flag
        num_workers = 0 if self.basic_run else 72

        # Set validation and save frequencies
        val_every = 10
        save_every = 10

        # Core arguments - each argument and value as separate list elements
        args = [
            "--dataset-cfg.dataset-json", self.task_config,
            "--shared-cfg.num_cameras", "2",
            "--dataset-cfg.num_repeat_traj", str(self.num_repeat_traj),
            "--model-cfg.policy-cfg.scratch-llama-config", f"config/model_config/{self.model_config}",
            "--shared-cfg.seq_length", str(self.seq_length),
            "--shared-cfg.seed", str(self.seed),
            "--shared-cfg.batch-size", str(self.batch_size),
            "--trainer-cfg.accum-iter", str(self.accum_iter),
            "--shared-cfg.num-pred-steps", "32",
            "--optimizer-cfg.warmup-epochs", "2",
            "--trainer-cfg.num-workers", str(num_workers),
            "--trainer-cfg.epochs", "200",
            "--trainer-cfg.val-every", str(val_every),
            "--shared-cfg.save-every", str(save_every),
            "--shared-cfg.use-language-conditioning",
            "--shared-cfg.pad-to-max-length",
            "--optimizer-cfg.lr", str(self.lr),
            "--model-cfg.policy-cfg.model_version", "v2",
            "--shared-cfg.downsample_obs", str(self.downsample_obs),
            "--logging-cfg.output-dir", f"{exp_storage_base_dir}/{self.exp_base_dir}/{self.exp_name}",
            "--trainer-cfg.wandb-project", self.wandb_project_name,  # Add wandb project name
        ]
        if self.resume is not None:
            args.extend(["--shared-cfg.resume", self.resume])

        # Add k_ptp argument if > 0
        if self.k_ptp > 0:
            args.extend([
                "--shared-cfg.k_ptp", str(self.k_ptp)
            ])

        # Add extra flags
        args.extend(self.extra_flags)

        return args

    def get_torchrun_command(self) -> List[str]:
        """Generate the torchrun command for training."""
        # Use python instead of torchrun for basic_run
        if self.basic_run:
            cmd = ["python", "scripts/train.py"]
        else:
            # Base command
            cmd = [
                f"CC=gcc CXX=g++ torchrun",
                f"--nproc_per_node={self.num_gpus}",
                f"--master_port={self.port_num}",
                "scripts/train.py"
            ]

        # Add training arguments
        cmd.extend(self.get_training_command())

        return cmd

    def print_config(self):
        """Print the current configuration."""
        print("=" * 50)
        print("TRAINING CONFIGURATION")
        print("=" * 50)
        print(f"DOWNSAMPLE_OBS: {self.downsample_obs}")
        print(f"BATCH_SIZE: {self.batch_size}")
        print(f"NUM_GPUS: {self.num_gpus}")
        print(f"MODEL_CONFIG: {self.model_config}")
        print(f"DATA_CONFIG: {self.data_config}")
        print(f"SEED: {self.seed}")
        print(f"NOCOMPRESS: {self.nocompress}")
        print(f"GATING_FLAG: {self.gating_flag}")
        print(f"SEQ_LENGTH: {self.seq_length}")
        print(f"VISION_ENCODER: {self.vision_encoder.value}")
        print(f"K_PTP: {self.k_ptp}")  # Add k_ptp to config print
        print(f"WANDB_PROJECT_NAME: {self.wandb_project_name}")  # Add wandb project name to config print
        print(f"ACCUM_ITER: {self.accum_iter}")
        print(f"NUM_REPEAT_TRAJ: {self.num_repeat_traj}")
        print(f"TASK_CONFIG: {self.task_config}")
        print(f"GBS: {self.gbs}")
        print(f"PORT_NUM: {self.port_num}")
        print(f"exp_name: {self.exp_name}")
        print(f"exp_base_dir: {self.exp_base_dir}")
        print(f"EXTRA_FLAGS: {' '.join(self.extra_flags)}")
        print("=" * 50)

    def run_training(self, dry_run: bool = False):
        """Run the training command locally."""
        self.print_config()
        self._run_local(dry_run)

    def _run_local(self, dry_run: bool = False):
        """Run training locally."""
        # Set CUDA_VISIBLE_DEVICES
        cuda_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "0,1,2,3,4,5,6,7")
        os.environ["CUDA_VISIBLE_DEVICES"] = cuda_devices

        # Get the command
        cmd = self.get_torchrun_command()

        print(f"\nExecuting command locally:")
        print(" ".join(cmd))
        print()

        if dry_run:
            print("DRY RUN - Command not executed")
            return

        # Execute the command with proper environment inheritance
        try:
            # Use shell=True to ensure proper environment inheritance
            subprocess.run(" ".join(cmd), shell=True, check=True)
        except subprocess.CalledProcessError as e:
            print(f"Training failed with exit code {e.returncode}")
            sys.exit(1)
        except KeyboardInterrupt:
            print("Training interrupted by user")
            sys.exit(1)



def main():
    """Main entry point for the script."""
    parser = argparse.ArgumentParser(
        description="Pythonic training script launcher for memory-visuomotor-policies",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
    Examples:
    # Basic training (local)
    DEBUG=1 python run_trainer.py -ds 4 -bs 8 -ng 2  -mc libero_1x.json -dc task_robocasa_mem_fruitsink.json -s 1 --gating-flag block_sigmoid_g5 -sl 8 --dry-run

    # With custom parameters (local)
    python run_trainer.py -ds 4 -bs 8 -ng 2  -mc libero_1x.json -dc task_robocasa_mem_four.json -s 1 --gating-flag block_sigmoid_g5 -sl 8 --dry-run

    # Dry run to see the command
    python run_trainer.py -ds 4 -bs 8 -ng 2  -mc libero_1x.json -dc task_robocasa_mem_fruitsink.json -s 1 --gating-flag block_sigmoid_g5 -sl 8 --dry-run

    # With k_ptp parameter
    python run_trainer.py -ds 4 -bs 8 -ng 2 -mc libero_1x.json -dc task_robocasa_mem_fruitsink.json -s 1 --gating-flag block_sigmoid_g5 -sl 8 --k-ptp 16 --dry-run
    """
    )

    # Required positional arguments
    parser.add_argument("-ds", "--downsample_obs", type=int, help="Downsample observations and actions by this factor (e.g., 8 means use every 8th frame). Also used to compute max sequence length for transformer.", default=8)
    parser.add_argument("-bs", "--batch_size", type=int, help="Batch size per GPU. Global batch size (GBS) is computed automatically based on this, num_gpus, and seq_length.")
    parser.add_argument("-ng", "--num_gpus", type=int, help="Number of GPUs for distributed training. Used to compute global batch size and accumulation steps.")
    parser.add_argument("-mc", "--model_config", type=str, help="Model configuration JSON file (e.g., libero_1x.json). Path relative to config/model_config/.", default='libero_1x.json')
    parser.add_argument("-dc", "--data_config", type=str, help="Dataset configuration JSON file (e.g., task_robocasa_mem_four.json). Path relative to config/.", default='task_robocasa_mem_four.json')
    parser.add_argument("-br", "--break-after-n-epochs", type=int, help="Stop training after N epochs (useful for early stopping).", default=101)
    parser.add_argument("-ad", "--attn-drop", type=float, default=0.0, help="Attention dropout probability for transformer layers (0.0 = no dropout).")

    # Optional arguments
    parser.add_argument("-s", "--seed", type=int, default=1, help="Random seed for reproducibility (default: 1)")
    parser.add_argument("--nocompress", action="store_true", help="Disable compression mechanism. By default, compression (hierarchical local + global attention) is enabled to reduce memory usage.")
    parser.add_argument("-gf", "--gating-flag", type=str, default="block_sigmoid_g5", help="Gating mechanism for attention layers. Options: 'block_sigmoid_g5', 'nogate', 'block_nogate', etc. Controls how information is filtered in transformer blocks.")
    parser.add_argument("-sl", "--seq-length", type=int, default=8, help="Sequence length in number of (state, action) pairs. Used to compute max sequence length for transformer: max_seq_len = (latent_len * seq_length) // downsample_obs")
    parser.add_argument("-ve", "--vision-encoder", type=VisionEncoder, default=VisionEncoder.CROSSMAE,
                       choices=list(VisionEncoder), help="Vision encoder backbone: CLIP, DINOv3, or CrossMAE (default: CrossMAE)")
    parser.add_argument("--use-lstm", action="store_true", help="Use LSTM instead of Transformer as the sequence model backbone")
    parser.add_argument("--use-mamba", action="store_true", help="Use Mamba instead of Transformer as the sequence model backbone")
    parser.add_argument("--use-transformer-xl", action="store_true", help="Use Transformer-XL instead of standard Transformer (enables memory mechanism for longer sequences)")
    parser.add_argument("--xl-chunk-factor", type=int, default=64, help="Chunking factor for Transformer-XL. Sequences are processed in chunks of this size for memory efficiency.")
    parser.add_argument("--block-attn-ind", nargs="+", type=int, default=None, help="Layer indices (0-based) to use block attention instead of full attention. Reduces memory for long sequences.")
    parser.add_argument("--full-attn-inds", nargs="+", type=int, default=None, help="Layer indices (0-based) to use full attention even when block attention is enabled elsewhere")
    parser.add_argument("--pool-block-inds", nargs="+", type=int, default=None, help="Layer indices (0-based) to use pool block attention (mean pooling instead of attention)")
    parser.add_argument("--use-linear-attention", action="store_true", help="Use linear attention instead of standard quadratic attention (reduces memory complexity)")
    parser.add_argument("-clen", "--compressor-latent-len", type=int, default=1, help="Number of tokens to allocate per step during compression. Higher values preserve more information but use more memory.")

    parser.add_argument("--load-in-mem", action="store_true", help="Load entire dataset into memory for faster data loading. Requires sufficient RAM.")
    parser.add_argument("--compile-model", action="store_true", help="Compile model with torch.compile() for faster training (PyTorch 2.0+). May increase compilation time initially.")
    parser.add_argument("--lr", type=float, default=5e-4, help="Learning rate for optimizer (default: 5e-4)")
    parser.add_argument("--resume", type=str, default=None, help="Path to checkpoint file to resume training from. If None, training starts from scratch.")

    # k_ptp parameter
    parser.add_argument("--k-ptp", type=int, default=0, help="Number of past tokens to predict per step (for multi-step prediction). Only supported with 'block_nogate' or 'nogate' gating flags. (default: 0, disabled)")

    # Wandb configuration
    parser.add_argument("--wandb-project-name", type=str, help="Wandb project name for experiment tracking (default: 'prism')", default="prism")

    # add the option to specify the exp_base_dir
    parser.add_argument("--exp-base-dir", type=str, default=None, help="Base directory name for organizing experiments (e.g., 'memory_exps', 'nopool_exps'). Auto-set based on model type if None.")

    # dry run parameter
    parser.add_argument("--basic-run", action="store_true", help="Use single-GPU training mode: runs with 'python' instead of 'torchrun', sets num_workers=0. Useful for debugging.")
    parser.add_argument("--dry-run", action="store_true", help="Print the training command without executing it. Useful for verifying configuration.")

    args = parser.parse_args()

    # Create training configuration
    config = TrainingConfig(
        downsample_obs=args.downsample_obs,
        batch_size=args.batch_size,
        num_gpus=args.num_gpus,
        break_after_n_epochs=args.break_after_n_epochs,
        model_config=args.model_config,
        data_config=args.data_config,
        seed=args.seed,
        nocompress=args.nocompress,
        gating_flag=args.gating_flag,
        attn_drop=args.attn_drop,
        seq_length=args.seq_length,
        vision_encoder=args.vision_encoder,
        dry_run=args.dry_run,  # Pass dry_run parameter
        basic_run=args.basic_run,  # Pass basic_run parameter
        k_ptp=args.k_ptp,  # Pass k_ptp parameter
        wandb_project_name=args.wandb_project_name,  # Pass wandb project name
        load_in_mem=args.load_in_mem,  # Pass load_in_mem parameter
        compile_model=args.compile_model,  # Pass compile_model parameter
        lr=args.lr,  # Pass lr parameter
        resume=args.resume,  # Pass resume parameter
        use_mamba=args.use_mamba,  # Pass use_mamba parameter
        use_transformer_xl=args.use_transformer_xl,  # Pass use_mamba parameter
        use_linear_attention=args.use_linear_attention,  # Pass use_linear_attention parameter
        xl_chunk_factor=args.xl_chunk_factor,
        use_lstm=args.use_lstm,  # Pass use_lstm parameter
        exp_base_dir=args.exp_base_dir,  # Pass exp_base_dir parameter
        block_attn_ind=args.block_attn_ind,  # Pass block_attn_ind parameter
        full_attn_inds=args.full_attn_inds,  # Pass full_attn_inds parameter
        pool_block_inds=args.pool_block_inds,  # Pass pool_block_inds parameter
        compressor_latent_len=args.compressor_latent_len,  # Pass compressor_latent_len parameter
    )

    # Run training
    config.run_training(dry_run=args.dry_run)

if __name__ == "__main__":
    main()