defaults:
  - system: go2
  - model: quadruped_world_model
  - loss: world_model
  - heads: all
  - _self_

run_name:

data:
  dataset_name:
  num_train_workers: 8
  num_test_workers: 4
  prefetch_factor: 2
  # Reference dataset size the run's "realised paper-scale fraction" is taken
  # against, logged as part of the run fingerprint. It is a config value rather
  # than a constant in code because "the paper's scale" is a claim about a
  # specific table, and a claim like that should be visible and editable.
  paper_scale_steps: 25000000

training:
  num_epochs: 2
  max_steps:  # will train until num_epochs, or max_steps (if given), whichever is first
  eval_interval: 10000  # evaluates and checkpoints at this interval
  learning_rate: 0.0002
  warmup_steps: 10000
  decay_steps: 400000
  end_learning_rate: 0.0001
  batch_size: 512
  training_data_ratio: 0.998  # portion of dataset to use for training; the rest is used for evaluation
  seed: 42
  # Clip gradients to this global norm. Null is upstream's and the paper's
  # behaviour: `grad_norm` is logged, never applied. See bead UWLab-436.19 --
  # the manipulation run diverges without it, and clipping bounds the rate of
  # the runaway while `model.encoder.latent_norm` removes its cause.
  grad_clip_norm:

checkpoint:
  # create a checkpoint in this directory with the given run name, if resume_checkpoint is not set
  enabled: True
  # keeps this many most recent checkpoints; leave blank to keep all
  max_to_keep: 3
  # if non-empty, resume from this checkpoint
  resume_checkpoint:

# Read-only diagnostics run at every `training.eval_interval`. Nothing under
# this block can change the training gradient: every forward it does is
# `deterministic=True` and none of it owns an optimizer over the world model.
eval:
  # Held-out batches folded into the streaming metrics (reward/value r2, latent
  # rollout error, BC error, latent variance). Bounded so the diagnostics cost
  # does not scale with the size of the test split.
  max_batches: 8
  # Points kept for the predicted-vs-actual next-latent scatter.
  scatter_points: 5000

  episodes:
    # The episode-level probe: `eval/value_auc_success` and the paper's Fig. 5
    # curves. Disables itself with a printed reason when the processed dataset
    # has no `episode_ids.hdf5` (the quadruped processor writes none).
    enabled: True
    # Whole episodes to reserve from BOTH splits, so the AUC has no step-level
    # leak. 0 keeps upstream's chunk-level `random_split` bit-identical, at the
    # cost of an AUC measured on episodes the model has seen other steps of.
    holdout: 0
    num_auc_episodes: 32
    num_curve_episodes: 4
    # Where in an episode the "matched timestep" sits. Matched by PROGRESS, not
    # by absolute step: a successful insertion terminates early, so comparing at
    # absolute step k would measure episode length rather than value.
    match_frac: 0.5
    # Every n-th held-out chunk along a Fig. 5 curve.
    curve_stride: 5
    batch_size: 32
    success:
      # The processed dataset carries no success flag, so the label is derived.
      # `episode_return` | `final_reward`.
      metric: episode_return
      # Null splits at the median of the probed episodes -- needs no knowledge
      # of the reward scale and balances the classes, but is not comparable
      # across runs. Set an absolute value once a generation run's reward scale
      # is known.
      threshold:

  diagnostics:
    # PNG + npz written to <checkpoint_dir>/<dirname>/step_<N>/ at every eval.
    enabled: True
    dirname: diagnostics

  decoder_probe:
    # OPTIONAL and VISUALISATION ONLY (paper Fig. 9). A separate MLP with its
    # own optimizer, fed stop_gradient latents. It is never a training loss:
    # paper Table I measures reconstruction-as-an-objective at 0.90 -> 0.32
    # manipulation success.
    enabled: False
    hidden_size: 128
    learning_rate: 0.001
    steps_per_eval: 20

wandb:
  # METRICS ONLY. No artifacts, no checkpoint upload, no source snapshot --
  # checkpoints stay on local disk and are moved to Drive/HF by hand at the end.
  # Enforced in code by simdist/utils/wandb_utils.py and by
  # tests/test_wandb_metrics_only.py, which fails if an upload API appears
  # anywhere in the package.
  log: False
  project: simdist_train_model
  entity:
  # online | offline | disabled
  mode: online
  # Local wandb run directory; null uses wandb's default (./wandb).
  dir: