diff --git a/egomimic/hydra_configs/experiment/pusht/action_flow_bc_usocket_latent_fm_sg_recon5_200m_muon_lr1e5_s42.yaml b/egomimic/hydra_configs/experiment/pusht/action_flow_bc_usocket_latent_fm_sg_recon5_200m_muon_lr1e5_s42.yaml new file mode 100644 index 000000000..bc42025d3 --- /dev/null +++ b/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} diff --git a/scripts/train/launch_action_flow_usocket.sbatch b/scripts/train/launch_action_flow_usocket.sbatch index 3e9ec368a..6f4b7dacb 100644 --- a/scripts/train/launch_action_flow_usocket.sbatch +++ b/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. diff --git a/scripts/train/verify_action_flow_training_smoke.py b/scripts/train/verify_action_flow_training_smoke.py index 29dc7842d..d4ef4b21d 100644 --- a/scripts/train/verify_action_flow_training_smoke.py +++ b/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, diff --git a/tests/test_action_flow_launcher.py b/tests/test_action_flow_launcher.py index d07920351..6866ce0fc 100644 --- a/tests/test_action_flow_launcher.py +++ b/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 diff --git a/tests/test_validate_action_flow_config.py b/tests/test_validate_action_flow_config.py index b2082b110..e35a27ca7 100644 --- a/tests/test_validate_action_flow_config.py +++ b/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 diff --git a/tests/test_verify_action_flow_training_smoke.py b/tests/test_verify_action_flow_training_smoke.py index 1a9215f28..92f95ac32 100644 --- a/tests/test_verify_action_flow_training_smoke.py +++ b/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) diff --git a/tools/validate_action_flow_config.py b/tools/validate_action_flow_config.py index 1bcf16f91..909e045e7 100644 --- a/tools/validate_action_flow_config.py +++ b/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,