Skip to content

ML Configuration

The experiment configuration is a strict, frozen dataclass contract. Unknown JSON fields and incompatible settings are rejected before data loading or model execution begins. Mapping-valued fields remain ordinary dictionaries and should be treated as read-only. Start from the bundled configs/ml_dry_run_cpu.json, configs/ml_dry_run_gpu.json, or configs/ml_baseline.json rather than constructing the full schema from memory.

Loading

load_experiment_config 

load_experiment_config(
    path: str | Path,
) -> ExperimentConfig

Load and validate an experiment JSON file.

Parameters:

  • path (str | Path) –

    UTF-8 JSON file containing the complete experiment configuration.

Returns:

  • ExperimentConfig –

    The strictly typed and cross-field-validated configuration.

Raises:

  • OSError –

    If the file cannot be read.

  • JSONDecodeError –

    If the file is not valid JSON.

  • TypeError –

    If a value has the wrong JSON type.

  • ValueError –

    If fields are missing, unknown, or mutually incompatible.

Examples:

Load the CPU dry-run configuration and inspect its model context:

from batgrad.ml.config import load_experiment_config

config = load_experiment_config("configs/ml_dry_run_cpu.json")
print(config.run.device, config.loader.seq_len)

parse_experiment_config 

parse_experiment_config(raw: object) -> ExperimentConfig

Strictly parse a decoded JSON-compatible experiment object.

Unknown fields are rejected rather than ignored, and nested lists and objects are converted to their immutable dataclass representations.

Parameters:

  • raw (object) –

    Decoded JSON-compatible root object.

Returns:

Raises:

  • TypeError –

    If the root or a nested value has the wrong type.

  • ValueError –

    If required fields are absent or validation fails.

resolve_store_root 

resolve_store_root(configured: str | None) -> str

Resolve the data-store root from configuration or DATA_ROOT.

Parameters:

  • configured (str | None) –

    Explicit root. When None, the DATA_ROOT environment variable is used.

Returns:

  • str –

    A non-empty data-store root.

Raises:

  • ValueError –

    If neither source provides a non-empty value.

Root Contract

ExperimentConfig dataclass 

ExperimentConfig(
    data: DataConfig,
    loader: LoaderTrainConfig,
    model: SequenceMixerConfig,
    train: TrainConfig,
    validation: ValidationConfig,
    optim: OptimizerConfig,
    scheduler: SchedulerConfig,
    run: RunConfig,
    logging: LoggingConfig,
    checkpoint: CheckpointConfig = CheckpointConfig(),
)

Validated, frozen dataclass contract for one ML experiment.

Constructing this object validates relationships spanning data, loader, model, objective, validation, device, output, and checkpoint settings. JSON callers should normally use load_experiment_config or parse_experiment_config so nested objects are coerced strictly.

Attributes:

Note

Frozen dataclasses prevent field reassignment, but mapping-valued fields remain ordinary dictionaries and should be treated as read-only.

Data And Loading

DataConfig dataclass 

DataConfig(
    manifest_paths: dict[str, str],
    protocols: tuple[str, ...],
    input_columns: tuple[str, ...],
    target_columns: tuple[str, ...],
    protocol_mode: Literal[
        "strict", "available"
    ] = "available",
    store_root: str | None = None,
    feedback_columns: tuple[str, ...] = (),
    scaling: tuple[ScalingRuleConfig, ...] = (),
)

Dataset selection and model column contract.

Attributes:

  • manifest_paths (dict[str, str]) –

    Mapping from canonical normalized manifest paths to the expected Git commit or commit prefix.

  • protocols (tuple[str, ...]) –

    Protocol values included in the experiment.

  • input_columns (tuple[str, ...]) –

    Ordered model input columns.

  • target_columns (tuple[str, ...]) –

    Ordered model target columns.

  • protocol_mode (Literal['strict', 'available']) –

    "strict" requires every selected dataset to contain every protocol; "available" retains the requested protocols that exist and warns about missing combinations.

  • store_root (str | None) –

    Data-store root, or None to use DATA_ROOT.

  • feedback_columns (tuple[str, ...]) –

    Columns that are both inputs and targets and may receive model predictions during rollout.

  • scaling (tuple[ScalingRuleConfig, ...]) –

    Scaling contracts for selected columns. Every input and target requires exactly one rule.

