diff --git a/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml index 1b7f4922b..c992e14b6 100644 --- a/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml @@ -126,6 +126,11 @@ robomimic_model: input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} output_dim: 840 widths: [840] + # Proprio observation history (as hpt_bc_mecka_6d_300M's state_ee_pose + # stem): K tokens, learned zero-init time-step embedding, whole-past + # dropout. Must equal the data config's `proprio_history`; inert at 1. + history_len: 1 + history_dropout: 0.2 specs: random_horizon_masking: false cross_attn: @@ -159,6 +164,10 @@ robomimic_model: mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225] +# LR multiplier for the image encoders' pretrained backbones; 1.0 keeps the +# single optimizer param group (ModelWrapper.configure_optimizers). +backbone_lr_scale: 1.0 + optimizer: _target_: torch.optim.AdamW _partial_: true diff --git a/tests/unit/test_hpt_kp_model_config.py b/tests/unit/test_hpt_kp_model_config.py new file mode 100644 index 000000000..51a951397 --- /dev/null +++ b/tests/unit/test_hpt_kp_model_config.py @@ -0,0 +1,24 @@ +"""The keypoint HPT model's stem knobs compose and agree with their data config.""" + +from __future__ import annotations + +RECIPE = "train_zarr_mecka_kp_wrist_hpt" + + +def test_hpt_kp_model_carries_the_history_and_backbone_lr_knobs(compose_resolve): + cfg = compose_resolve("train_zarr_cartesian", ["model=hpt_bc_keypoints_wrist_300M"]) + stem = cfg.model.robomimic_model.stem_specs.human_bimanual.state_keypoints + assert stem.history_len == 1 + assert stem.history_dropout == 0.2 + assert cfg.model.backbone_lr_scale == 1.0 + + +# The two knobs are one setting in two files: the keymap decides how many +# proprio steps a sample carries, the stem how many tokens consume them, and +# HPTModel.stem_process raises when they disagree. Composition alone cannot see +# it, so the keypoint recipe this model pairs with is checked here. +def test_keymap_proprio_history_matches_the_stem_that_consumes_it(compose_resolve): + cfg = compose_resolve(RECIPE, []) + key_map = cfg.data.train_datasets.human_bimanual.resolver.key_map + stem = cfg.model.robomimic_model.stem_specs.human_bimanual.state_keypoints + assert key_map.get("proprio_history", 1) == stem.history_len