Add 200M Muon reconstruction-5 experiment

egomimic/hydra_configs/experiment/pusht/action_flow_bc_usocket_latent_fm_sg_recon5_200m_muon_lr1e5_s42.yaml

@@ -0,0 +1,14 @@
+# @package _global_
+defaults:
+  - /experiment/pusht/action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42
+  - _self_
+
+name: action_flow_bc_usocket_latent_fm_sg_recon5_200m_muon_lr1e5_s42
+description: "Option-A 200M latent FM stopgrad with 1M codecs, Muon+AdamW, the matched 1e-5 schedule, and 5x reconstruction loss"
+
+model:
+  reconstruction_weight: 5.0
+
+run_provenance:
+  objective:
+    reconstruction_weight: ${model.reconstruction_weight}

scripts/train/launch_action_flow_usocket.sbatch

@@ -122,6 +122,13 @@ case "$AF_EXPERIMENT" in
     AF_EXPECTED_CONFIG_NAME=action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42
     AF_FULL_RECONSTRUCTION_ONLY_WARMUP_STEPS=0
     ;;
+  pusht/action_flow_bc_usocket_latent_fm_sg_recon5_200m_muon_lr1e5_s42)
+    AF_METHOD=latent_fm_stopgrad
+    AF_RECONSTRUCTION_WEIGHT=5
+    AF_FLOW_WEIGHT=1
+    AF_EXPECTED_CONFIG_NAME=action_flow_bc_usocket_latent_fm_sg_recon5_200m_muon_lr1e5_s42
+    AF_FULL_RECONSTRUCTION_ONLY_WARMUP_STEPS=0
+    ;;
   pusht/action_flow_bc_usocket_bridge_likelihood_s42)
     AF_METHOD=gaussian_bridge_likelihood
     # Transport placeholders only; no FM/reconstruction flags enter this model.

scripts/train/verify_action_flow_training_smoke.py

@@ -68,6 +68,13 @@
 SCALED_MUON_EXPERIMENT = (
     "pusht/action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42"
 )
