opentau.configs.train

Training pipeline configuration module.

This module provides the TrainPipelineConfig class which contains all configuration parameters needed to run a training pipeline, including dataset settings, policy configuration, training hyperparameters, and evaluation settings.

Functions

warn(*args, **kwargs)

Print a warning message to stderr.

Classes

TrainPipelineConfig(dataset_mixture, policy, ...)

Configuration for the training pipeline.

class opentau.configs.train.TrainPipelineConfig(dataset_mixture: ~opentau.configs.default.DatasetMixtureConfig, policy: ~opentau.configs.policies.PreTrainedConfig | None = None, output_dir: ~pathlib.Path | None = None, job_name: str | None = None, resume: bool = False, seed: int | None = 1000, resolution: tuple[int, int] = (224, 224), num_cams: int = 2, max_state_dim: int = 32, max_action_dim: int = 32, action_chunk: int = 50, loss_weighting: dict[str, float] = <factory>, num_workers: int = 4, batch_size: int | None = None, gradient_accumulation_steps: int = 1, dataloader_batch_size: int | None = None, prefetch_factor: int | None = None, steps: int = 100000, log_freq: int = 200, save_checkpoint: bool = True, save_freq: int = 20000, use_policy_training_preset: bool = False, optimizer: ~opentau.optim.optimizers.OptimizerConfig | None = None, scheduler: ~opentau.optim.schedulers.LRSchedulerConfig | None = None, wandb: ~opentau.configs.default.WandBConfig = <factory>, debug: bool = False, trace_nans: bool = False, env: ~opentau.envs.configs.EnvConfig | None = None, eval: ~opentau.configs.default.EvalConfig | None = <factory>, eval_freq: int = 0, val_freq: int = 0, last_checkpoint_only: bool = True, running_best_count: int = 0, running_best_metric: str = 'auto', server: ~opentau.configs.deployment.ServerConfig = <factory>, planner: ~opentau.configs.deployment.PlannerConfig = <factory>)

Bases: HubMixin

Configuration for the training pipeline.

This class contains all configuration parameters needed to run a training pipeline, including dataset settings, policy configuration, training hyperparameters, and evaluation settings.

Parameters:

• dataset_mixture – Configuration for the dataset mixture to use during training.

• policy – Configuration for the policy model. If None, must be set via CLI or from a pretrained checkpoint.

• output_dir – Directory where all run outputs will be saved. If another training session uses the same directory, its contents will be overwritten unless resume is set to True.

• job_name – Name identifier for the training job. If not provided, defaults to the policy type.

• resume – If True, resume a previous run. Requires output_dir to point to an existing run directory with at least one checkpoint. When resuming, the configuration from the checkpoint is used by default, regardless of command-line arguments.

• seed – Random seed used for training (model initialization, dataset shuffling) and for evaluation environments. Defaults to 1000.

• resolution – Resolution of images (height, width) in data samples. Defaults to (224, 224).

• num_cams – Number of cameras for the cloud VLM in each data sample. Defaults to 2.

• max_state_dim – Maximum dimension of the state vector. Defaults to 32.

• max_action_dim – Maximum dimension of the action vector. Defaults to 32.

• action_chunk – Size of action chunk. Defaults to 50.

• loss_weighting – Dictionary mapping loss type names to their weights. Defaults to {“MSE”: 1, “CE”: 1}.

• num_workers – Number of workers for the dataloader. Defaults to 4.

• batch_size – Total batch size for training. If None, calculated from dataloader_batch_size * gradient_accumulation_steps.

• gradient_accumulation_steps – Number of gradient accumulation steps. Defaults to 1.

• dataloader_batch_size – Batch size used by the dataloader. If None, calculated from batch_size // gradient_accumulation_steps.

• prefetch_factor – Prefetch factor for the dataloader. If None, uses default.

• steps – Total number of training steps. Defaults to 100,000.

• log_freq – Frequency of logging in training iterations. Defaults to 200.

• save_checkpoint – Whether to save checkpoints during training. Defaults to True.

• save_freq – Frequency of checkpoint saving in training iterations. Checkpoints are saved every save_freq steps and after the last training step. Defaults to 20,000.

• use_policy_training_preset – If True, use optimizer and scheduler presets from the policy configuration. Defaults to False.

• optimizer – Configuration for the optimizer. Required if use_policy_training_preset is False.

• scheduler – Configuration for the learning rate scheduler. Required if use_policy_training_preset is False.

• wandb – Configuration for Weights & Biases logging. Defaults to WandBConfig().

• debug – If True, set logging level to DEBUG. Defaults to False.

• trace_nans – Enable anomaly detection for debugging NaN/Inf values. Warning: causes large computational overhead. Defaults to False.

• env – Optional environment configuration for evaluation. Defaults to None.

• eval – Configuration for evaluation settings. Defaults to EvalConfig().

• eval_freq – Frequency of evaluation in training steps. If 0, evaluation is disabled. Defaults to 0.

• last_checkpoint_only – If True, only evaluate the last checkpoint. Defaults to True.

• running_best_count – Number of “running best” checkpoints to keep, i.e. the checkpoint(s) that achieved the best metric so far. If 0, the feature is disabled (like eval_freq/val_freq). When >= 1, after an eval/validation step that hits a new best the current weights are saved (or, if that step already wrote a regular save_freq checkpoint, that directory is reused); older running-bests are evicted down to running_best_count, but one that coincided with a regular checkpoint is left to the normal retention logic. Running-bests are protected from last_checkpoint_only pruning, and work even when save_checkpoint is False (only the best-performing checkpoints are then written, so resume granularity is limited to running-best steps). Defaults to 0.

• running_best_metric – Which metric drives the running best. One of “auto”, “eval_success” (maximize sim-eval success rate), or “val_loss” (minimize validation loss). “auto” resolves to “eval_success” when sim eval is configured (env set and eval_freq > 0), otherwise “val_loss”. Only used when running_best_count >= 1. Defaults to “auto”.

• server – Configuration for the gRPC inference server. Defaults to ServerConfig().

• planner – Configuration for the high-level planner of the gRPC inference server (server.py). Disabled by default, in which case only the VLA policy is served. Defaults to PlannerConfig().

__init__(dataset_mixture: ~opentau.configs.default.DatasetMixtureConfig, policy: ~opentau.configs.policies.PreTrainedConfig | None = None, output_dir: ~pathlib.Path | None = None, job_name: str | None = None, resume: bool = False, seed: int | None = 1000, resolution: tuple[int, int] = (224, 224), num_cams: int = 2, max_state_dim: int = 32, max_action_dim: int = 32, action_chunk: int = 50, loss_weighting: dict[str, float] = <factory>, num_workers: int = 4, batch_size: int | None = None, gradient_accumulation_steps: int = 1, dataloader_batch_size: int | None = None, prefetch_factor: int | None = None, steps: int = 100000, log_freq: int = 200, save_checkpoint: bool = True, save_freq: int = 20000, use_policy_training_preset: bool = False, optimizer: ~opentau.optim.optimizers.OptimizerConfig | None = None, scheduler: ~opentau.optim.schedulers.LRSchedulerConfig | None = None, wandb: ~opentau.configs.default.WandBConfig = <factory>, debug: bool = False, trace_nans: bool = False, env: ~opentau.envs.configs.EnvConfig | None = None, eval: ~opentau.configs.default.EvalConfig | None = <factory>, eval_freq: int = 0, val_freq: int = 0, last_checkpoint_only: bool = True, running_best_count: int = 0, running_best_metric: str = 'auto', server: ~opentau.configs.deployment.ServerConfig = <factory>, planner: ~opentau.configs.deployment.PlannerConfig = <factory>) → None

action_chunk: int = 50

batch_size: int | None = None

dataloader_batch_size: int | None = None

dataset_mixture: DatasetMixtureConfig

debug: bool = False

env: EnvConfig | None = None

eval: EvalConfig | None

eval_freq: int = 0

classmethod from_pretrained(pretrained_name_or_path: str | Path, *, force_download: bool = False, resume_download: bool | None = None, proxies: dict | None = None, token: str | bool | None = None, cache_dir: str | Path | None = None, local_files_only: bool = False, revision: str | None = None, **kwargs) → TrainPipelineConfig

Load a training configuration from a pretrained model or local path.

Parameters:

• cls – The class to instantiate.

• pretrained_name_or_path –

  Can be either:

  • A string, the model id of a pretrained config hosted inside a model repo on huggingface.co.

  • A path to a directory containing a configuration file saved using the _save_pretrained method.

  • A path to a saved configuration JSON file.

• force_download – Whether to force (re-)downloading the config files and configuration from the HuggingFace Hub. Defaults to False.

• resume_download – Whether to resume downloading the config files. Defaults to None.

• proxies – Dictionary of proxies to use for requests. Defaults to None.

• token – The token to use as HTTP bearer authorization. If True, will use the token generated when running huggingface-cli login. Defaults to None.

• cache_dir – Path to a directory in which a downloaded pretrained model configuration should be cached. Defaults to None.

• local_files_only – Whether to only look at local files (i.e., do not try to download the config). Defaults to False.

• revision – The specific model version to use. It can be a branch name, a tag name, or a commit id. Defaults to None.

• **kwargs – Additional keyword arguments passed to the parser.

Returns:

An instance of TrainPipelineConfig loaded from the specified path.

Raises:

FileNotFoundError – If the configuration file is not found on the HuggingFace Hub or in the local path.

gradient_accumulation_steps: int = 1

job_name: str | None = None

last_checkpoint_only: bool = True

log_freq: int = 200

loss_weighting: dict[str, float]

max_action_dim: int = 32

max_state_dim: int = 32

num_cams: int = 2

num_workers: int = 4

optimizer: OptimizerConfig | None = None

output_dir: Path | None = None

planner: PlannerConfig

policy: PreTrainedConfig | None = None

prefetch_factor: int | None = None

resolution: tuple[int, int] = (224, 224)

resume: bool = False

running_best_count: int = 0

running_best_metric: str = 'auto'

save_checkpoint: bool = True

save_freq: int = 20000

scheduler: LRSchedulerConfig | None = None

seed: int | None = 1000

server: ServerConfig

steps: int = 100000

to_dict() → dict

Convert the configuration to a dictionary.

Returns:

Dictionary representation of the configuration.

trace_nans: bool = False

use_policy_training_preset: bool = False

val_freq: int = 0

validate()

Validate and finalize the training configuration.

This method performs several validation and setup tasks: - Loads policy configuration from CLI arguments or pretrained path if specified - Sets up checkpoint paths for resuming training - Validates output directory and creates default if needed - Sets up optimizer and scheduler from presets if enabled - Updates policy configuration with training parameters

Raises:

• ValueError – If required configurations are missing or invalid.

• FileExistsError – If output directory exists and resume is False.

• NotADirectoryError – If config_path for resuming doesn’t exist locally.

wandb: WandBConfig

opentau.configs.train.warn(*args, **kwargs)

Print a warning message to stderr.

This function is used instead of logging.warning() to avoid setting the logger level to WARNING.

Parameters:

• *args – Variable length argument list to print.

• **kwargs – Arbitrary keyword arguments passed to print().