ScalingRuleConfig dataclass 

ScalingRuleConfig(
    column: str,
    input_min: float,
    input_max: float,
    output_min: float = -1.0,
    output_max: float = 1.0,
    clip: bool = False,
    transform: ScalingTransform = "linear",
)

Serializable scaling rule for one selected column.

Linear scaling maps [input_min, input_max] to [output_min, output_max]. The log1p transform is applied before the linear map and consequently requires non-negative input bounds.

Attributes:

  • column (str) –

    Canonical selected column name.

  • input_min (float) –

    Lower bound in physical units.

  • input_max (float) –

    Upper bound in physical units.

  • output_min (float) –

    Lower model-space bound.

  • output_max (float) –

    Upper model-space bound.

  • clip (bool) –

    Whether forward scaling clips to the output bounds.

  • transform (ScalingTransform) –

    Pre-scaling transform.

Note

Manifest validation still rejects observed values outside the configured input bounds when clip is enabled. Clipping protects runtime values; it does not relax the dataset contract.

LoaderTrainConfig dataclass 

LoaderTrainConfig(
    batch_size: int = 32,
    seq_len: int = 1024,
    strategy: Literal[
        "sequential", "shuffled_protocol_groups"
    ] = "shuffled_protocol_groups",
    stateful_n_windows: int = 1,
    cross_protocol_state_carry: Literal["chain"]
    | None = None,
    data_access: Literal[
        "windowed", "full_in_mem"
    ] = "windowed",
    num_workers: int = 0,
    prefetch_to_device: bool = False,
)

Training-loader settings stored in an experiment configuration.

Attributes:

  • batch_size (int) –

    Number of stream lanes per batch.

  • seq_len (int) –

    Input context length before any roll-forward extension.

  • strategy (Literal['sequential', 'shuffled_protocol_groups']) –

    Sequential windows or shuffled protocol-group streams.

  • stateful_n_windows (int) –

    Consecutive windows per stateful group, or -1 for whole-stream mode.

  • cross_protocol_state_carry (Literal['chain'] | None) –

    "chain" permits aligned protocol streams to share recurrent state.

  • data_access (Literal['windowed', 'full_in_mem']) –

    Windowed parquet reads or a full CPU tensor cache.

  • num_workers (int) –

    PyTorch data-loader worker count.

  • prefetch_to_device (bool) –

    Whether to asynchronously prefetch to the CUDA device.

Note

full_in_mem trades RAM for throughput and requires num_workers=0. Whole-stream batches are limited by the shortest lane and can omit longer lane tails. Finite stateful groups discard a final group shorter than stateful_n_windows. Shuffled protocol groups support cycling, HPPC, and RPT but not EIS.

Training And Validation

TrainConfig dataclass 

TrainConfig(
    epochs: float = 1.0,
    log_per_epoch: int = 10,
    log_every_steps: int | None = None,
    validate_per_epoch: int = 1,
    validate_every_steps: int | None = None,
    loss: Literal["categorical_ce"] = "categorical_ce",
    grad_clip_norm: float = 1.0,
    max_steps: int | None = None,
    masked_suffix: MaskedSuffixConfig = MaskedSuffixConfig(),
)

Optimization-loop controls.

Attributes:

  • epochs (float) –

    Fractional or whole passes over planned training batches.

  • log_per_epoch (int) –

    Derived logging frequency when log_every_steps is unset.

  • log_every_steps (int | None) –

    Optional explicit logging interval.

  • validate_per_epoch (int) –

    Derived validation frequency; zero disables validation.

  • validate_every_steps (int | None) –

    Optional explicit validation interval.

  • loss (Literal['categorical_ce']) –

    Training objective. Only "categorical_ce" is supported.

  • grad_clip_norm (float) –

    Maximum global gradient norm after AMP unscaling.

  • max_steps (int | None) –

    Optional hard step limit overriding the epoch-derived count.

  • masked_suffix (MaskedSuffixConfig) –

    Autoregressive training objective.

MaskedSuffixConfig dataclass 

MaskedSuffixConfig(
    enabled: bool = True,
    channels: tuple[str, ...] = (),
    suffix_steps: int = 128,
    loss_on_masked_only: bool = True,
    carry_mamba_state: bool = True,
    detach_between_windows: bool = True,
    roll_forward_steps: int = 0,
)