+SCALED_MUON_RECON5_EXPERIMENT = (
+    "pusht/action_flow_bc_usocket_latent_fm_sg_recon5_200m_muon_lr1e5_s42"
+)
+SCALED_MUON_EXPERIMENTS = {
+    SCALED_MUON_EXPERIMENT,
+    SCALED_MUON_RECON5_EXPERIMENT,
+}
 SCALED_MUON_PARAMETER_COUNT = 199_754_837
 APPROVED_EXPERIMENTS = {
     "pusht/action_flow_bc_usocket_latent_fm_sg_recon1_s42": (
@@ -85,6 +92,11 @@
         1.0,
         1.0,
     ),
+    SCALED_MUON_RECON5_EXPERIMENT: (
+        "action_flow_bc_usocket_latent_fm_sg_recon5_200m_muon_lr1e5_s42",
+        5.0,
+        1.0,
+    ),
     "pusht/action_flow_bc_usocket_bridge_likelihood_s42": (
         "action_flow_bc_usocket_bridge_likelihood_s42",
         0.0,
@@ -319,7 +331,7 @@ def _validate_config(
     _require(
         targets == method_stage_targets(method), f"unexpected stage topology: {targets}"
     )
-    scaled_muon = experiment == SCALED_MUON_EXPERIMENT
+    scaled_muon = experiment in SCALED_MUON_EXPERIMENTS
     field_hidden_dim = 1_024 if scaled_muon else 512
     field_depth = 14 if scaled_muon else 12
     field_num_heads = 16 if scaled_muon else 8
@@ -991,7 +1003,10 @@ def _validate_optimizer_state(
     scaled_muon = (
         config is not None
         and str(config.get("name", ""))
-        == APPROVED_EXPERIMENTS[SCALED_MUON_EXPERIMENT][0]
+        in {
+            APPROVED_EXPERIMENTS[experiment][0]
+            for experiment in SCALED_MUON_EXPERIMENTS
+        }
     )
     if not scaled_muon:
         _require(bool(optimizer_state.get("state")), "AdamW optimizer state is empty")
@@ -1161,9 +1176,9 @@ def _validate_checkpoint(
             CODEC98K_EXPERIMENT
         ][0]:
             expected_parameter_count = CODEC98K_PARAMETER_COUNT
-        if config is not None and str(config.get("name", "")) == APPROVED_EXPERIMENTS[
-            SCALED_MUON_EXPERIMENT
-        ][0]:
+        if config is not None and str(config.get("name", "")) in {
+            APPROVED_EXPERIMENTS[item][0] for item in SCALED_MUON_EXPERIMENTS
+        }:
             expected_parameter_count = SCALED_MUON_PARAMETER_COUNT
         _require(
             parameter_count == expected_parameter_count,

tests/test_action_flow_launcher.py

@@ -71,6 +71,7 @@ def test_launcher_accepts_only_the_approved_sweep_and_pins_training_semantics():
     assert "pusht/action_flow_bc_usocket_recon100_s42" in source
     assert "pusht/action_flow_bc_usocket_latent_fm_sg_recon1_codec98k_s42" in source
     assert "pusht/action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42" in source
+    assert "pusht/action_flow_bc_usocket_latent_fm_sg_recon5_200m_muon_lr1e5_s42" in source
     assert "AF_EXPECTED_CONFIG_NAME=action_flow_bc_usocket_recon1_s42" in source
     assert "AF_EXPECTED_CONFIG_NAME=action_flow_bc_usocket_recon10_s42" in source
     assert "AF_EXPECTED_CONFIG_NAME=action_flow_bc_usocket_recon100_s42" in source

tests/test_validate_action_flow_config.py

@@ -113,6 +113,23 @@ def test_option_a_200m_muon_config_has_exact_capacity_optimizer_and_schedule():
     assert optimization["parameter_groups"]["disjoint"] is True
 
 
+def test_option_a_200m_muon_recon5_changes_only_reconstruction_weight():
+    baseline, _ = preflight.validate_experiment(
+        "pusht/action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42",
+        config_root=CONFIG_ROOT,
+    )
+    recon5, _ = preflight.validate_experiment(
+        "pusht/action_flow_bc_usocket_latent_fm_sg_recon5_200m_muon_lr1e5_s42",
+        config_root=CONFIG_ROOT,
+    )
+
+    assert recon5["status"] == "PASS"
+    assert recon5["parameters"] == baseline["parameters"]
+    assert recon5["optimization"] == baseline["optimization"]
+    assert recon5["objective"]["reconstruction_weight"] == pytest.approx(5.0)
+    assert baseline["objective"]["reconstruction_weight"] == pytest.approx(1.0)
+
+
 def test_resolved_hash_is_stable_and_uses_runtime_sentinels():
     first = preflight.compose_experiment(
         "pusht/action_flow_bc_usocket_recon1_s42", config_root=CONFIG_ROOT

tests/test_verify_action_flow_training_smoke.py

@@ -597,6 +597,25 @@ def test_cli_accepts_option_a_200m_muon_experiment():
     assert args.experiment.endswith("recon1_200m_muon_lr1e5_s42")
 
 
+def test_cli_accepts_option_a_200m_muon_recon5_experiment():
+    args = MODULE._parser().parse_args(
+        [
+            "/tmp/run",
+            "--expected-head",
+            HEAD,
+            "--expected-experiment",
+            "pusht/action_flow_bc_usocket_latent_fm_sg_recon5_200m_muon_lr1e5_s42",
+            "--expected-reconstruction-weight",
+            "5",
+            "--expected-preflight-sha256",
+            "d" * 64,
+        ]
+    )
+
+    assert args.experiment.endswith("recon5_200m_muon_lr1e5_s42")
+    assert args.expected_reconstruction_weight == 5.0
+
+
 def test_gpu_probe_gate_requires_real_single_h100_or_h200_bf16(tmp_path):
     path = tmp_path / "provenance/restart-0/gpu_probe.json"
     path.parent.mkdir(parents=True)
@@ -664,11 +683,16 @@ def _scaled_muon_optimizer_state():
     }
 
 
-def test_optimizer_state_gate_accepts_scaled_muon_composite_state():
+@pytest.mark.parametrize(
+    "experiment_name",
+    [
+        "action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42",
+        "action_flow_bc_usocket_latent_fm_sg_recon5_200m_muon_lr1e5_s42",
+    ],
+)
+def test_optimizer_state_gate_accepts_scaled_muon_composite_state(experiment_name):
     config = OmegaConf.create(
-        {
-            "name": "action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42"
-        }
+        {"name": experiment_name}
     )
 
     MODULE._validate_optimizer_state(_scaled_muon_optimizer_state(), config)

tools/validate_action_flow_config.py

@@ -93,13 +93,20 @@
 LIKELIHOOD_METHOD = "gaussian_bridge_likelihood"
 GRAPH_METHOD = "graph_section_diagnostic"
 STOPGRAD_METHOD = "latent_fm_stopgrad"
-SCALED_MUON_CONFIG_NAME = (
-    "action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42"
+SCALED_MUON_RECON5_CONFIG_NAME = (
+    "action_flow_bc_usocket_latent_fm_sg_recon5_200m_muon_lr1e5_s42"
+)
+SCALED_MUON_CONFIG_NAMES = frozenset(
+    {
+        "action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42",
+        SCALED_MUON_RECON5_CONFIG_NAME,
+    }
 )
 CANDIDATE_METHODS = {
     "pusht/action_flow_bc_usocket_latent_fm_sg_recon1_s42": STOPGRAD_METHOD,
     "pusht/action_flow_bc_usocket_latent_fm_sg_recon1_codec98k_s42": STOPGRAD_METHOD,
     "pusht/action_flow_bc_usocket_latent_fm_sg_recon1_200m_muon_lr1e5_s42": STOPGRAD_METHOD,
+    "pusht/action_flow_bc_usocket_latent_fm_sg_recon5_200m_muon_lr1e5_s42": STOPGRAD_METHOD,
     "pusht/action_flow_bc_usocket_bridge_likelihood_s42": LIKELIHOOD_METHOD,
     "pusht/action_flow_bc_usocket_graph_section_s42": GRAPH_METHOD,
 }
@@ -167,8 +174,15 @@ def validate_method_contract(config: DictConfig, experiment: str | None = None)
             "all_stopgrad",
             "FM-only reference detachment",
         )
+        expected_reconstruction_weight = (
+            5.0
+            if str(config.name) == SCALED_MUON_RECON5_CONFIG_NAME
+            else 1.0
+        )
         _float(
-            config.model.reconstruction_weight, 1.0, "candidate reconstruction weight"
+            config.model.reconstruction_weight,
+            expected_reconstruction_weight,
+            "candidate reconstruction weight",
         )
     elif method == GRAPH_METHOD:
         _exact(
@@ -512,7 +526,7 @@ def _validate_dimensions_and_modules(
     field = field_stage.field
     codec_profile = (int(encoder.hidden_dim), int(encoder.feedforward_dim))
     config_name = str(config.name)
-    if config_name == SCALED_MUON_CONFIG_NAME:
+    if config_name in SCALED_MUON_CONFIG_NAMES:
         expected_codec_profile = (204, 816)
     elif config_name == "action_flow_bc_usocket_latent_fm_sg_recon1_codec98k_s42":
         expected_codec_profile = (44, 176)
@@ -553,7 +567,7 @@ def _validate_dimensions_and_modules(
         "decoder g context-free forward signature",
     )
 
-    scaled_muon = config_name == SCALED_MUON_CONFIG_NAME
+    scaled_muon = config_name in SCALED_MUON_CONFIG_NAMES
     field_expected = {
         "input_dim": 8,
         "output_dim": 8,
@@ -621,7 +635,7 @@ def _validate_dimensions_and_modules(
         )
         _require(
             float(config.model.reconstruction_weight)
-            in ({0.0} if method == GRAPH_METHOD else {1.0, 10.0, 100.0}),
+            in ({0.0} if method == GRAPH_METHOD else {1.0, 5.0, 10.0, 100.0}),
             "unsupported reconstruction weight",
         )
 
@@ -731,7 +745,7 @@ def _validate_topology(
 
 def _validate_optimization(config: DictConfig) -> dict[str, Any]:
     optimizer = config.model.optimizer
-    scaled_muon = str(config.name) == SCALED_MUON_CONFIG_NAME
+    scaled_muon = str(config.name) in SCALED_MUON_CONFIG_NAMES
     expected_optimizer = (
         "egomimic.utils.unite_optim.ReleasedUniteCompositeOptimizer"
         if scaled_muon
@@ -1143,7 +1157,7 @@ def _validate_data_and_launch(
                 if action_flow_method(config) == GRAPH_METHOD
                 else (
                     {0: 0, 1: 13}
-                    if str(config.name) == SCALED_MUON_CONFIG_NAME
+                    if str(config.name) in SCALED_MUON_CONFIG_NAMES
                     else {0: 0, 1: 11}
                 )
             ),
@@ -1296,7 +1310,7 @@ def validate_config(
         if name != "pipeline_total"
     )
     _exact(accounted, parameters["pipeline_total"]["total"], "parameter accounting")
-    if str(config.name) == SCALED_MUON_CONFIG_NAME:
+    if str(config.name) in SCALED_MUON_CONFIG_NAMES:
         _exact(
             parameters["pipeline_total"]["total"],
             199_754_837,