File size: 4,455 Bytes
1c5b1f4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | 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: |