Autoregressive masked-suffix objective and rollout settings.

Attributes:

  • enabled (bool) –

    Whether feedback inputs in a suffix are predicted rather than supplied.

  • channels (tuple[str, ...]) –

    Feedback columns to mask and regenerate.

  • suffix_steps (int) –

    Maximum destination rows generated per model call.

  • loss_on_masked_only (bool) –

    Score only feedback targets in the suffix when true.

  • carry_mamba_state (bool) –

    Carry recurrent Mamba state between aligned windows.

  • detach_between_windows (bool) –

    Truncate gradients between roll-forward windows.

  • roll_forward_steps (int) –

    Additional training positions traversed with generated feedback.

Note

Generated feedback is detached before reuse. Enabled suffixes must be shorter than the loader sequence length.

ValidationConfig dataclass 

ValidationConfig(
    split: ValidationSplitConfig = ValidationSplitConfig(),
    max_tf_batches: int = 1,
    rollout_steps: int = 0,
    log_rollout_plots: bool = True,
    masked_suffix: ValidationMaskedSuffixConfig = ValidationMaskedSuffixConfig(),
    rollout_extension: RolloutExtensionConfig = RolloutExtensionConfig(),
)

Held-out-window and anchored-rollout validation configuration.

Attributes:

  • split (ValidationSplitConfig) –

    Group-aware train/validation split.

  • max_tf_batches (int) –

    Maximum held-out batches evaluated at each validation; zero disables this pass.

  • rollout_steps (int) –

    Number of observed future rows scored per anchor.

  • log_rollout_plots (bool) –

    Whether supported loggers receive trajectory plots.

  • masked_suffix (ValidationMaskedSuffixConfig) –

    Validation-time masked-suffix overrides.

  • rollout_extension (RolloutExtensionConfig) –

    Optional unscored continuation after observed rows.

ValidationSplitConfig dataclass 

ValidationSplitConfig(
    strategy: Literal[
        "sample", "provide", "merge"
    ] = "sample",
    fraction: float = 0.2,
    group_by: tuple[str, ...] = (
        "dataset id",
        "cell id",
        "cycle index",
    ),
    groups: tuple[ValidationGroupConfig, ...] = (),
)

Group-aware validation split policy.

sample deterministically hashes groups, provide uses only explicit selectors, and merge combines both sets. Splitting groups rather than rows reduces leakage between related measurements. Sampling selects int(group_count * fraction) groups, so small datasets can produce no sampled validation groups.

Attributes:

  • strategy (Literal['sample', 'provide', 'merge']) –

    Group selection strategy.

  • fraction (float) –

    Fraction selected by deterministic sampling.

  • group_by (tuple[str, ...]) –

    Manifest columns defining one indivisible group.

  • groups (tuple[ValidationGroupConfig, ...]) –

    Explicit group selectors used by provide and merge.

ValidationGroupConfig dataclass 

ValidationGroupConfig(
    match: dict[str, object],
    rollout_start_offsets: tuple[int, ...] = (),
)

Explicit held-out group and optional rollout anchors.

Attributes:

  • match (dict[str, object]) –

    Values matched against the validation split's grouping columns.

  • rollout_start_offsets (tuple[int, ...]) –

    Zero-based source-row offsets identifying the last observed input row at which anchored rollouts begin.

ValidationMaskedSuffixConfig dataclass 

ValidationMaskedSuffixConfig(
    enabled: bool | None = None,
    suffix_steps: int | None = None,
    carry_mamba_state: bool | None = None,
)

Validation overrides for masked-suffix execution.

None values inherit their training counterparts. Validation always disables training roll-forward, regardless of the training configuration.

Attributes:

  • enabled (bool | None) –

    Override masked-suffix validation.

  • suffix_steps (int | None) –

    Override the number of predicted suffix rows per call.

  • carry_mamba_state (bool | None) –

    Override recurrent state carry during validation.

RolloutExtensionConfig dataclass 

RolloutExtensionConfig(
    enabled: bool = False,
    steps: int = 0,
    input_values: dict[str, float] = dict(),
)

Unscored rollout continuation beyond observed rows.

Attributes:

  • enabled (bool) –

    Whether to append synthetic future control rows.

  • steps (int) –

    Number of unobserved rows to generate.

  • input_values (dict[str, float]) –

    Physical-unit values for known future input controls.

