configs: history and backbone-LR knobs on the keypoint wrist HPT model

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

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