Note

Extension rows have no targets and never contribute to validation metrics. Inputs omitted from input_values repeat their final observed value.

Runtime And Outputs

OptimizerConfig dataclass 

OptimizerConfig(
    kind: Literal["adamw"] = "adamw",
    lr: float = 0.0001,
    weight_decay: float = 0.01,
    beta1: float = 0.9,
    beta2: float = 0.95,
    eps: float = 1e-08,
)

AdamW optimizer parameters.

Attributes:

  • kind (Literal['adamw']) –

    Optimizer implementation; currently only "adamw".

  • lr (float) –

    Peak learning rate.

  • weight_decay (float) –

    Decoupled weight decay.

  • beta1 (float) –

    First-moment coefficient.

  • beta2 (float) –

    Second-moment coefficient.

  • eps (float) –

    Numerical stability term.

SchedulerConfig dataclass 

SchedulerConfig(
    kind: Literal[
        "linear_warmup_cosine", "none"
    ] = "linear_warmup_cosine",
    warmup_ratio: float = 0.05,
    min_lr_ratio: float = 0.01,
)

Learning-rate scheduler parameters.

Attributes:

  • kind (Literal['linear_warmup_cosine', 'none']) –

    Cosine decay with linear warmup, or no scheduler.

  • warmup_ratio (float) –

    Fraction of total steps spent warming up.

  • min_lr_ratio (float) –

    Final cosine learning rate as a fraction of peak rate.

RunConfig dataclass 

RunConfig(
    device: str = "cuda",
    seed: int = 69,
    use_amp: bool = True,
    compile_model: bool = False,
    init_from: str | None = None,
    output_dir: str | None = "outputs/runs",
    name: str | None = None,
)

Runtime and output settings.

Attributes:

  • device (str) –

    PyTorch device string.

  • seed (int) –

    Random seed used by model initialization and data planning.

  • use_amp (bool) –

    Enable CUDA automatic mixed precision.

  • compile_model (bool) –

    Wrap the model with torch.compile.

  • init_from (str | None) –

    Optional checkpoint used to initialize compatible model weights.

  • output_dir (str | None) –

    Parent run directory, or None for output-free stdout runs.

  • name (str | None) –

    Stable run directory name. Existing directories with this name are replaced before training starts.

Warning

init_from is weight initialization, not training resume: optimizer, scheduler, scaler, and cursor state are not restored.

LoggingConfig dataclass 

LoggingConfig(
    backend: Literal["stdout", "jsonl", "wandb"] = "jsonl",
    mode: Literal["offline", "online"] = "offline",
    mirror_stdout: bool = True,
    wandb: WandbConfig = WandbConfig(),
)

Metric logging configuration.

Attributes:

  • backend (Literal['stdout', 'jsonl', 'wandb']) –

    Output backend. JSONL and W&B require a run output directory.

  • mode (Literal['offline', 'online']) –

    W&B synchronization mode; ignored by other backends.

  • mirror_stdout (bool) –

    Whether file and W&B backends also emit concise console metrics.

  • wandb (WandbConfig) –

    W&B-specific metadata.

WandbConfig dataclass 

WandbConfig(
    project: str | None = None,
    entity: str | None = None,
    group: str | None = None,
    name: str | None = None,
    tags: tuple[str, ...] = (),
)

Weights & Biases run metadata.

Attributes:

  • project (str | None) –

    W&B project name. The W&B client default is used when omitted.

  • entity (str | None) –

    Optional team or user account.

  • group (str | None) –

    Optional run group.

  • name (str | None) –

    Optional display name.

  • tags (tuple[str, ...]) –

    Searchable run tags.

CheckpointConfig dataclass 

CheckpointConfig(
    save_latest: bool = False,
    save_best: bool = False,
    save_final: bool = False,
    monitors: tuple[str, ...] = (),
)

Checkpoint persistence policy.

Attributes:

  • save_latest (bool) –

    Replace the latest-step checkpoint after validation events.

  • save_best (bool) –

    Keep the lowest value observed for each monitored metric at validation events. Unknown or unavailable metrics only emit a warning.

  • save_final (bool) –

    Save model and training state after the last step.

  • monitors (tuple[str, ...]) –

    Metric names minimized by best-checkpoint tracking.