diff --git a/docs/action-flow-usocket-candidates.md b/docs/action-flow-usocket-candidates.md new file mode 100644 index 000000000..202a6d22a --- /dev/null +++ b/docs/action-flow-usocket-candidates.md @@ -0,0 +1,104 @@ +# Conditional U-Socket ports of the torus candidates + +These are three research candidates, not three equivalent reconstruction-weight +ablations. All use U-Socket BC, fresh seed42, H16 normalized `[x,y,cosθ,sinθ]`, +latent8 per token, one observation frame, a shared 12×512 AdaLN field, private +context-free two-layer width20 codecs, AdamW3e-5, 8K warmup, cosine floor3e-6, +240K updates, batch32 and BF16 on one H100 or H200. No warm-start or +reconstruction-only phase. The generic pipeline runner is unchanged. + +## Configs and gradients + +### FM-only clean stop-gradient (method 2) + +Config: `pusht/action_flow_bc_usocket_latent_fm_sg_recon1_s42`. + +- Reconstruction: `A → E → g → MSE(A)`; updates E and g, weight1. +- Latent FM: `A → E → stop-gradient → latent bridge → v(z,t,O) → latent MSE`; + both the bridge endpoint and velocity target detach. Updates v and the + observation encoder, not E or g. +- Action Flow: `A → E → attached latent bridge → v → residual → J_g(residual)`; + squared residual mean updates E, v, g and the observation encoder, weight1. + +The two v forwards share time, Gaussian noise and conditioning-dropout mask. +Fourteen samples per content means 28 field sample-equivalents, not 14, for +this variant. A zero FM/reconstruction parameter intersection is intentional; +its gradient cosine is undefined, not evidence of broken training. +Inference: `Gaussian latent → reverse Euler16(v,O) → g → action`. + +### Learned Gaussian-bridge likelihood (method 3) + +Config: `pusht/action_flow_bc_usocket_bridge_likelihood_s42`. + +Private reference mean `mu(A,t)=(1-t)E_raw(A,t)` ends exactly at zero. The +reference stage, Gaussian noising stage, conditional reverse-mean transformer, +decoder and objective are separate nodes. + +Interior: `A → mu_k,mu_(k-1) → Gaussian z_k and attached reverse target → +M(z,k,O)=z+field(z,k/32,O) → equal-variance Gaussian KL`. +Boundary: `A → mu_1 + .1*noise → M(z,1,O) → g → action Gaussian NLL`. +Both mean and target gradients remain attached. Both terms update the private +mean encoder, field and observation encoder; the boundary also updates g. +There is no latent-FM, clean reconstruction, JVP or decoded-noise scale loss. + +K32, sigma1=.1, sigma32=1, geometric schedule, rho=.95, tau=.02. Draw 14 +interior levels/noises per content plus one independent boundary noise. The +batched field sees 15 sample-equivalents in one forward. Interior loss uses the +31× uniform-level correction and variance `(1-rho²)*sigma_(k-1)²`. + +All squared coordinate residuals are **summed over the full H16 chunk**, then +averaged over examples/level samples. Boundary weighting is therefore 80000× +per-coordinate action MSE for H16×D4, not the toy three-coordinate coefficient. +Logged likelihood terms omit parameter-independent Gaussian constants; they +are not exact reported marginal log-likelihood estimates. + +Inference: `z32~N(0,I) → 31 stochastic Gaussian reverse transitions → +final learned mean M(z1,1,O) → g → add tau*action_noise`. Exactly32 field calls, +31 latent innovations, and one action-noise draw. No deterministic Euler +substitution. Finite output noise models a smoothed action distribution. + +`Train/MSE` is decoded noisy-boundary-mean error; `Valid/MSE` is actual generated +action error. Neither is clean-codec reconstruction. Ordinary CFM trajectory/ +clean-latent/Jacobian diagnostics are disabled; normalized/native generated +errors and fixed-bank stochastic EnergyScore@32 remain required. + +### Exact graph section (method 4, restricted diagnostic) + +Config: `pusht/action_flow_bc_usocket_graph_section_s42`. + +`E(A)=(A,f(A))`; split each latent token into x4,h4 and decode +`g(x,h)=x + [R(x,h)-R(x,f(x))]`. The subtraction is grouped before adding x. +The same f parameter objects are used at both boundaries. A model-specific +Hydra factory injects them once: repeating a Hydra config is not weight sharing. +E owns only f; decoder-only R is not included in E's parameter set. + +Training: `A → E → latent bridge → v → latent FM`, plus +`latent residual → J_g → Action Flow loss`, both fully attached. No reconstruction +optimizer term; `g(E(A))=A` is checked and logged as a diagnostic. Inference uses +Gaussian latent → reverse Euler16 → g. Fixed-level/Jacobian diagnostics remain; +activation matching is disabled with an empty map and cknna_k=0. + +**R6 FAIL / restricted diagnostic:** explicit action coordinates are embedded in +the latent. Exact reconstruction is not a general flexible-interface solution +and is not evidence of good generation. + +## Comparison and operational boundaries + +Methods2/4 retain the historical U-Socket scale-regularizer weight0; the torus +winners used scale1. That difference is explicit, not a faithful toy ablation. +Conditional quality, scale stability and transfer remain OPEN. R2–R5/R8 are +implemented contracts, not demonstrated cross-embodiment transfer; R7 requires +observed GPU cost. R1/R9/R10 are not passed merely by finite smoke metrics. + +Use the maintained `scripts/train/launch_action_flow_usocket.sbatch` typed +entrypoint with exact source/runtime/config/data/norm bindings. Reuse hashed +2999-episode content, 2970/29 split, and train-only normalization receipts; +never recompute unchanged statistics to prepare another row. Validation10K, +immutable checkpoint40K, all scheduled/signal saves retained. Full training +requires that candidate's real optimizer+validation smoke and strict reload. + +An explicitly authorized ICE/Phoenix queue race needs distinct attempt paths and +tracked scheduler IDs. Keep the first allocation, confirm loser cancellation, +and never let two writers share W&B/output. Preserve the original queue deadline. +Restart testing must prove checkpoint-based optimizer continuation and the same +W&B ID; a passing ordinary smoke alone does not prove requeue continuity. diff --git a/docs/cluster-pipeline.md b/docs/cluster-pipeline.md index ff18b7a30..339534ac4 100644 --- a/docs/cluster-pipeline.md +++ b/docs/cluster-pipeline.md @@ -32,6 +32,12 @@ branch as its parent and carries shared runtime fixes without replacing its model code. Do not treat the synthetic Action Flow scripts on `main` as the real-data implementation or reset this stack to `main`. +The `codex/torus-winners-usocket-system-test-20260907` child adds three explicit +conditional candidate configs to that same launcher. Read +[their objective and sampler contracts](action-flow-usocket-candidates.md) +before selecting one. The likelihood arm is not ordinary latent FM; the exact +graph-section arm is a restricted diagnostic, not an unrestricted solution. + ## Cluster entry points | Host | Execution authority | Current-task discovery | diff --git a/egomimic/hydra_configs/experiment/pusht/action_flow_bc_usocket_bridge_likelihood_s42.yaml b/egomimic/hydra_configs/experiment/pusht/action_flow_bc_usocket_bridge_likelihood_s42.yaml new file mode 100644 index 000000000..6694d7be8 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pusht/action_flow_bc_usocket_bridge_likelihood_s42.yaml @@ -0,0 +1,44 @@ +# @package _global_ +defaults: + - /experiment/pusht/action_flow_usocket_candidate_common + - override /model: bf/us_action_flow_bc_bridge_likelihood + - _self_ + +name: action_flow_bc_usocket_bridge_likelihood_s42 +description: "Learned Gaussian bridge likelihood; K32 stochastic reverse chain; full H16 coordinate-sum reduction" +train: + action_flow_method: gaussian_bridge_likelihood + +evaluator: + # Straight-CFM clean endpoints, velocities and Euler trajectories are not + # defined for this model. Native MSE and stochastic EnergyScore32 remain on. + action_flow_diagnostics: + enabled: false + +run_provenance: + objective: + method: gaussian_bridge_likelihood + components: [interior_bridge_kl, action_boundary_nll] + target_mean_gradients: attached + interior_samples_per_content: ${model.interior_samples_per_content} + num_levels: ${model.num_levels} + sigma_min: ${model.sigma_min} + sigma_max: ${model.sigma_max} + rho: ${model.rho} + tau: ${model.tau} + reduction: sum_all_horizon_coordinates_then_batch_mean + interior_level_multiplier: 31 + # H16 × D4 / (2 × .02²) = 80000 times per-coordinate MSE. + boundary_per_coordinate_mse_multiplier: 80000.0 + reconstruction_only_warmup_steps: 0 + decoded_noise_scale_weight: 0.0 + inference: + sampler: gaussian_bridge_reverse_chain + steps: ${model.num_inference_steps} + latent_innovations: 31 + output_noise_std: ${model.tau} + classifier_free_guidance: false + requirements: + r6: PASS_design + conditional_quality: OPEN + observation_noise_caveat: finite_tau_models_a_smoothed_action_distribution diff --git a/egomimic/hydra_configs/experiment/pusht/action_flow_bc_usocket_graph_section_s42.yaml b/egomimic/hydra_configs/experiment/pusht/action_flow_bc_usocket_graph_section_s42.yaml new file mode 100644 index 000000000..0dc91b0ad --- /dev/null +++ b/egomimic/hydra_configs/experiment/pusht/action_flow_bc_usocket_graph_section_s42.yaml @@ -0,0 +1,72 @@ +# @package _global_ +defaults: + - /experiment/pusht/action_flow_usocket_candidate_common + - override /model: bf/us_action_flow_bc_graph_section + - _self_ + +name: action_flow_bc_usocket_graph_section_s42 +description: "Exact graph section; attached FM and Action Flow; R6 FAIL restricted diagnostic" +train: + action_flow_method: graph_section_diagnostic + +evaluator: + action_flow_diagnostics: + enabled: true + noise_seed_bank_path: ${hydra:runtime.cwd}/egomimic/hydra_configs/evaluator/energy_score_seed_bank_k32_v1.json + noise_seed_bank_sha256: 88657b829905d4374823db145ded19b99cec4735f76694734473bcee068bb5b6 + raw_noise_levels: [0.0, 0.25, 0.5, 0.75, 1.0] + max_batches_per_rank: 1 + max_samples: 16 + jacobian_samples: 2 + capture_activations: false + activation_layer_map: {} + cknna_k: 0 + native_error: + enabled: true + type: usocket_native_xy_wrapped_theta_mse_v1 + space: decoded_native_x_y_theta_radians + native_theta_index: 2 + wrap_period_radians: 6.283185307179586 + reduction: mean_squared_error_over_horizon_and_native_coordinates + normalizer: bound_train_only_evaluator_normalizer + native_decoder: egomimic.pipeline.pushshapes.USocketRotVecNativeDecoder + artifact_root: ${paths.output_dir}/validation_predictions/action_flow_diagnostics + validation_view: + definition: first_deterministic_validation_batch_per_rank + split_manifest_sha256: ${run_provenance.split_manifest_sha256} + per_rank_batch_size: 16 + world_size: 1 + provenance: + source_commit: ${run_provenance.source_commit} + normalization_sha256: ${run_provenance.normalization_sha256} + split_manifest_sha256: ${run_provenance.split_manifest_sha256} + dataset_content: + manifest_sha256: ${run_provenance.content_manifest_sha256} + aggregate_sha256: ${run_provenance.dataset_content_aggregate_sha256} + action_representation: x_y_cos_theta_sin_theta + clean_endpoint: t0 + gaussian_endpoint: t1 + latent_shape: [16, 8] + flow_samples_per_content: ${model.flow_samples_per_content} + condition_dropout_probability: ${model.condition_dropout_probability} + sampler: reverse_euler + sampler_steps: ${model.num_inference_steps} + fixed_level_coupling: one_clean_latent_and_one_gaussian_per_condition + decoder_jacobian_evaluation: declared_bridge_state_at_each_fixed_noise_level + activation_capture: disabled_for_graph_section + +run_provenance: + objective: + method: graph_section_diagnostic + flow_clean_gradient_mode: full + action_velocity_clean_gradient_mode: full + reconstruction_is_optimizer_objective: false + flow_weight: ${model.flow_weight} + reconstruction_weight: ${model.reconstruction_weight} + action_velocity_weight: 1.0 + flow_samples_per_content: ${model.flow_samples_per_content} + decoded_noise_scale_weight: 0.0 + monotonic_weight: 0.0 + requirements: + r6: FAIL_restricted_diagnostic + conditional_quality: OPEN diff --git a/egomimic/hydra_configs/experiment/pusht/action_flow_bc_usocket_latent_fm_sg_recon1_s42.yaml b/egomimic/hydra_configs/experiment/pusht/action_flow_bc_usocket_latent_fm_sg_recon1_s42.yaml new file mode 100644 index 000000000..86330f914 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pusht/action_flow_bc_usocket_latent_fm_sg_recon1_s42.yaml @@ -0,0 +1,72 @@ +# @package _global_ +defaults: + - /experiment/pusht/action_flow_usocket_candidate_common + - override /model: bf/us_action_flow_bc_latent_fm_sg + - _self_ + +name: action_flow_bc_usocket_latent_fm_sg_recon1_s42 +description: "FM-only endpoint detach; fully attached Action Flow; recon1" +train: + action_flow_method: latent_fm_stopgrad + +evaluator: + action_flow_diagnostics: + enabled: true + noise_seed_bank_path: ${hydra:runtime.cwd}/egomimic/hydra_configs/evaluator/energy_score_seed_bank_k32_v1.json + noise_seed_bank_sha256: 88657b829905d4374823db145ded19b99cec4735f76694734473bcee068bb5b6 + raw_noise_levels: [0.0, 0.25, 0.5, 0.75, 1.0] + max_batches_per_rank: 1 + max_samples: 16 + jacobian_samples: 2 + capture_activations: true + activation_layer_map: {0: 0, 1: 11} + cknna_k: 10 + native_error: + enabled: true + type: usocket_native_xy_wrapped_theta_mse_v1 + space: decoded_native_x_y_theta_radians + native_theta_index: 2 + wrap_period_radians: 6.283185307179586 + reduction: mean_squared_error_over_horizon_and_native_coordinates + normalizer: bound_train_only_evaluator_normalizer + native_decoder: egomimic.pipeline.pushshapes.USocketRotVecNativeDecoder + artifact_root: ${paths.output_dir}/validation_predictions/action_flow_diagnostics + validation_view: + definition: first_deterministic_validation_batch_per_rank + split_manifest_sha256: ${run_provenance.split_manifest_sha256} + per_rank_batch_size: 16 + world_size: 1 + provenance: + source_commit: ${run_provenance.source_commit} + normalization_sha256: ${run_provenance.normalization_sha256} + split_manifest_sha256: ${run_provenance.split_manifest_sha256} + dataset_content: + manifest_sha256: ${run_provenance.content_manifest_sha256} + aggregate_sha256: ${run_provenance.dataset_content_aggregate_sha256} + action_representation: x_y_cos_theta_sin_theta + clean_endpoint: t0 + gaussian_endpoint: t1 + latent_shape: [16, 8] + flow_samples_per_content: ${model.flow_samples_per_content} + condition_dropout_probability: ${model.condition_dropout_probability} + sampler: reverse_euler + sampler_steps: ${model.num_inference_steps} + fixed_level_coupling: one_clean_latent_and_one_gaussian_per_condition + decoder_jacobian_evaluation: declared_bridge_state_at_each_fixed_noise_level + activation_capture: real_forward_hooks_without_replay_or_replacement + +run_provenance: + objective: + method: latent_fm_stopgrad + flow_clean_gradient_mode: all_stopgrad + action_velocity_clean_gradient_mode: full + reconstruction_is_optimizer_objective: true + flow_weight: ${model.flow_weight} + reconstruction_weight: ${model.reconstruction_weight} + action_velocity_weight: 1.0 + flow_samples_per_content: ${model.flow_samples_per_content} + decoded_noise_scale_weight: 0.0 + monotonic_weight: 0.0 + requirements: + r6: PASS_design + conditional_quality: OPEN diff --git a/egomimic/hydra_configs/experiment/pusht/action_flow_usocket_candidate_common.yaml b/egomimic/hydra_configs/experiment/pusht/action_flow_usocket_candidate_common.yaml new file mode 100644 index 000000000..4a527b020 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pusht/action_flow_usocket_candidate_common.yaml @@ -0,0 +1,130 @@ +# @package _global_ +defaults: + - /experiment/pusht/planar_v2_base + - override /data: pusht/action_flow_usocket_val01_h16_common4 + - _self_ + +seed: 42 + +launch_params: {gpus_per_node: 1, nodes: 1} + +trainer: + max_steps: 240000 + max_epochs: -1 + min_epochs: null + limit_train_batches: 1.0 + val_check_interval: 10000 + limit_val_batches: 1.0 + check_val_every_n_epoch: null + num_sanity_val_steps: 0 + accumulate_grad_batches: 1 + strategy: auto + sync_batchnorm: false + log_every_n_steps: 10 + gradient_clip_val: 3.0 + gradient_clip_algorithm: norm + +norm_stats: + norm_mode: quantile + sample_frac: 1.0 + num_workers: 4 + precomputed_norm_path: null + +callbacks: + model_checkpoint: + filename: "epoch-{epoch}-step-{step}" + # A symlink avoids two byte-distinct serializations at the same semantic + # step, which the preemption runner correctly treats as a checkpoint fork. + save_last: link + save_top_k: -1 + every_n_train_steps: 40000 + every_n_epochs: null + +logger: + wandb: + project: pushshapes-action-flow + +planar: + action_dims: {pushshapes_sim_u_socket: 4} + action_horizon: 16 + observation_horizon: 1 + batch_size: 32 + eval_native_decoder: + _target_: egomimic.pipeline.pushshapes.USocketRotVecNativeDecoder + +evaluator: + native_decoder: ${planar.eval_native_decoder} + semantic_blocks: [[0, 2], [2, 4]] + # Typed EnergyScore distance over the complete H16 chunk. Translation uses + # normalized XY; rotation is recovered through the bound train normalizer and + # native atan2 decoder, wrapped on S1, and divided by pi to be dimensionless. + energy_score_distance: + type: usocket_normalized_xy_wrapped_theta_v1 + complete_normalized_chunk_shape: [16, 4] + normalized_translation_indices: [0, 1] + native_theta_index: 2 + rotation_scale_radians: 3.141592653589793 + semantic_weights: {translation: 0.5, rotation: 0.5} + native_decoder: egomimic.pipeline.pushshapes.USocketRotVecNativeDecoder + deterministic_seed: 420042 + energy_score_max_batches_per_rank: 1 + energy_score_validation_view: + definition: first_deterministic_validation_batch_per_ddp_rank + split_manifest_sha256: ${run_provenance.split_manifest_sha256} + per_rank_batch_size: 16 + world_size: 1 + energy_score_provenance: + source_commit: ${run_provenance.source_commit} + normalization_sha256: ${run_provenance.normalization_sha256} + split_manifest_sha256: ${run_provenance.split_manifest_sha256} + resolved_config_path: ${paths.output_dir}/.hydra/config.yaml + wandb: + entity: ${logger.wandb.entity} + project: ${logger.wandb.project} + run_id: ${logger.wandb.id} + distance_contract: ${evaluator.energy_score_distance} + dataset_content: + manifest_path: ${run_provenance.content_manifest_path} + manifest_sha256: ${run_provenance.content_manifest_sha256} + aggregate_sha256: ${run_provenance.dataset_content_aggregate_sha256} + action_representation: x_y_cos_theta_sin_theta + prediction_horizon: 16 + sampler: ${run_provenance.inference.sampler} + sampler_steps: ${model.num_inference_steps} + model_autocast_precision: bf16 + +run_provenance: + # The portable launcher replaces these with exact immutable identities before + # any optimizer/validation run. Null keeps read-only local Hydra composition + # deterministic without inventing a source or normalization identity. + source_commit: null + normalization_sha256: null + content_manifest_path: egomimic/hydra_configs/data/pusht/manifests/usocket_3000_v2_clean_content_v1.json + content_manifest_sha256: a1c81fb0ce8967aba795383a293180f9ba08a0ecfdd6f4a878afb20b39733761 + dataset_content_aggregate_sha256: 80f835ad37c3d5c5b7b2d5c3e1656c307ee567a1f63f51081165bf404b8ceb52 + split_manifest_path: egomimic/hydra_configs/data/pusht/planar_v2_usocket_dp_3k_split_seed42_v1.json + split_manifest_sha256: 3683e3461596eef8df2432fa865779b3c77b2a2057dabd0fea125595729cf313 + split_seed: 42 + valid_ratio: 0.01 + train_episode_count_per_domain: 2970 + valid_episode_count_per_domain: 29 + union_episode_count_per_domain: 2999 + id_overlap_count: 0 + resolved_path_overlap_count: 0 + dataset_observation_alignment: pre_step + action_contract: + representation: x_y_cos_theta_sin_theta + prediction_horizon: 16 + native_decoder: atan2 + inference: + sampler: reverse_euler + steps: ${model.num_inference_steps} + classifier_free_guidance: false + energy_score_contract: + sample_count: 32 + seed_bank_sha256: 88657b829905d4374823db145ded19b99cec4735f76694734473bcee068bb5b6 + distance: ${evaluator.energy_score_distance} + +deployment: + action_decoder: + _target_: egomimic.pipeline.pushshapes.USocketRotVecNativeDecoder diff --git a/egomimic/hydra_configs/model/bf/us_action_flow_bc_bridge_likelihood.yaml b/egomimic/hydra_configs/model/bf/us_action_flow_bc_bridge_likelihood.yaml new file mode 100644 index 000000000..0325e4700 --- /dev/null +++ b/egomimic/hydra_configs/model/bf/us_action_flow_bc_bridge_likelihood.yaml @@ -0,0 +1,120 @@ +# Context-free H16 codec around one observation-conditioned latent flow field. +_target_: egomimic.pl_utils.pl_model_action_flow_likelihood.ActionFlowLikelihoodModelWrapper + +action_horizon: 16 +action_dim: 4 +latent_dim: 8 +condition_dim: 67 +action_flow_method: gaussian_bridge_likelihood +interior_samples_per_content: 14 +num_levels: 32 +sigma_min: 0.1 +sigma_max: 1.0 +rho: 0.95 +tau: 0.02 +condition_dropout_probability: 0.3 +num_inference_steps: 32 + +pipeline: + _target_: egomimic.pipeline.algo.PipelineAlgo + stages: + - _target_: egomimic.pipeline.stages_sampler.FusedObsEncoder + n_obs_steps: 1 + inputs: + front_img_1: front_img_1 + state_agent_obj: state_agent_obj + encoder: + _target_: egomimic.pipeline.stages_sampler.DPStyleObsEncoder + obs_specs: + state_agent_obj: {input_dim: 3, input_slice: [0, 3]} + img_encoders: + front_img_1: + _target_: egomimic.models.stems.visual_core.VisualCore + in_channels: 3 + image_size: 96 + num_kp: 32 + feature_dimension: 64 + pretrained: false + crop_aug: true + crop_height: 84 + crop_width: 84 + crop_eval_mode: center + crop_sample_mode: v02 + crop_scope: frame + norm_layer: group + pool_type: spatial_softmax + - _target_: egomimic.pipeline.stages_sampler.GaussianLatentNoise + num_tokens: ${model.action_horizon} + latent_dim: ${model.latent_dim} + - _target_: egomimic.pipeline.stages_io.ActionTargetBuilder + - _target_: egomimic.pipeline.stages_action_flow_likelihood.LikelihoodReferenceStage + num_levels: ${model.num_levels} + interior_samples_per_content: ${model.interior_samples_per_content} + mean_encoder: + _target_: egomimic.models.action_flow_likelihood.TimeDependentSequenceMean + input_dim: ${model.action_dim} + latent_dim: ${model.latent_dim} + horizon: ${model.action_horizon} + hidden_dim: 20 + depth: 2 + num_heads: 4 + feedforward_dim: 80 + - _target_: egomimic.pipeline.stages_action_flow_likelihood.GaussianBridgeNoisingStage + num_levels: ${model.num_levels} + sigma_min: ${model.sigma_min} + sigma_max: ${model.sigma_max} + rho: ${model.rho} + condition_dropout_probability: ${model.condition_dropout_probability} + - _target_: egomimic.pipeline.stages_action_flow_likelihood.ConditionalReverseMeanStage + num_levels: ${model.num_levels} + sigma_min: ${model.sigma_min} + sigma_max: ${model.sigma_max} + rho: ${model.rho} + field: + _target_: egomimic.models.action_flow_transformer.AdaLNSequenceField + input_dim: ${model.latent_dim} + output_dim: ${model.latent_dim} + horizon: ${model.action_horizon} + condition_dim: ${model.condition_dim} + hidden_dim: 512 + depth: 12 + num_heads: 8 + feedforward_dim: 2048 + time_embedding_dim: 512 + time_scale: 1000.0 + dropout: 0.0 + condition_dropout_probability: ${model.condition_dropout_probability} + - _target_: egomimic.pipeline.stages_action_flow_likelihood.LikelihoodDecoderStage + tau: ${model.tau} + decoder: + _target_: egomimic.models.action_flow_codec.ContextFreeSequenceDecoder + latent_dim: ${model.latent_dim} + output_dim: ${model.action_dim} + horizon: ${model.action_horizon} + hidden_dim: 20 + depth: 2 + num_heads: 4 + feedforward_dim: 80 + dropout: 0.0 + - _target_: egomimic.pipeline.stages_action_flow_likelihood.GaussianBridgeObjectiveStage + num_levels: ${model.num_levels} + tau: ${model.tau} + +enable_grad_norm: false +gradient_telemetry_cadence: 100 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 3.0e-5 + betas: [0.9, 0.999] + eps: 1.0e-8 + weight_decay: 1.0e-4 + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 240000 + warmup_steps: 8000 + warmup_start_factor: 0.1 + eta_min: 3.0e-6 diff --git a/egomimic/hydra_configs/model/bf/us_action_flow_bc_graph_section.yaml b/egomimic/hydra_configs/model/bf/us_action_flow_bc_graph_section.yaml new file mode 100644 index 000000000..a5d097b7a --- /dev/null +++ b/egomimic/hydra_configs/model/bf/us_action_flow_bc_graph_section.yaml @@ -0,0 +1,101 @@ +# Context-free H16 codec around one observation-conditioned latent flow field. +_target_: egomimic.pl_utils.pl_model_action_flow.ActionFlowModelWrapper + +action_horizon: 16 +action_dim: 4 +latent_dim: 8 +condition_dim: 67 +flow_samples_per_content: 14 +condition_dropout_probability: 0.3 +num_inference_steps: 16 +reconstruction_weight: 0.0 +flow_weight: 1.0 +action_flow_method: graph_section_diagnostic + +pipeline: + _target_: egomimic.models.action_flow_graph.build_graph_section_pipeline + _recursive_: false + codec: + _target_: egomimic.models.action_flow_codec.GraphSectionSequenceCodec + action_dim: ${model.action_dim} + latent_dim: ${model.latent_dim} + horizon: ${model.action_horizon} + hidden_dim: 20 + depth: 2 + num_heads: 4 + feedforward_dim: 80 + dropout: 0.0 + stages: + - _target_: egomimic.pipeline.stages_sampler.FusedObsEncoder + n_obs_steps: 1 + inputs: + front_img_1: front_img_1 + state_agent_obj: state_agent_obj + encoder: + _target_: egomimic.pipeline.stages_sampler.DPStyleObsEncoder + obs_specs: + state_agent_obj: {input_dim: 3, input_slice: [0, 3]} + img_encoders: + front_img_1: + _target_: egomimic.models.stems.visual_core.VisualCore + in_channels: 3 + image_size: 96 + num_kp: 32 + feature_dimension: 64 + pretrained: false + crop_aug: true + crop_height: 84 + crop_width: 84 + crop_eval_mode: center + crop_sample_mode: v02 + crop_scope: frame + norm_layer: group + pool_type: spatial_softmax + - _target_: egomimic.pipeline.stages_sampler.GaussianLatentNoise + num_tokens: ${model.action_horizon} + latent_dim: ${model.latent_dim} + - _target_: egomimic.pipeline.stages_io.ActionTargetBuilder + - _target_: egomimic.pipeline.stages_action_flow.ContentEncoderStage + - _target_: egomimic.pipeline.stages_action_flow.LatentBridgeStage + samples_per_content: ${model.flow_samples_per_content} + condition_dropout_probability: ${model.condition_dropout_probability} + - _target_: egomimic.pipeline.stages_action_flow.ConditionalVelocityStage + num_inference_steps: ${model.num_inference_steps} + field: + _target_: egomimic.models.action_flow_transformer.AdaLNSequenceField + input_dim: ${model.latent_dim} + output_dim: ${model.latent_dim} + horizon: ${model.action_horizon} + condition_dim: ${model.condition_dim} + hidden_dim: 512 + depth: 12 + num_heads: 8 + feedforward_dim: 2048 + time_embedding_dim: 512 + time_scale: 1000.0 + dropout: 0.0 + condition_dropout_probability: ${model.condition_dropout_probability} + - _target_: egomimic.pipeline.stages_action_flow.ContentDecoderStage + - _target_: egomimic.pipeline.stages_action_flow.ActionFlowObjectiveStage + flow_weight: ${model.flow_weight} + reconstruction_weight: ${model.reconstruction_weight} + action_velocity_weight: 1.0 + +enable_grad_norm: false +gradient_telemetry_cadence: 100 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 3.0e-5 + betas: [0.9, 0.999] + eps: 1.0e-8 + weight_decay: 1.0e-4 + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 240000 + warmup_steps: 8000 + warmup_start_factor: 0.1 + eta_min: 3.0e-6 diff --git a/egomimic/hydra_configs/model/bf/us_action_flow_bc_latent_fm_sg.yaml b/egomimic/hydra_configs/model/bf/us_action_flow_bc_latent_fm_sg.yaml new file mode 100644 index 000000000..8b1e34687 --- /dev/null +++ b/egomimic/hydra_configs/model/bf/us_action_flow_bc_latent_fm_sg.yaml @@ -0,0 +1,112 @@ +# Context-free H16 codec around one observation-conditioned latent flow field. +_target_: egomimic.pl_utils.pl_model_action_flow.ActionFlowModelWrapper + +action_horizon: 16 +action_dim: 4 +latent_dim: 8 +condition_dim: 67 +flow_samples_per_content: 14 +condition_dropout_probability: 0.3 +num_inference_steps: 16 +reconstruction_weight: 1.0 +flow_weight: 1.0 +action_flow_method: latent_fm_stopgrad + +pipeline: + _target_: egomimic.pipeline.algo.PipelineAlgo + stages: + - _target_: egomimic.pipeline.stages_sampler.FusedObsEncoder + n_obs_steps: 1 + inputs: + front_img_1: front_img_1 + state_agent_obj: state_agent_obj + encoder: + _target_: egomimic.pipeline.stages_sampler.DPStyleObsEncoder + obs_specs: + state_agent_obj: {input_dim: 3, input_slice: [0, 3]} + img_encoders: + front_img_1: + _target_: egomimic.models.stems.visual_core.VisualCore + in_channels: 3 + image_size: 96 + num_kp: 32 + feature_dimension: 64 + pretrained: false + crop_aug: true + crop_height: 84 + crop_width: 84 + crop_eval_mode: center + crop_sample_mode: v02 + crop_scope: frame + norm_layer: group + pool_type: spatial_softmax + - _target_: egomimic.pipeline.stages_sampler.GaussianLatentNoise + num_tokens: ${model.action_horizon} + latent_dim: ${model.latent_dim} + - _target_: egomimic.pipeline.stages_io.ActionTargetBuilder + - _target_: egomimic.pipeline.stages_action_flow.ContentEncoderStage + encoder: + _target_: egomimic.models.action_flow_codec.ContextFreeSequenceEncoder + input_dim: ${model.action_dim} + latent_dim: ${model.latent_dim} + horizon: ${model.action_horizon} + hidden_dim: 20 + depth: 2 + num_heads: 4 + feedforward_dim: 80 + dropout: 0.0 + - _target_: egomimic.pipeline.stages_action_flow.LatentBridgeStage + samples_per_content: ${model.flow_samples_per_content} + condition_dropout_probability: ${model.condition_dropout_probability} + - _target_: egomimic.pipeline.stages_action_flow.ConditionalVelocityStage + flow_clean_gradient_mode: all_stopgrad + num_inference_steps: ${model.num_inference_steps} + field: + _target_: egomimic.models.action_flow_transformer.AdaLNSequenceField + input_dim: ${model.latent_dim} + output_dim: ${model.latent_dim} + horizon: ${model.action_horizon} + condition_dim: ${model.condition_dim} + hidden_dim: 512 + depth: 12 + num_heads: 8 + feedforward_dim: 2048 + time_embedding_dim: 512 + time_scale: 1000.0 + dropout: 0.0 + condition_dropout_probability: ${model.condition_dropout_probability} + - _target_: egomimic.pipeline.stages_action_flow.ContentDecoderStage + decoder: + _target_: egomimic.models.action_flow_codec.ContextFreeSequenceDecoder + latent_dim: ${model.latent_dim} + output_dim: ${model.action_dim} + horizon: ${model.action_horizon} + hidden_dim: 20 + depth: 2 + num_heads: 4 + feedforward_dim: 80 + dropout: 0.0 + - _target_: egomimic.pipeline.stages_action_flow.ActionFlowObjectiveStage + residual_key: action_flow/fm_velocity_residual + flow_weight: ${model.flow_weight} + reconstruction_weight: ${model.reconstruction_weight} + action_velocity_weight: 1.0 + +enable_grad_norm: false +gradient_telemetry_cadence: 100 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 3.0e-5 + betas: [0.9, 0.999] + eps: 1.0e-8 + weight_decay: 1.0e-4 + +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 240000 + warmup_steps: 8000 + warmup_start_factor: 0.1 + eta_min: 3.0e-6 diff --git a/egomimic/models/action_flow_codec.py b/egomimic/models/action_flow_codec.py index 13100a07b..36c51d559 100644 --- a/egomimic/models/action_flow_codec.py +++ b/egomimic/models/action_flow_codec.py @@ -207,3 +207,74 @@ def __init__( feedforward_dim=feedforward_dim, dropout=dropout, ) + + +class GraphSectionSequenceCodec(nn.Module): + """Restricted exact-section diagnostic with one shared learned graph. + + ``encode(A) = (A, f(A))`` and + ``decode(x, h) = x + R(x, h) - R(x, f(x))``. Both small sequence + Transformers are context-free. This keeps raw coordinates in the latent + and is a restricted diagnostic, not a general action-interface solution. + Construct this codec once and share it between encoder/decoder stages; + independently instantiating two copies would break the section identity. + """ + + def __init__( + self, + action_dim: int, + latent_dim: int, + horizon: int, + hidden_dim: int = 20, + depth: int = 2, + num_heads: int = 4, + feedforward_dim: int = 80, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.action_dim = int(action_dim) + self.latent_dim = int(latent_dim) + self.horizon = int(horizon) + self.input_dim = self.latent_dim + self.output_dim = self.action_dim + if self.action_dim <= 0 or self.latent_dim <= self.action_dim: + raise ValueError("graph section requires latent_dim > action_dim > 0") + if float(dropout) != 0.0: + raise ValueError("graph section requires dropout=0 for exact reconstruction") + self.graph = ContextFreeSequenceEncoder( + input_dim=self.action_dim, + latent_dim=self.latent_dim - self.action_dim, + horizon=self.horizon, + hidden_dim=hidden_dim, + depth=depth, + num_heads=num_heads, + feedforward_dim=feedforward_dim, + dropout=0.0, + ) + self.residual = ContextFreeSequenceDecoder( + latent_dim=self.latent_dim, + output_dim=self.action_dim, + horizon=self.horizon, + hidden_dim=hidden_dim, + depth=depth, + num_heads=num_heads, + feedforward_dim=feedforward_dim, + dropout=0.0, + ) + + def encode(self, content: torch.Tensor) -> torch.Tensor: + return torch.cat((content, self.graph(content)), dim=-1) + + def forward(self, latent: torch.Tensor) -> torch.Tensor: + if latent.ndim != 3 or tuple(latent.shape[1:]) != ( + self.horizon, + self.latent_dim, + ): + raise ValueError( + f"expected sequence shape (B, {self.horizon}, {self.latent_dim}), " + f"got {tuple(latent.shape)}" + ) + content = latent[..., : self.action_dim] + section = self.encode(content) + # Subtract before adding content so equal residuals cancel exactly. + return content + (self.residual(latent) - self.residual(section)) diff --git a/egomimic/models/action_flow_graph.py b/egomimic/models/action_flow_graph.py new file mode 100644 index 000000000..262e466f0 --- /dev/null +++ b/egomimic/models/action_flow_graph.py @@ -0,0 +1,49 @@ +"""Construct an exact-section graph with one shared private codec instance. + +Hydra interpolation of a module configuration constructs independent modules; +it does not share their weights. This model-specific factory binds the encoder +and decoder explicitly without introducing model semantics into PipelineAlgo. +""" + +from __future__ import annotations + +import hydra +import torch +import torch.nn as nn + +from egomimic.pipeline.algo import PipelineAlgo + + +class GraphSectionSequenceEncoder(nn.Module): + def __init__(self, codec: nn.Module): + super().__init__() + # E owns only f; decoder-only R must not appear in E.parameters(). + self.graph = codec.graph + + def forward(self, content): + return torch.cat((content, self.graph(content)), dim=-1) + + +def build_graph_section_pipeline(*, codec, stages, device=None): + """Instantiate declared nodes, injecting the same codec at both boundaries.""" + shared = hydra.utils.instantiate(codec) + bound_stages = [] + counts = {"encoder": 0, "decoder": 0} + for config in stages: + target = str(config.get("_target_", "")) + if target == "egomimic.pipeline.stages_action_flow.ContentEncoderStage": + if "encoder" in config: + raise ValueError("graph encoder is injected; do not configure a second codec") + stage = hydra.utils.instantiate(config, encoder=GraphSectionSequenceEncoder(shared)) + counts["encoder"] += 1 + elif target == "egomimic.pipeline.stages_action_flow.ContentDecoderStage": + if "decoder" in config: + raise ValueError("graph decoder is injected; do not configure a second codec") + stage = hydra.utils.instantiate(config, decoder=shared) + counts["decoder"] += 1 + else: + stage = hydra.utils.instantiate(config) + bound_stages.append(stage) + if counts != {"encoder": 1, "decoder": 1}: + raise ValueError(f"exact-section graph requires one encoder and decoder: {counts}") + return PipelineAlgo(bound_stages, device=device) diff --git a/egomimic/models/action_flow_likelihood.py b/egomimic/models/action_flow_likelihood.py new file mode 100644 index 000000000..e8ea48d48 --- /dev/null +++ b/egomimic/models/action_flow_likelihood.py @@ -0,0 +1,88 @@ +"""Context-free sequence reference means for a discrete Gaussian latent model. + +These are learned variational means, not ground-truth clean latents. The +terminal mean is exactly zero, so sigma_K=1 gives an action-independent prior. +""" + +from __future__ import annotations + +import math + +import torch +from torch import nn + +from egomimic.models.action_flow_codec import ContextFreeSequenceEncoder + + +class TimeDependentSequenceMean(nn.Module): + """mu(A,t)=(1-t) E_raw(A,t); no observation/context input exists.""" + + def __init__( + self, + input_dim: int = 4, + latent_dim: int = 8, + horizon: int = 16, + hidden_dim: int = 20, + depth: int = 2, + num_heads: int = 4, + feedforward_dim: int = 80, + dropout: float = 0.0, + ): + super().__init__() + self.input_dim = int(input_dim) + self.latent_dim = int(latent_dim) + self.horizon = int(horizon) + self.hidden_dim = int(hidden_dim) + self.depth = int(depth) + self.num_heads = int(num_heads) + self.feedforward_dim = int(feedforward_dim) + self.network = ContextFreeSequenceEncoder( + input_dim=self.input_dim + 1, + latent_dim=self.latent_dim, + horizon=self.horizon, + hidden_dim=hidden_dim, + depth=depth, + num_heads=num_heads, + feedforward_dim=feedforward_dim, + dropout=dropout, + ) + + def forward(self, action: torch.Tensor, time: torch.Tensor) -> torch.Tensor: + if action.ndim != 3 or tuple(action.shape[1:]) != ( + self.horizon, + self.input_dim, + ): + raise ValueError("action must have shape (B, horizon, input_dim)") + if time.shape != (len(action),) or time.device != action.device: + raise ValueError("time must have shape (B,) on the action device") + fraction = time.to(action).view(-1, 1, 1) + raw = self.network( + torch.cat((action, fraction.expand(-1, self.horizon, 1)), -1) + ) + # Compute the gate in FP32: t=1 remains exactly zero under autocast. + return (1.0 - time.float().view(-1, 1, 1)) * raw.float() + + +class GaussianBridgeSchedule(nn.Module): + """Fixed variances of the learned reference and reverse transitions.""" + + def __init__(self, num_levels=32, sigma_min=0.1, sigma_max=1.0, rho=0.95): + super().__init__() + self.num_levels = int(num_levels) + self.sigma_min, self.sigma_max, self.rho = map( + float, (sigma_min, sigma_max, rho) + ) + if self.num_levels < 2: + raise ValueError("num_levels must be at least two") + if not 0 < self.sigma_min <= self.sigma_max or self.sigma_max != 1.0: + raise ValueError( + "require 0 < sigma_min <= sigma_max=1 for the standard Gaussian prior" + ) + if not math.isfinite(self.rho) or not 0 <= self.rho < 1: + raise ValueError("rho must lie in [0, 1)") + self.register_buffer( + "sigmas", torch.logspace(math.log10(self.sigma_min), 0, self.num_levels) + ) + + def previous_variance(self, levels: torch.Tensor) -> torch.Tensor: + return (1.0 - self.rho**2) * self.sigmas[levels - 2].float().square() diff --git a/egomimic/pipeline/stages_action_flow.py b/egomimic/pipeline/stages_action_flow.py index a4c447f31..c36e515b5 100644 --- a/egomimic/pipeline/stages_action_flow.py +++ b/egomimic/pipeline/stages_action_flow.py @@ -191,7 +191,12 @@ def forward(self, batch: dict) -> dict: class ConditionalVelocityStage(Stage): - """Predict bridge velocity in training and integrate it during inference.""" + """Predict bridge velocity in training and integrate it during inference. + + ``all_stopgrad`` isolates only the latent-FM clean-state/target routes. + The original state, prediction, and residual remain fully attached for + the decoder JVP. Both field calls use the same sampled bridge and mask. + """ def __init__( self, @@ -209,12 +214,17 @@ def __init__( generated_latent_key: str = "action_flow/generated_latent", trajectory_key: str = "action_flow/trajectory", inference_steps_log_key: str = "log/action_flow_inference_steps", + flow_clean_gradient_mode: str = "full", + flow_residual_key: str = "action_flow/fm_velocity_residual", ): super().__init__() self.field = _module(field, label="field") self.num_inference_steps = int(num_inference_steps) if self.num_inference_steps <= 0: raise ValueError("num_inference_steps must be positive") + if flow_clean_gradient_mode not in {"full", "all_stopgrad"}: + raise ValueError("flow_clean_gradient_mode must be full|all_stopgrad") + self.flow_clean_gradient_mode = flow_clean_gradient_mode self.state_key = _key(state_key, label="state_key") self.time_key = _key(time_key, label="time_key") @@ -229,6 +239,9 @@ def __init__( predicted_velocity_key, label="predicted_velocity_key" ) self.residual_key = _key(residual_key, label="residual_key") + self.flow_residual_key = _key(flow_residual_key, label="flow_residual_key") + if len({self.predicted_velocity_key, self.residual_key, self.flow_residual_key}) != 3: + raise ValueError("prediction, residual, and FM residual keys must be distinct") self.inference_noise_key = _key( inference_noise_key, label="inference_noise_key" ) @@ -250,7 +263,11 @@ def __init__( self.condition_drop_mask_key, self.target_velocity_key, ) - self.writes = (self.predicted_velocity_key, self.residual_key) + self.writes = ( + self.predicted_velocity_key, + self.residual_key, + self.flow_residual_key, + ) self.reads_by_mode = { "inference": (self.inference_noise_key, self.inference_condition_key) } @@ -314,6 +331,16 @@ def _forward_train(self, batch: dict) -> dict: prediction = self._predict(state, time, condition, drop_mask) batch[self.predicted_velocity_key] = prediction batch[self.residual_key] = prediction - target_velocity + if self.flow_clean_gradient_mode == "all_stopgrad": + # The bridge consists only of the learned clean endpoint and + # action-independent Gaussian noise. Detaching its state and + # target removes both clean routes from FM, not from Action Flow. + flow_prediction = self._predict( + state.detach(), time, condition, drop_mask + ) + batch[self.flow_residual_key] = flow_prediction - target_velocity.detach() + else: + batch[self.flow_residual_key] = batch[self.residual_key] return batch def _forward_inference(self, batch: dict) -> dict: diff --git a/egomimic/pipeline/stages_action_flow_likelihood.py b/egomimic/pipeline/stages_action_flow_likelihood.py new file mode 100644 index 000000000..e6322b6a9 --- /dev/null +++ b/egomimic/pipeline/stages_action_flow_likelihood.py @@ -0,0 +1,415 @@ +"""Discrete learned Gaussian-bridge likelihood stages; not a CFM/Euler model. + +The objective is the fixed-variance negative ELBO up to parameter-independent +constants: all reference-mean and target gradients remain attached. Coordinates +are summed across the complete chunk, then examples/sampled levels are averaged. +""" + +from __future__ import annotations + +import math + +import torch +from torch import nn + +from egomimic.models.action_flow_likelihood import GaussianBridgeSchedule +from egomimic.pipeline.core import Stage + + +def _tensor(batch, key): + value = batch[key] + if not torch.is_tensor(value): + raise TypeError(f"{key} must be a tensor") + return value + + +class LikelihoodReferenceStage(Stage): + """Sample integer levels and evaluate attached private reference means.""" + + train_only = True + + def __init__( + self, + mean_encoder: nn.Module, + num_levels=32, + interior_samples_per_content=14, + input_key="target", + prefix="likelihood/", + ): + super().__init__() + if not isinstance(mean_encoder, nn.Module): + raise TypeError("mean_encoder must be a module") + self.mean_encoder = mean_encoder + self.num_levels = int(num_levels) + self.interior_samples_per_content = int(interior_samples_per_content) + if self.num_levels < 2 or self.interior_samples_per_content < 1: + raise ValueError("need >=2 levels and >=1 sampled interior per example") + self.input_key, self.prefix = input_key, prefix + self.reads = (input_key,) + self.writes = tuple( + prefix + key + for key in ( + "levels", + "base_index", + "current_mean", + "previous_mean", + "boundary_mean", + ) + ) + + def forward(self, batch): + action = _tensor(batch, self.input_key) + if action.ndim != 3 or len(action) == 0: + raise ValueError("target must be a nonempty (B,H,A) sequence") + index = torch.arange(len(action), device=action.device).repeat_interleave( + self.interior_samples_per_content + ) + levels = torch.randint( + 2, self.num_levels + 1, (len(index),), device=action.device + ) + repeated = action.index_select(0, index) + current = self.mean_encoder(repeated, levels.float() / self.num_levels) + previous = self.mean_encoder(repeated, (levels - 1).float() / self.num_levels) + boundary = self.mean_encoder( + action, + torch.full((len(action),), 1.0 / self.num_levels, device=action.device), + ) + if ( + current.shape != previous.shape + or current.ndim != 3 + or boundary.shape[1:] != current.shape[1:] + ): + raise ValueError("reference means must be aligned (B,H,D) sequences") + for key, value in zip( + self.writes, (levels, index, current, previous, boundary), strict=True + ): + batch[key] = value + return batch + + +class GaussianBridgeNoisingStage(Stage): + """Independent Gaussian reparameterizations with exact reverse targets.""" + + train_only = True + + def __init__( + self, + num_levels=32, + sigma_min=0.1, + sigma_max=1.0, + rho=0.95, + condition_dropout_probability=0.3, + condition_key="condition", + prefix="likelihood/", + ): + super().__init__() + self.schedule = GaussianBridgeSchedule(num_levels, sigma_min, sigma_max, rho) + self.num_levels = self.schedule.num_levels + self.sigma_min, self.sigma_max, self.rho = ( + self.schedule.sigma_min, + self.schedule.sigma_max, + self.schedule.rho, + ) + self.condition_dropout_probability = float(condition_dropout_probability) + if not 0 <= self.condition_dropout_probability <= 1: + raise ValueError("condition dropout must lie in [0,1]") + self.condition_key, self.prefix = condition_key, prefix + self.reads = tuple( + prefix + key + for key in ( + "levels", + "base_index", + "current_mean", + "previous_mean", + "boundary_mean", + ) + ) + (condition_key,) + self.writes = tuple( + prefix + key + for key in ( + "state", + "time", + "condition", + "condition_drop_mask", + "posterior_target", + "posterior_variance", + "interior_noise", + "boundary_noise", + ) + ) + + def forward(self, batch): + p = self.prefix + levels, index, current, previous, boundary = ( + _tensor(batch, p + key) + for key in ( + "levels", + "base_index", + "current_mean", + "previous_mean", + "boundary_mean", + ) + ) + condition = _tensor(batch, self.condition_key) + if len(condition) != len(boundary) or condition.device != current.device: + raise ValueError("condition and reference means must align") + if levels.dtype != torch.long or tuple(levels.shape) != (len(current),): + raise ValueError("sampled levels must be an aligned int64 vector") + current_scale = self.schedule.sigmas[levels - 1].float().view(-1, 1, 1) + previous_scale = self.schedule.sigmas[levels - 2].float().view(-1, 1, 1) + noise = torch.randn_like(current, dtype=torch.float32) + boundary_noise = torch.randn_like(boundary, dtype=torch.float32) + current_state = current.float() + current_scale * noise + boundary_state = ( + boundary.float() + self.schedule.sigmas[0].float() * boundary_noise + ) + # This equals mu_prev + rho*sigma_prev/sigma_k*(z_k-mu_k), + # without cancellation. The target's learned mean is NOT detached. + target = previous.float() + self.rho * previous_scale * noise + probability = self.condition_dropout_probability + mask = torch.rand(len(boundary), device=boundary.device) < probability + values = ( + torch.cat((current_state, boundary_state)), + torch.cat( + ( + levels.float() / self.num_levels, + torch.full( + (len(boundary),), 1.0 / self.num_levels, device=boundary.device + ), + ) + ), + torch.cat((condition.index_select(0, index), condition)), + torch.cat((mask.index_select(0, index), mask)), + target, + self.schedule.previous_variance(levels), + noise, + boundary_noise, + ) + for key, value in zip(self.writes, values, strict=True): + batch[key] = value + return batch + + +class ConditionalReverseMeanStage(Stage): + """M(z,k,O)=z+field(z,k/K,O), with the actual stochastic K-level sampler.""" + + def __init__( + self, + field: nn.Module, + num_levels=32, + sigma_min=0.1, + sigma_max=1.0, + rho=0.95, + inference_noise_key="sampler/noise", + inference_condition_key="condition", + prefix="likelihood/", + ): + super().__init__() + if not isinstance(field, nn.Module): + raise TypeError("field must be a module") + self.field = field + self.schedule = GaussianBridgeSchedule(num_levels, sigma_min, sigma_max, rho) + self.num_levels = self.schedule.num_levels + self.sigma_min, self.sigma_max, self.rho = ( + self.schedule.sigma_min, + self.schedule.sigma_max, + self.schedule.rho, + ) + self.prefix = prefix + self.inference_noise_key, self.inference_condition_key = ( + inference_noise_key, + inference_condition_key, + ) + self.reads = tuple( + prefix + key + for key in ( + "state", + "time", + "condition", + "condition_drop_mask", + "posterior_target", + ) + ) + self.writes = (prefix + "interior_prediction", prefix + "boundary_latent") + self.reads_by_mode = { + "inference": (inference_noise_key, inference_condition_key) + } + self.writes_by_mode = { + "inference": ( + prefix + "generated_latent", + prefix + "trajectory", + "log/ActionFlow/SamplerCalls", + "log/ActionFlow/LatentInnovations", + ) + } + + def _mean(self, state, time, condition, mask): + if state.ndim != 3 or len(state) == 0 or time.shape != (len(state),): + raise ValueError( + "reverse mean requires nonempty sequences and aligned times" + ) + if ( + len(condition) != len(state) + or mask.shape != (len(state),) + or mask.dtype != torch.bool + ): + raise ValueError("reverse mean conditioning/mask must align") + delta = self.field(state, time, condition, condition_drop_mask=mask) + if not torch.is_tensor(delta) or delta.shape != state.shape: + raise ValueError("reverse field output must match latent sequence shape") + return state.float() + delta.float() + + def forward(self, batch): + p = self.prefix + state, time, condition, mask = ( + _tensor(batch, p + key) + for key in ( + "state", + "time", + "condition", + "condition_drop_mask", + ) + ) + prediction = self._mean(state, time, condition, mask) + n_interior = len(_tensor(batch, p + "posterior_target")) + batch[p + "interior_prediction"] = prediction[:n_interior] + batch[p + "boundary_latent"] = prediction[n_interior:] + return batch + + @torch.no_grad() + def _inference(self, batch): + state = _tensor(batch, self.inference_noise_key).float() + condition = _tensor(batch, self.inference_condition_key) + mask = torch.zeros(len(state), dtype=torch.bool, device=state.device) + trajectory = [state] + for level in range(self.num_levels, 1, -1): + time = torch.full( + (len(state),), level / self.num_levels, device=state.device + ) + mean = self._mean(state, time, condition, mask) + scale = self.schedule.sigmas[level - 2].float() * math.sqrt( + 1.0 - self.rho**2 + ) + state = mean + scale * torch.randn_like(mean) + trajectory.append(state) + time = torch.full((len(state),), 1.0 / self.num_levels, device=state.device) + state = self._mean(state, time, condition, mask) + trajectory.append(state) + batch[self.prefix + "generated_latent"] = state + batch[self.prefix + "trajectory"] = torch.stack(trajectory) + batch["log/ActionFlow/SamplerCalls"] = float(self.num_levels) + batch["log/ActionFlow/LatentInnovations"] = float(self.num_levels - 1) + return batch + + def execute(self, batch, *, mode): + if mode == "inference": + return self._inference(batch) + return self(batch) + + +class LikelihoodDecoderStage(Stage): + """Private boundary likelihood mean; inference samples its Gaussian output.""" + + def __init__( + self, + decoder: nn.Module, + tau=0.02, + prefix="likelihood/", + prediction_key="pred_action", + ): + super().__init__() + if not isinstance(decoder, nn.Module): + raise TypeError("decoder must be a module") + self.decoder, self.tau, self.prefix, self.prediction_key = ( + decoder, + float(tau), + prefix, + prediction_key, + ) + if not math.isfinite(self.tau) or self.tau <= 0: + raise ValueError("tau must be finite and positive") + self.reads = (prefix + "boundary_latent",) + self.writes = (prefix + "boundary_prediction",) + self.reads_by_mode = {"inference": (prefix + "generated_latent",)} + self.writes_by_mode = {"inference": (prediction_key, prefix + "action_mean")} + + def forward(self, batch): + batch[self.prefix + "boundary_prediction"] = self.decoder( + _tensor(batch, self.prefix + "boundary_latent") + ).float() + return batch + + def execute(self, batch, *, mode): + if mode != "inference": + return self(batch) + mean = self.decoder(_tensor(batch, self.prefix + "generated_latent")).float() + batch[self.prefix + "action_mean"] = mean + batch[self.prediction_key] = mean + self.tau * torch.randn_like(mean) + return batch + + +class GaussianBridgeObjectiveStage(Stage): + """Gaussian negative-ELBO terms, excluding fixed additive constants. + + InteriorBridgeNLL is the equal-covariance KL sum (not ordinary latent FM). + BoundaryNLL is ||A-action_mean||^2/(2*tau^2), not clean reconstruction. + Both sum ALL sequence coordinates. No extra reconstruction/scale/JVP loss. + """ + + train_only = True + reduction = "sum_chunk_coordinates_mean_examples_constant_free_gaussian_bound" + + def __init__( + self, num_levels=32, tau=0.02, target_key="target", prefix="likelihood/" + ): + super().__init__() + self.num_levels, self.tau = int(num_levels), float(tau) + if self.num_levels < 2 or not math.isfinite(self.tau) or self.tau <= 0: + raise ValueError("require num_levels>=2 and finite positive tau") + self.target_key, self.prefix = target_key, prefix + self.reads = tuple( + prefix + key + for key in ( + "interior_prediction", + "posterior_target", + "posterior_variance", + "boundary_prediction", + ) + ) + (target_key,) + self.writes = ( + "loss/likelihood", + "log/ActionFlow/InteriorBridgeNLL", + "log/ActionFlow/BoundaryNLL", + "log/ActionFlow/TotalLoss", + "log/MSE", + ) + + def forward(self, batch): + p = self.prefix + prediction, target, variance, boundary = ( + _tensor(batch, p + key) + for key in ( + "interior_prediction", + "posterior_target", + "posterior_variance", + "boundary_prediction", + ) + ) + action = _tensor(batch, self.target_key) + if prediction.shape != target.shape or prediction.ndim != 3: + raise ValueError( + "interior predictions/targets must align as complete sequences" + ) + if variance.shape != (len(prediction),) or boundary.shape != action.shape: + raise ValueError("likelihood variance or action-boundary shape mismatch") + interior = (self.num_levels - 1) * ( + (prediction.float() - target.float()).square().sum(dim=(-2, -1)) + / (2 * variance.float()) + ).mean() + squared = (boundary.float() - action.float()).square() + boundary_nll = squared.sum(dim=(-2, -1)).mean() / (2 * self.tau**2) + total = interior + boundary_nll + values = (total, interior, boundary_nll, total, squared.mean()) + for key, value in zip(self.writes, values, strict=True): + batch[key] = value + return batch diff --git a/egomimic/pl_utils/pl_model_action_flow.py b/egomimic/pl_utils/pl_model_action_flow.py index 9750277e0..1cbf339a2 100644 --- a/egomimic/pl_utils/pl_model_action_flow.py +++ b/egomimic/pl_utils/pl_model_action_flow.py @@ -436,6 +436,10 @@ def _log_gradient_telemetry(self, components: Mapping[str, torch.Tensor]) -> Non gradients = OrderedDict() routes = OrderedDict() for label, component_name in self._gradient_components: + # An exact-section codec still reports its reconstruction error, + # but that diagnostic is not a trained objective. + if label == "Reconstruction" and self._objective_reconstruction_weight() == 0: + continue active = self._component_gradients(components[component_name], named, label) gradients[label] = active routes[label] = [ @@ -464,9 +468,17 @@ def _log_gradient_telemetry(self, components: Mapping[str, torch.Tensor]) -> Non index for index in left_gradients if index in right_gradients ) if not shared: - raise RuntimeError( - f"Action Flow {left} and {right} have no shared gradient path" - ) + # FM-only endpoint detachment deliberately separates FM from + # clean reconstruction. Do not invent a cosine for that pair. + if {left, right} != {"FM", "Reconstruction"} or not self._fm_endpoint_detached(): + raise RuntimeError( + f"Action Flow {left} and {right} have no shared gradient path" + ) + pair = f"{left}__{right}" + self._log_telemetry(f"GradientCosine/{pair}", 0.0) + self._log_telemetry(f"GradientCosineDefined/{pair}", 0.0) + self._log_telemetry(f"GradientIntersectionParameterCount/{pair}", 0) + continue zero = left_gradients[shared[0]].new_zeros(()) dot = sum( (left_gradients[index] * right_gradients[index]).sum() @@ -533,14 +545,22 @@ def _log_gradient_telemetry(self, components: Mapping[str, torch.Tensor]) -> Non ).hexdigest(), } + def _fm_endpoint_detached(self) -> bool: + stages = getattr(getattr(self.model, "pipeline", None), "stages", ()) + return any( + getattr(stage, "flow_clean_gradient_mode", "full") == "all_stopgrad" + for stage in stages + ) + def _log_compute_contract(self) -> None: if self.flow_samples_per_content is None: return + field_calls = 2 if self._fm_endpoint_detached() else 1 for name, value in ( - ("Compute/FieldForwardCallsPerStep", 1), + ("Compute/FieldForwardCallsPerStep", field_calls), ( "Compute/FieldSampleEquivalentsPerStep", - self.flow_samples_per_content, + field_calls * self.flow_samples_per_content, ), ("Compute/DecoderJVPCallsPerStep", 1), ): diff --git a/egomimic/pl_utils/pl_model_action_flow_likelihood.py b/egomimic/pl_utils/pl_model_action_flow_likelihood.py new file mode 100644 index 000000000..d9e3684f5 --- /dev/null +++ b/egomimic/pl_utils/pl_model_action_flow_likelihood.py @@ -0,0 +1,281 @@ +"""Thin Lightning integration for the discrete Gaussian likelihood graph. + +No FM, clean-reconstruction, Euler, or decoder-JVP diagnostics are borrowed. +The existing evaluator still measures actual generated normalized/native errors +and EnergyScore@32. Train/MSE is the noisy level-one boundary mean's MSE. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping + +import torch + +from egomimic.pl_utils.pl_model import ModelWrapper + + +def _json_hash(value): + return hashlib.sha256( + json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode() + ).hexdigest() + + +class ActionFlowLikelihoodModelWrapper(ModelWrapper): + def __init__( + self, *, action_flow_method=None, gradient_telemetry_cadence=None, **kwargs + ): + super().__init__(**kwargs) + tree = getattr(self.hparams, "config_tree", None) + cfg = {} if tree is None else self._as_config(tree).model + method = action_flow_method or cfg.get( + "action_flow_method", "gaussian_bridge_likelihood" + ) + if method != "gaussian_bridge_likelihood": + raise ValueError("likelihood wrapper requires gaussian_bridge_likelihood") + self.action_flow_method = method + cadence = gradient_telemetry_cadence + if cadence is None: + cadence = cfg.get("gradient_telemetry_cadence", 100) + if isinstance(cadence, bool) or not isinstance(cadence, int) or cadence < 0: + raise ValueError("gradient_telemetry_cadence must be a nonnegative integer") + self.gradient_telemetry_cadence = cadence + self._gradient_route_manifest = None + self.save_hyperparameters( + {"action_flow_method": method, "gradient_telemetry_cadence": cadence} + ) + + def _log_telemetry(self, name, value): + self.log( + f"Train/ActionFlow/{name}", + value, + on_step=True, + on_epoch=False, + sync_dist=False, + ) + + def _capture_gradient_routes(self, components): + named = sorted( + (name, parameter) + for name, parameter in self.nets.named_parameters( + prefix="nets", remove_duplicate=True + ) + if parameter.requires_grad + ) + if not named: + raise RuntimeError("likelihood graph has no trainable parameters") + routes, vectors = {}, {} + for label, loss in components.items(): + gradients = torch.autograd.grad( + loss, tuple(p for _, p in named), retain_graph=True, allow_unused=True + ) + active = { + i: g.detach().float() for i, g in enumerate(gradients) if g is not None + } + if not active: + raise RuntimeError( + f"likelihood {label} has no generative gradient route" + ) + norm = torch.stack([g.square().sum() for g in active.values()]).sum().sqrt() + if not bool(torch.isfinite(norm)): + raise RuntimeError(f"non-finite likelihood {label} gradient") + self._log_telemetry(f"GradientNorm/{label}", norm) + self._log_telemetry( + f"GradientParameterCount/{label}", + sum(named[i][1].numel() for i in active), + ) + routes[label] = [ + { + "name": named[i][0], + "dtype": str(named[i][1].dtype), + "numel": named[i][1].numel(), + "shape": list(named[i][1].shape), + } + for i in active + ] + vectors[label] = active + left, right = "InteriorBridgeNLL", "BoundaryNLL" + indices = sorted(set(vectors[left]) & set(vectors[right])) + if not indices: + raise RuntimeError( + "interior and boundary likelihood must share a gradient pathway" + ) + dot = torch.stack( + [(vectors[left][i] * vectors[right][i]).sum() for i in indices] + ).sum() + norms = [ + torch.stack([vectors[label][i].square().sum() for i in indices]) + .sum() + .sqrt() + for label in (left, right) + ] + denominator = norms[0] * norms[1] + defined = bool(denominator > 0) + pair = left + "__" + right + self._log_telemetry( + f"GradientCosine/{pair}", + (dot / denominator).clamp(-1, 1) if defined else dot.new_zeros(()), + ) + self._log_telemetry(f"GradientCosineDefined/{pair}", float(defined)) + self._log_telemetry( + f"GradientIntersectionParameterCount/{pair}", + sum(named[i][1].numel() for i in indices), + ) + core = { + "schema_version": 1, + "routes": routes, + "route_sha256": { + label: _json_hash(route) for label, route in routes.items() + }, + "intersections": {pair: [named[i][0] for i in indices]}, + } + self._gradient_route_manifest = {**core, "manifest_sha256": _json_hash(core)} + + def _log_prediction_metrics(self, predictions, reference): + # Aggregate exactly as the generic pipeline objective: equal sources; + # each source already averages its examples and sampled interior levels. + for metric, source_values in self._prediction_log_metrics( + predictions, reference + ).items(): + for source, value in source_values: + self.log( + f"Train/{metric}/{source}", + value, + sync_dist=True, + on_step=True, + on_epoch=True, + ) + self.log( + f"Train/{metric}", + torch.stack([x for _, x in source_values]).mean(), + sync_dist=True, + on_step=True, + on_epoch=True, + ) + self._log_telemetry("OptimizerStep", float(self.global_step)) + stages = self.model.pipeline.stages + samples = next( + s.interior_samples_per_content + for s in stages + if hasattr(s, "interior_samples_per_content") + ) + self._log_telemetry("Compute/FieldForwardCallsPerStep", 1) + self._log_telemetry("Compute/FieldSampleEquivalentsPerStep", samples + 1) + self._log_telemetry("Compute/DecoderJVPCallsPerStep", 0) + if ( + self.gradient_telemetry_cadence + and (int(self.global_step) + 1) % self.gradient_telemetry_cadence == 0 + ): + self._capture_gradient_routes( + { + label: torch.stack( + [ + result[f"log/ActionFlow/{label}"] + for result in predictions.values() + ] + ).mean() + for label in ("InteriorBridgeNLL", "BoundaryNLL") + } + ) + + def on_after_backward(self): + self._log_telemetry( + "Compute/PeakAllocatedBytes", + ( + torch.cuda.max_memory_allocated(self.device) + if self.device.type == "cuda" + else 0 + ), + ) + super().on_after_backward() + + def on_before_optimizer_step(self, optimizer): + if ( + self.gradient_telemetry_cadence + and (int(self.global_step) + 1) % self.gradient_telemetry_cadence == 0 + ): + pieces = [ + p.grad.detach().float().square().sum() + for p in self.parameters() + if p.grad is not None + ] + if pieces: + self._log_telemetry( + "GradientNorm/TotalPreclip", torch.stack(pieces).sum().sqrt() + ) + super().on_before_optimizer_step(optimizer) + + def validation_step(self, batch, batch_idx, dataloader_idx=0): + """Measure the held-out graph objective, then run the normal evaluator. + + This retains the configured noising/condition-dropout objective, with + modules in evaluation mode and no gradients. It is not clean recon or + FM. Isolated draws cannot perturb the evaluator's actual sample bank. + """ + if isinstance(batch, Mapping): + batch = {key: value for key, value in batch.items() if value is not None} + if not batch: + return + processed = self.model.process_batch_for_training(batch) + devices = [self.device.index] if self.device.type == "cuda" else [] + seed = 420_042 + int(batch_idx) + int(dataloader_idx) * 100_003 + seed += int(self.global_rank) * 1_000_003 + with torch.random.fork_rng(devices=devices), torch.no_grad(): + torch.manual_seed(seed) + predictions = self.model.forward_training(processed) + reference = next(iter(predictions.values()))["loss/likelihood"] + metrics = self._prediction_log_metrics(predictions, reference) + count = sum(len(result["target"]) for result in predictions.values()) + for label in ("InteriorBridgeNLL", "BoundaryNLL", "TotalLoss"): + values = metrics[f"ActionFlow/{label}"] + if len(values) != len(predictions): + raise RuntimeError(f"missing held-out likelihood component {label}") + # As in training, average the per-source objective means equally. + self.log( + f"Valid/ActionFlow/{label}", + torch.stack([value for _, value in values]).mean(), + on_step=False, + on_epoch=True, + batch_size=count, + sync_dist=True, + add_dataloader_idx=False, + ) + if self.evaluator is not None: + self.evaluator.on_validation_step(processed, batch_idx, dataloader_idx) + + def on_save_checkpoint(self, checkpoint): + if self._gradient_route_manifest is not None: + checkpoint["action_flow_gradient_route_manifest"] = ( + self._gradient_route_manifest + ) + stages = self.model.pipeline.stages + reference = next( + s for s in stages if hasattr(s, "interior_samples_per_content") + ) + noising = next( + s + for s in stages + if hasattr(s, "schedule") and hasattr(s, "condition_dropout_probability") + ) + objective = next(s for s in stages if getattr(s, "reduction", None)) + checkpoint["action_flow_likelihood_contract"] = { + "schema_version": 1, + "method": self.action_flow_method, + "num_levels": reference.num_levels, + "interior_samples_per_content": reference.interior_samples_per_content, + "sigma_min": noising.sigma_min, + "sigma_max": noising.sigma_max, + "rho": noising.rho, + "tau": objective.tau, + "reduction": objective.reduction, + "learned_reference_targets": "attached", + "sampler": "stochastic_reverse_gaussian_chain_with_action_output_noise", + "clean_reconstruction_objective": False, + "latent_fm_objective": False, + } diff --git a/scripts/ice/ice_requeue_runner.py b/scripts/ice/ice_requeue_runner.py index f28e080ff..41fdd9986 100644 --- a/scripts/ice/ice_requeue_runner.py +++ b/scripts/ice/ice_requeue_runner.py @@ -900,6 +900,12 @@ def handler(received: int, _frame: object) -> None: # A parent shell must not be able to leave the runner deaf to Slurm's # boundary/termination contract through an inherited signal mask. signal.pthread_sigmask(signal.SIG_UNBLOCK, handled_signals) + signal_ready_file = os.environ.pop("ICE_RUNNER_SIGNAL_READY_FILE", None) + if signal_ready_file: + ready_path = Path(signal_ready_file) + if not ready_path.is_absolute(): + raise SystemExit("ICE_RUNNER_SIGNAL_READY_FILE must be absolute") + atomic_json_once(ready_path, {"pid": os.getpid(), "signal_handlers_ready": True}) if not args.command: raise SystemExit("a child command is required after --") diff --git a/scripts/ice/relay_batch_signals.sh b/scripts/ice/relay_batch_signals.sh new file mode 100644 index 000000000..e8568d1cc --- /dev/null +++ b/scripts/ice/relay_batch_signals.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# Source this helper when post-run verification prevents exec-ing the runner. +# The runner publishes ICE_RUNNER_SIGNAL_READY_FILE after installing handlers. +ice_run_with_signal_relay() { + local ready_file=$1 + shift + test ! -e "$ready_file" || return 73 + ICE_BATCH_BOUNDARY_FORWARDED=0 + local relay_pid= relay_usr1=0 relay_cancel= relay_interrupted=0 relay_status=0 + _ice_relay_pending() { + test -n "$relay_pid" || return 0 + test -f "$ready_file" || return 0 + if test -n "$relay_cancel"; then + kill -s "$relay_cancel" "$relay_pid" 2>/dev/null || true + relay_cancel= + relay_usr1=0 + elif test "$relay_usr1" = 1; then + if kill -USR1 "$relay_pid" 2>/dev/null; then + ICE_BATCH_BOUNDARY_FORWARDED=1 + fi + relay_usr1=0 + fi + } + trap 'relay_interrupted=1; relay_usr1=1; _ice_relay_pending' USR1 + trap 'relay_interrupted=1; relay_cancel=TERM; _ice_relay_pending' TERM + trap 'relay_interrupted=1; relay_cancel=INT; _ice_relay_pending' INT + ICE_RUNNER_SIGNAL_READY_FILE="$ready_file" "$@" & + relay_pid=$! + # Latch an early boundary instead of killing Python before its handlers exist. + while test ! -f "$ready_file" && kill -0 "$relay_pid" 2>/dev/null; do + sleep 0.02 + done + _ice_relay_pending + while :; do + relay_interrupted=0 + if wait "$relay_pid"; then relay_status=0; else relay_status=$?; fi + # A trapped signal interrupts bash wait, not necessarily the child. A + # repeated wait also retains an already-exited child's actual status. + test "$relay_interrupted" = 1 || break + done + trap - USR1 TERM INT + unset -f _ice_relay_pending + return "$relay_status" +} diff --git a/scripts/train/launch_action_flow_usocket.sbatch b/scripts/train/launch_action_flow_usocket.sbatch index beb9d4907..86dff099d 100644 --- a/scripts/train/launch_action_flow_usocket.sbatch +++ b/scripts/train/launch_action_flow_usocket.sbatch @@ -41,6 +41,16 @@ sha256() { sha256sum "$1" | awk '{print $1}' } +partition_allowed() { + local requested=$1 selected=$2 + [[ "$requested" =~ ^[A-Za-z0-9_-]+(,[A-Za-z0-9_-]+)*$ ]] || return 1 + [[ "$selected" =~ ^[A-Za-z0-9_-]+$ ]] || return 1 + case ",$requested," in + *",$selected,"*) return 0 ;; + *) return 1 ;; + esac +} + for variable in \ AF_REPO AF_EXPECTED_HEAD AF_PYTHON AF_DATASET_DIR AF_OUTPUT_DIR \ AF_EXPERIMENT AF_WANDB_ENTITY AF_WANDB_PROJECT AF_WANDB_RUN_ID \ @@ -88,7 +98,30 @@ if test "$AF_LAUNCH_MODE" = run; then required AF_EXPECTED_MEMORY required AF_EXPECTED_TIME_LIMIT fi +AF_METHOD=action_flow_joint case "$AF_EXPERIMENT" in + pusht/action_flow_bc_usocket_latent_fm_sg_recon1_s42) + AF_METHOD=latent_fm_stopgrad + AF_RECONSTRUCTION_WEIGHT=1 + AF_FLOW_WEIGHT=1 + AF_EXPECTED_CONFIG_NAME=action_flow_bc_usocket_latent_fm_sg_recon1_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. + AF_RECONSTRUCTION_WEIGHT=0 + AF_FLOW_WEIGHT=0 + AF_EXPECTED_CONFIG_NAME=action_flow_bc_usocket_bridge_likelihood_s42 + AF_FULL_RECONSTRUCTION_ONLY_WARMUP_STEPS=0 + ;; + pusht/action_flow_bc_usocket_graph_section_s42) + AF_METHOD=graph_section_diagnostic + AF_RECONSTRUCTION_WEIGHT=0 + AF_FLOW_WEIGHT=1 + AF_EXPECTED_CONFIG_NAME=action_flow_bc_usocket_graph_section_s42 + AF_FULL_RECONSTRUCTION_ONLY_WARMUP_STEPS=0 + ;; pusht/action_flow_bc_usocket_recon1_s42) AF_RECONSTRUCTION_WEIGHT=1 AF_FLOW_WEIGHT=1 @@ -119,10 +152,10 @@ case "$AF_EXPERIMENT" in AF_EXPECTED_CONFIG_NAME=action_flow_bc_usocket_recon100_s42 AF_FULL_RECONSTRUCTION_ONLY_WARMUP_STEPS=0 ;; - *) die "AF_EXPERIMENT must select an approved recon-1, recon-10, or recon-100 config" ;; + *) die "AF_EXPERIMENT must select an explicitly approved Action Flow config" ;; esac -if [[ "$AF_EXPERIMENT" = *warmup10k* ]]; then +if test "$AF_METHOD" != action_flow_joint || [[ "$AF_EXPERIMENT" = *warmup10k* ]]; then AF_COMPARISON_EXPERIMENT= elif test "$AF_EXPERIMENT" = pusht/action_flow_bc_usocket_recon1_s42; then AF_COMPARISON_EXPERIMENT=pusht/action_flow_bc_usocket_recon10_s42 @@ -287,9 +320,6 @@ BASE_OVERRIDES=( trainer.num_sanity_val_steps=0 "trainer.log_every_n_steps=$LOG_EVERY" model.gradient_telemetry_cadence="$TELEMETRY_EVERY" - ++model.reconstruction_only_warmup_steps="$RECONSTRUCTION_ONLY_WARMUP_STEPS" - "++run_provenance.objective.requested_full_reconstruction_only_warmup_steps=$AF_FULL_RECONSTRUCTION_ONLY_WARMUP_STEPS" - "++run_provenance.objective.effective_reconstruction_only_warmup_steps=$RECONSTRUCTION_ONLY_WARMUP_STEPS" ++callbacks.model_checkpoint.monitor=null callbacks.model_checkpoint.save_top_k=-1 callbacks.model_checkpoint.save_last=link @@ -325,6 +355,14 @@ BASE_OVERRIDES=( "++run_provenance.dataset_content_aggregate_sha256=$AF_EXPECTED_DATASET_CONTENT_AGGREGATE_SHA256" ) +if test "$AF_METHOD" != gaussian_bridge_likelihood; then + BASE_OVERRIDES+=( + ++model.reconstruction_only_warmup_steps="$RECONSTRUCTION_ONLY_WARMUP_STEPS" + "++run_provenance.objective.requested_full_reconstruction_only_warmup_steps=$AF_FULL_RECONSTRUCTION_ONLY_WARMUP_STEPS" + "++run_provenance.objective.effective_reconstruction_only_warmup_steps=$RECONSTRUCTION_ONLY_WARMUP_STEPS" + ) +fi + if test -n "${AF_NORM_STATS_PATH:-}"; then absolute_path AF_NORM_STATS_PATH "$AF_NORM_STATS_PATH" if test -d "$AF_NORM_STATS_PATH"; then @@ -355,7 +393,7 @@ if test "$AF_LAUNCH_MODE" = preflight; then required "$variable" done test "$SLURM_JOB_ACCOUNT" = "$AF_EXPECTED_ACCOUNT" || die "unexpected Slurm account" - test "$SLURM_JOB_PARTITION" = "$AF_EXPECTED_PARTITION" || die "unexpected Slurm partition" + partition_allowed "$AF_EXPECTED_PARTITION" "$SLURM_JOB_PARTITION" || die "unexpected Slurm partition" test "$SLURM_JOB_QOS" = "$AF_EXPECTED_QOS" || die "unexpected Slurm QoS" test "$SLURM_JOB_NUM_NODES" = 1 || die "preflight requires one node" test "$SLURM_NTASKS" = 1 || die "preflight requires one task" @@ -631,7 +669,7 @@ for variable in \ required "$variable" done test "$SLURM_JOB_ACCOUNT" = "$AF_EXPECTED_ACCOUNT" || die "unexpected Slurm account" -test "$SLURM_JOB_PARTITION" = "$AF_EXPECTED_PARTITION" || die "unexpected Slurm partition" +partition_allowed "$AF_EXPECTED_PARTITION" "$SLURM_JOB_PARTITION" || die "unexpected Slurm partition" test "$SLURM_JOB_QOS" = "$AF_EXPECTED_QOS" || die "unexpected Slurm QoS" test "$SLURM_JOB_NUM_NODES" = 1 || die "launcher requires one node" test "$SLURM_NTASKS" = 1 || die "launcher requires one task" @@ -696,7 +734,7 @@ scontrol show job -dd -o "$SLURM_JOB_ID" > "$ATTEMPT/slurm_job.txt" --record "$ATTEMPT/slurm_job.txt" \ --expected-job-id "$SLURM_JOB_ID" \ --expected-account "$AF_EXPECTED_ACCOUNT" \ - --expected-partition "$AF_EXPECTED_PARTITION" \ + --expected-partition "$SLURM_JOB_PARTITION" \ --expected-qos "$AF_EXPECTED_QOS" \ --expected-cpus 8 \ --expected-memory "$AF_EXPECTED_MEMORY" \ @@ -779,42 +817,24 @@ if run_kind == "full": assert cfg.run_provenance.smoke_result_sha256 == smoke_sha else: assert smoke_sha == "" -assert cfg.model._target_.endswith("ActionFlowModelWrapper") +from tools.validate_action_flow_config import ( + LEGACY_METHOD, LIKELIHOOD_METHOD, GRAPH_METHOD, STOPGRAD_METHOD, + validate_method_contract, _validate_dimensions_and_modules, _validate_topology, +) +method = validate_method_contract(cfg, experiment) assert cfg.model.action_horizon == 16 and cfg.model.action_dim == 4 assert cfg.model.latent_dim == 8 and cfg.model.condition_dim == 67 -assert cfg.model.flow_samples_per_content == 14 assert math.isclose(cfg.model.condition_dropout_probability, 0.3) -assert cfg.model.num_inference_steps == 16 -assert math.isclose(cfg.model.reconstruction_weight, reconstruction) -assert math.isclose(cfg.model.flow_weight, flow_weight) -assert cfg.model.reconstruction_only_warmup_steps == effective_warmup_steps -assert ( - cfg.run_provenance.objective.effective_reconstruction_only_warmup_steps - == effective_warmup_steps -) -assert ( - cfg.run_provenance.objective.requested_full_reconstruction_only_warmup_steps - == full_warmup_steps -) -stages = cfg.model.pipeline.stages -assert [str(stage._target_).rsplit(".", 1)[-1] for stage in stages] == [ - "FusedObsEncoder", - "GaussianLatentNoise", - "ActionTargetBuilder", - "ContentEncoderStage", - "LatentBridgeStage", - "ConditionalVelocityStage", - "ContentDecoderStage", - "ActionFlowObjectiveStage", -] -assert stages[1].num_tokens == 16 and stages[1].latent_dim == 8 -assert stages[4].samples_per_content == 14 -assert math.isclose(stages[4].condition_dropout_probability, 0.3) -assert stages[5].num_inference_steps == 16 -assert stages[5].field.time_scale == 1000.0 -assert math.isclose(stages[7].flow_weight, flow_weight) -assert stages[7].action_velocity_weight == 1.0 -assert math.isclose(stages[7].reconstruction_weight, reconstruction) +if method != LIKELIHOOD_METHOD: + assert cfg.model.flow_samples_per_content == 14 + assert cfg.model.num_inference_steps == 16 + assert math.isclose(cfg.model.reconstruction_weight, reconstruction) + assert math.isclose(cfg.model.flow_weight, flow_weight) + assert cfg.model.reconstruction_only_warmup_steps == effective_warmup_steps + assert cfg.run_provenance.objective.effective_reconstruction_only_warmup_steps == effective_warmup_steps + assert cfg.run_provenance.objective.requested_full_reconstruction_only_warmup_steps == full_warmup_steps +else: + assert effective_warmup_steps == full_warmup_steps == 0 optimizer, scheduler = cfg.model.optimizer, cfg.model.scheduler assert optimizer._target_ == "torch.optim.AdamW" and optimizer._partial_ is True assert optimizer.lr == 3e-5 and list(optimizer.betas) == [0.9, 0.999] @@ -850,15 +870,20 @@ assert cfg.data.valid_dataloader_params[source].batch_size == 16 assert str(cfg.norm_stats.precomputed_norm_path) == norm assert cfg.norm_stats.save_cache_dir is None diagnostics = cfg.evaluator.action_flow_diagnostics -assert diagnostics.enabled is True and diagnostics.max_batches_per_rank == 1 -assert list(diagnostics.raw_noise_levels) == [0.0, 0.25, 0.5, 0.75, 1.0] -assert diagnostics.max_samples == 16 and diagnostics.jacobian_samples == 2 -assert diagnostics.capture_activations is True and diagnostics.cknna_k == 10 -assert diagnostics.validation_view.world_size == 1 -assert diagnostics.validation_view.per_rank_batch_size == 16 +if method != LIKELIHOOD_METHOD: + assert diagnostics.enabled is True and diagnostics.max_batches_per_rank == 1 + assert list(diagnostics.raw_noise_levels) == [0.0, 0.25, 0.5, 0.75, 1.0] + assert diagnostics.max_samples == 16 and diagnostics.jacobian_samples == 2 + assert diagnostics.capture_activations is (method != GRAPH_METHOD) + assert diagnostics.cknna_k == (0 if method == GRAPH_METHOD else 10) + assert diagnostics.validation_view.world_size == 1 + assert diagnostics.validation_view.per_rank_batch_size == 16 pipeline = instantiate(cfg.model.pipeline, device="cpu") +_validate_topology(cfg, pipeline) +_validate_dimensions_and_modules(cfg, tuple(pipeline.pipeline.stages)) count = sum(parameter.numel() for parameter in pipeline.nets.parameters()) -assert count == 50_725_221, count +if method in (LEGACY_METHOD, STOPGRAD_METHOD): + assert count == 50_725_221, count print(f"[run-preflight] PASS kind={run_kind} parameters={count}") PY sha256sum "$ATTEMPT/resolved_config.yaml" > "$ATTEMPT/resolved_config.sha256" @@ -936,9 +961,20 @@ RUNNER_ARGS=( --confirm-child-requeue-disabled --completion-sentinel "$AF_OUTPUT_DIR/COMPLETE.json" ) -"$AF_PYTHON" "$AF_REQUEUE_RUNNER" "${RUNNER_ARGS[@]}" -- \ +# Slurm B:USR1 reaches this batch shell only. Preserve the post-run verifier +# while relaying boundary/cancellation signals to the runner's installed handlers. +source "$AF_REPO/scripts/ice/relay_batch_signals.sh" +ice_run_with_signal_relay "$ATTEMPT/runner-signal-ready.json" \ + "$AF_PYTHON" "$AF_REQUEUE_RUNNER" "${RUNNER_ARGS[@]}" -- \ /usr/bin/env bash -c "$CHILD_CODE" action-flow-usocket "${BASE_OVERRIDES[@]}" +# A successful requeue request also returns zero, but is not training completion. +# The runner writes COMPLETE.json only for successful child completion. +if test "$ICE_BATCH_BOUNDARY_FORWARDED" = 1 && test ! -s "$AF_OUTPUT_DIR/COMPLETE.json"; then + printf '[requeue] boundary handled; completion verification belongs to the final attempt\n' + exit 0 +fi + if test "$AF_RUN_KIND" = smoke; then SMOKE_CHECKPOINT=$AF_OUTPUT_DIR/checkpoints/last.ckpt test -s "$SMOKE_CHECKPOINT" || die "smoke checkpoint is unavailable" @@ -952,11 +988,14 @@ if test "$AF_RUN_KIND" = smoke; then --measured-checkpoint-bytes "$SMOKE_CHECKPOINT_BYTES" \ --safety-reserve-bytes "$AF_STORAGE_SAFETY_RESERVE_BYTES" \ --output "$ATTEMPT/CHECKPOINT_STORAGE.json" + SMOKE_WEIGHT_ARGS=() + if test "$AF_METHOD" != gaussian_bridge_likelihood; then + SMOKE_WEIGHT_ARGS=(--expected-reconstruction-weight "$AF_RECONSTRUCTION_WEIGHT" --expected-flow-weight "$AF_FLOW_WEIGHT") + fi "$AF_PYTHON" "$AF_SMOKE_VERIFIER" "$AF_OUTPUT_DIR" \ --expected-head "$AF_EXPECTED_HEAD" \ --expected-experiment "$AF_EXPERIMENT" \ - --expected-reconstruction-weight "$AF_RECONSTRUCTION_WEIGHT" \ - --expected-flow-weight "$AF_FLOW_WEIGHT" \ + "${SMOKE_WEIGHT_ARGS[@]}" \ --expected-split-sha256 "$AF_EXPECTED_SPLIT_MANIFEST_SHA256" \ --expected-normalization-sha256 "$AF_EXPECTED_NORM_SHA256" \ --expected-content-manifest-sha256 \ diff --git a/scripts/train/verify_action_flow_training_smoke.py b/scripts/train/verify_action_flow_training_smoke.py index 3a924f121..41b53460d 100644 --- a/scripts/train/verify_action_flow_training_smoke.py +++ b/scripts/train/verify_action_flow_training_smoke.py @@ -45,10 +45,38 @@ from egomimic.pl_utils.pl_model_action_flow import ( # noqa: E402 ActionFlowModelWrapper, ) +from tools.validate_action_flow_config import ( # noqa: E402 + CANDIDATE_METHODS, + LEGACY_METHOD, + LIKELIHOOD_METHOD, + GRAPH_METHOD, + STOPGRAD_METHOD, + PreflightError, + action_flow_method, + method_stage_targets, + method_wrapper_target, + validate_method_contract, + _validate_dimensions_and_modules, +) SCHEMA_VERSION = 1 EXPECTED_PARAMETER_COUNT = 50_725_221 APPROVED_EXPERIMENTS = { + "pusht/action_flow_bc_usocket_latent_fm_sg_recon1_s42": ( + "action_flow_bc_usocket_latent_fm_sg_recon1_s42", + 1.0, + 1.0, + ), + "pusht/action_flow_bc_usocket_bridge_likelihood_s42": ( + "action_flow_bc_usocket_bridge_likelihood_s42", + 0.0, + 0.0, + ), + "pusht/action_flow_bc_usocket_graph_section_s42": ( + "action_flow_bc_usocket_graph_section_s42", + 0.0, + 1.0, + ), "pusht/action_flow_bc_usocket_recon1_s42": ( "action_flow_bc_usocket_recon1_s42", 1.0, @@ -258,15 +286,21 @@ def _validate_config( _require(experiment in APPROVED_EXPERIMENTS, f"unapproved experiment: {experiment}") expected_name, reconstruction_weight, flow_weight = APPROVED_EXPERIMENTS[experiment] config = OmegaConf.load(config_path) + try: + method = validate_method_contract(config, experiment) + except PreflightError as error: + raise SmokeVerificationError(str(error)) from error _exact(config, "name", expected_name) _exact( config, "model._target_", - "egomimic.pl_utils.pl_model_action_flow.ActionFlowModelWrapper", + method_wrapper_target(method), ) targets = tuple(str(stage._target_) for stage in config.model.pipeline.stages) - _require(targets == EXPECTED_STAGE_TARGETS, f"unexpected stage topology: {targets}") + _require( + targets == method_stage_targets(method), f"unexpected stage topology: {targets}" + ) for path, expected in ( ("model.action_horizon", 16), @@ -320,6 +354,15 @@ def _validate_config( ("evaluator.energy_score_validation_view.world_size", 1), ("evaluator.energy_score_validation_view.per_rank_batch_size", 16), ): + if method == LIKELIHOOD_METHOD and path in { + "model.flow_samples_per_content", + "model.num_inference_steps", + "model.pipeline.stages.4.samples_per_content", + "model.pipeline.stages.5.num_inference_steps", + "run_provenance.objective.flow_samples_per_content", + "run_provenance.inference.steps", + }: + continue # Validated against the discrete-chain contract above. _exact(config, path, expected) if "_warmup10k_" in experiment or experiment.endswith("_warmup10k_s42"): @@ -359,6 +402,12 @@ def _validate_config( ("run_provenance.objective.decoded_noise_scale_weight", 0.0), ("run_provenance.objective.monotonic_weight", 0.0), ): + if method == LIKELIHOOD_METHOD and ( + path.startswith("model.pipeline.stages.7.") + or path in {"model.reconstruction_weight", "model.flow_weight"} + or path.startswith("run_provenance.objective.") + ): + continue # This method has NLL components, not FM/reconstruction. _float(config, path, expected) _exact(config, "mode", "train") @@ -387,7 +436,15 @@ def _validate_config( ) _exact(config, "logger.wandb.offline", False) _exact(config, "evaluator.energy_score_enabled", True) - _exact(config, "run_provenance.inference.sampler", "reverse_euler") + _exact( + config, + "run_provenance.inference.sampler", + ( + "gaussian_bridge_reverse_chain" + if method == LIKELIHOOD_METHOD + else "reverse_euler" + ), + ) _exact(config, "run_provenance.inference.classifier_free_guidance", False) _exact(config, "run_provenance.action_contract.prediction_horizon", 16) _exact( @@ -419,17 +476,25 @@ def _validate_config( normalized_distance == USOCKET_ENERGY_DISTANCE_CONFIG, "typed USocket EnergyScore distance contract differs", ) - try: - native_error_contract = normalize_usocket_native_error_config( - _plain_mapping( - _select(config, "evaluator.action_flow_diagnostics.native_error"), - label="Action Flow native-error contract", - ) + if method == LIKELIHOOD_METHOD: + _exact( + config, + "evaluator.native_decoder._target_", + "egomimic.pipeline.pushshapes.USocketRotVecNativeDecoder", ) - except (TypeError, ValueError) as error: - raise SmokeVerificationError( - f"invalid Action Flow native-error contract: {error}" - ) from error + native_error_contract = dict(USOCKET_NATIVE_ERROR_CONFIG) + else: + try: + native_error_contract = normalize_usocket_native_error_config( + _plain_mapping( + _select(config, "evaluator.action_flow_diagnostics.native_error"), + label="Action Flow native-error contract", + ) + ) + except (TypeError, ValueError) as error: + raise SmokeVerificationError( + f"invalid Action Flow native-error contract: {error}" + ) from error _require( native_error_contract == USOCKET_NATIVE_ERROR_CONFIG, "Action Flow native-error contract differs", @@ -607,22 +672,30 @@ def _validate_config( "evaluator.energy_score_provenance", "evaluator.action_flow_diagnostics.provenance", ): + if ( + method == LIKELIHOOD_METHOD + and evaluator_prefix == "evaluator.action_flow_diagnostics.provenance" + ): + continue _exact(config, f"{evaluator_prefix}.source_commit", expected_head) _exact(config, f"{evaluator_prefix}.normalization_sha256", normalization_hash) _exact(config, f"{evaluator_prefix}.split_manifest_sha256", split_hash) - _exact(config, "evaluator.action_flow_diagnostics.enabled", True) - _exact(config, "evaluator.action_flow_diagnostics.max_batches_per_rank", 1) - _exact( - config, - "evaluator.action_flow_diagnostics.validation_view.world_size", - 1, - ) - diagnostic_split = _select( - config, - "evaluator.action_flow_diagnostics.validation_view.split_manifest_sha256", - ) - _require(str(diagnostic_split) == split_hash, "diagnostic split identity mismatch") + if method != LIKELIHOOD_METHOD: + _exact(config, "evaluator.action_flow_diagnostics.enabled", True) + _exact(config, "evaluator.action_flow_diagnostics.max_batches_per_rank", 1) + _exact( + config, + "evaluator.action_flow_diagnostics.validation_view.world_size", + 1, + ) + diagnostic_split = _select( + config, + "evaluator.action_flow_diagnostics.validation_view.split_manifest_sha256", + ) + _require( + str(diagnostic_split) == split_hash, "diagnostic split identity mismatch" + ) for path, label in ( ("evaluator.artifact_root", "EnergyScore artifact root"), @@ -631,6 +704,10 @@ def _validate_config( "Action Flow diagnostic artifact root", ), ): + if method == LIKELIHOOD_METHOD and path.startswith( + "evaluator.action_flow_diagnostics." + ): + continue root = Path(str(_select(config, path))).expanduser() if not root.is_absolute(): root = run_dir / root @@ -702,6 +779,7 @@ def _canonical_json_sha256(value: Any) -> str: def _validate_gradient_route_manifest( manifest: Any, named_parameters: Sequence[tuple[str, torch.nn.Parameter]], + method: str = LEGACY_METHOD, ) -> dict[str, Any]: _require(isinstance(manifest, Mapping), "gradient route manifest is missing") _require( @@ -719,7 +797,15 @@ def _validate_gradient_route_manifest( routes = manifest.get("routes") route_hashes = manifest.get("route_sha256") intersections = manifest.get("intersections") - expected_labels = ("FM", "Reconstruction", "ActionVelocity") + expected_labels = ( + ("InteriorBridgeNLL", "BoundaryNLL") + if method == LIKELIHOOD_METHOD + else ( + ("FM", "ActionVelocity") + if method == GRAPH_METHOD + else ("FM", "Reconstruction", "ActionVelocity") + ) + ) _require( isinstance(routes, Mapping) and tuple(routes) == expected_labels, "gradient route labels differ", @@ -760,23 +846,22 @@ def _validate_gradient_route_manifest( route_names[label] = names expected_intersections = {} - for left, right in ( - ("FM", "Reconstruction"), - ("FM", "ActionVelocity"), - ("Reconstruction", "ActionVelocity"), - ): - right_names = set(route_names[right]) - expected_intersections[f"{left}__{right}"] = [ - name for name in route_names[left] if name in right_names - ] + for index, left in enumerate(expected_labels): + for right in expected_labels[index + 1 :]: + right_names = set(route_names[right]) + expected_intersections[f"{left}__{right}"] = [ + name for name in route_names[left] if name in right_names + ] _require( intersections == expected_intersections, "gradient route intersections do not match route entries", ) - _require( - all(expected_intersections.values()), - "one or more required shared gradient pathways are empty", - ) + for pair, names in expected_intersections.items(): + expected_empty = method == STOPGRAD_METHOD and pair == "FM__Reconstruction" + _require( + bool(names) is not expected_empty, + f"unexpected shared gradient pathway: {pair}", + ) stage_prefixes = { "observation": "nets.pipeline.stages.0.", @@ -789,6 +874,15 @@ def _validate_gradient_route_manifest( "Reconstruction": ("encoder", "decoder"), "ActionVelocity": ("observation", "encoder", "field", "decoder"), } + if method == STOPGRAD_METHOD: + expected_reachability["FM"] = ("observation", "field") + elif method == GRAPH_METHOD: + del expected_reachability["Reconstruction"] + elif method == LIKELIHOOD_METHOD: + expected_reachability = { + "InteriorBridgeNLL": ("observation", "encoder", "field"), + "BoundaryNLL": ("observation", "encoder", "field", "decoder"), + } for label, active_groups in expected_reachability.items(): for group, prefix in stage_prefixes.items(): observed = any(name.startswith(prefix) for name in route_names[label]) @@ -817,8 +911,45 @@ def _validate_gradient_route_manifest( } +def _validate_checkpoint_loss_schedule( + loss_schedule: Any, + config: DictConfig | None, + *, + reconstruction_weight: float, + flow_weight: float, +) -> int: + _require(config is not None, "checkpoint loss schedule needs its exact config") + warmup_steps = OmegaConf.select( + config, "model.reconstruction_only_warmup_steps", default=0 + ) + _require( + isinstance(warmup_steps, int) + and not isinstance(warmup_steps, bool) + and warmup_steps in (0, 1), + "two-update smoke requires a configured reconstruction warmup of 0 or 1", + ) + _require( + loss_schedule + == { + "joint_objective_begins_at_global_step": warmup_steps, + "reconstruction_only_optimizer_steps": warmup_steps, + "joint_flow_weight": flow_weight, + "joint_reconstruction_weight": reconstruction_weight, + "joint_action_velocity_weight": 1.0, + "schema_version": 1, + }, + f"unexpected Action Flow loss schedule: {loss_schedule}", + ) + return warmup_steps + + def _validate_checkpoint( - run_dir: Path, *, reconstruction_weight: float, flow_weight: float + run_dir: Path, + *, + reconstruction_weight: float, + flow_weight: float, + method: str = LEGACY_METHOD, + config: DictConfig | None = None, ) -> dict[str, Any]: checkpoint_dir = run_dir / "checkpoints" last_path = checkpoint_dir / "last.ckpt" @@ -843,6 +974,7 @@ def _validate_checkpoint( loops = payload.get("loops") gradient_route_manifest = payload.get("action_flow_gradient_route_manifest") loss_schedule = payload.get("action_flow_loss_schedule") + likelihood_contract = payload.get("action_flow_likelihood_contract") _require( isinstance(state_dict, Mapping) and state_dict, "checkpoint has no state_dict" ) @@ -861,17 +993,11 @@ def _validate_checkpoint( ) _require(isinstance(loops, Mapping) and loops, "checkpoint loop state is empty") if loss_schedule is not None: - _require( - loss_schedule - == { - "joint_objective_begins_at_global_step": 1, - "reconstruction_only_optimizer_steps": 1, - "joint_flow_weight": flow_weight, - "joint_reconstruction_weight": reconstruction_weight, - "joint_action_velocity_weight": 1.0, - "schema_version": 1, - }, - f"unexpected Action Flow loss schedule: {loss_schedule}", + expected_warmup_steps = _validate_checkpoint_loss_schedule( + loss_schedule, + config, + reconstruction_weight=reconstruction_weight, + flow_weight=flow_weight, ) state_tensors, state_scalars = _finite_tree(state_dict, "checkpoint.state_dict") optimizer_tensors, optimizer_scalars = _finite_tree( @@ -901,8 +1027,38 @@ def _validate_checkpoint( ) del immutable_payload, payload + wrapper_type = ActionFlowModelWrapper + if method == LIKELIHOOD_METHOD: + from egomimic.pl_utils.pl_model_action_flow_likelihood import ( + ActionFlowLikelihoodModelWrapper, + ) + + wrapper_type = ActionFlowLikelihoodModelWrapper + _require( + isinstance(likelihood_contract, Mapping), + "likelihood checkpoint lacks scientific contract", + ) + _require( + likelihood_contract + == { + "schema_version": 1, + "method": LIKELIHOOD_METHOD, + "num_levels": 32, + "interior_samples_per_content": 14, + "sigma_min": 0.1, + "sigma_max": 1.0, + "rho": 0.95, + "tau": 0.02, + "reduction": "sum_chunk_coordinates_mean_examples_constant_free_gaussian_bound", + "learned_reference_targets": "attached", + "sampler": "stochastic_reverse_gaussian_chain_with_action_output_noise", + "clean_reconstruction_objective": False, + "latent_fm_objective": False, + }, + "likelihood checkpoint scientific contract mismatch", + ) try: - restored = ActionFlowModelWrapper.load_from_checkpoint( + restored = wrapper_type.load_from_checkpoint( last_path, map_location="cpu", strict=True, @@ -913,21 +1069,29 @@ def _validate_checkpoint( f"strict ActionFlowModelWrapper reload failed: {last_path}" ) from error _require( - type(restored) is ActionFlowModelWrapper, + type(restored) is wrapper_type, f"checkpoint restored unexpected wrapper {type(restored)!r}", ) if loss_schedule is not None: _require( - restored.reconstruction_only_warmup_steps == 1, + restored.reconstruction_only_warmup_steps == expected_warmup_steps, "strict reload lost the reconstruction-only warmup", ) parameter_count = sum(parameter.numel() for parameter in restored.parameters()) - _require( - parameter_count == EXPECTED_PARAMETER_COUNT, - f"parameter count mismatch: {parameter_count} != {EXPECTED_PARAMETER_COUNT}", - ) + if method in (LEGACY_METHOD, STOPGRAD_METHOD): + _require( + parameter_count == EXPECTED_PARAMETER_COUNT, + f"parameter count mismatch: {parameter_count} != {EXPECTED_PARAMETER_COUNT}", + ) + else: + _require(config is not None, "typed candidate reload needs its exact config") + _validate_dimensions_and_modules(config, tuple(restored.model.pipeline.stages)) # These properties fail closed on duplicated or disconnected owners. - owners = (restored.encoder_e, restored.field_v, restored.decoder_g) + if method == LIKELIHOOD_METHOD: + stages = restored.model.pipeline.stages + owners = (stages[3].mean_encoder, stages[5].field, stages[6].decoder) + else: + owners = (restored.encoder_e, restored.field_v, restored.decoder_g) _require( len({id(owner) for owner in owners}) == 3, "Action Flow owners are aliased" ) @@ -939,7 +1103,7 @@ def _validate_checkpoint( if parameter.requires_grad ) gradient_routes = _validate_gradient_route_manifest( - gradient_route_manifest, trainable + gradient_route_manifest, trainable, method ) del restored @@ -949,6 +1113,7 @@ def _validate_checkpoint( "gradient_routes": gradient_routes, "global_step": 2, "loss_schedule": loss_schedule, + "likelihood_contract": likelihood_contract, "immutable_checkpoint_path": str(immutable_path), "immutable_checkpoint_sha256": _sha256(immutable_path), "optimizer_state_count": 1, @@ -1050,6 +1215,7 @@ def _validate_history( reconstruction_weight: float = 1.0, flow_weight: float = 1.0, expect_reconstruction_warmup: bool = False, + method: str = LEGACY_METHOD, ) -> dict[str, Any]: component_names = ( "TotalLoss", @@ -1058,13 +1224,23 @@ def _validate_history( "ReconstructionL1", "ActionVelocityLoss", ) + if method == LIKELIHOOD_METHOD: + component_names = ("TotalLoss", "InteriorBridgeNLL", "BoundaryNLL") components = tuple(f"Train/ActionFlow/{name}" for name in component_names) per_source_components = tuple(f"{name}/{SOURCE_LABEL}" for name in components) - gradient_labels = ("FM", "Reconstruction", "ActionVelocity") - gradient_pairs = ( - "FM__Reconstruction", - "FM__ActionVelocity", - "Reconstruction__ActionVelocity", + gradient_labels = ( + ("InteriorBridgeNLL", "BoundaryNLL") + if method == LIKELIHOOD_METHOD + else ( + ("FM", "ActionVelocity") + if method == GRAPH_METHOD + else ("FM", "Reconstruction", "ActionVelocity") + ) + ) + gradient_pairs = tuple( + f"{left}__{right}" + for i, left in enumerate(gradient_labels) + for right in gradient_labels[i + 1 :] ) telemetry = [ *(f"Train/ActionFlow/GradientNorm/{label}" for label in gradient_labels), @@ -1088,6 +1264,8 @@ def _validate_history( "Train/ActionFlow/Schedule/EffectiveFlowWeight", "Train/ActionFlow/Schedule/EffectiveActionVelocityWeight", ] + if method == LIKELIHOOD_METHOD: + telemetry = [name for name in telemetry if "/Schedule/" not in name] train_step, train = _complete_row( rows, (*components, *per_source_components, *telemetry), @@ -1095,20 +1273,40 @@ def _validate_history( label="Action Flow training/gradient telemetry", ) for name in telemetry: + empty_pair = method == STOPGRAD_METHOD and name.endswith("/FM__Reconstruction") if ( "GradientNorm" in name or "GradientParameterCount" in name or "IntersectionParameterCount" in name ): - _require(train[name] > 0.0, f"gradient reachability is empty: {name}") + _require( + train[name] == 0.0 if empty_pair else train[name] > 0.0, + f"gradient reachability differs: {name}", + ) if "GradientCosineDefined" in name: - _require(train[name] == 1.0, f"gradient cosine is undefined: {name}") + _require( + train[name] == (0.0 if empty_pair else 1.0), + f"gradient cosine defined flag differs: {name}", + ) if "GradientCosine/" in name: _require(-1.0 <= train[name] <= 1.0, f"gradient cosine is invalid: {name}") for name, expected in ( - ("Train/ActionFlow/Compute/FieldForwardCallsPerStep", 1.0), - ("Train/ActionFlow/Compute/FieldSampleEquivalentsPerStep", 14.0), - ("Train/ActionFlow/Compute/DecoderJVPCallsPerStep", 1.0), + ( + "Train/ActionFlow/Compute/FieldForwardCallsPerStep", + 2.0 if method == STOPGRAD_METHOD else 1.0, + ), + ( + "Train/ActionFlow/Compute/FieldSampleEquivalentsPerStep", + ( + 28.0 + if method == STOPGRAD_METHOD + else 15.0 if method == LIKELIHOOD_METHOD else 14.0 + ), + ), + ( + "Train/ActionFlow/Compute/DecoderJVPCallsPerStep", + 0.0 if method == LIKELIHOOD_METHOD else 1.0, + ), ): _require(train[name] == expected, f"unexpected compute telemetry: {name}") _require( @@ -1116,28 +1314,36 @@ def _validate_history( "CUDA peak allocation telemetry is empty", ) _require( - train["Train/ActionFlow/Schedule/ReconstructionOnly"] == 0.0, + method == LIKELIHOOD_METHOD + or train["Train/ActionFlow/Schedule/ReconstructionOnly"] == 0.0, "latest smoke optimizer step is not joint", ) _require( - math.isclose( - train["Train/ActionFlow/Schedule/EffectiveFlowWeight"], - flow_weight, - rel_tol=0.0, - abs_tol=1.0e-6, - ) - and train[ - "Train/ActionFlow/Schedule/EffectiveActionVelocityWeight" - ] - == 1.0, + method == LIKELIHOOD_METHOD + or ( + math.isclose( + train["Train/ActionFlow/Schedule/EffectiveFlowWeight"], + flow_weight, + rel_tol=0.0, + abs_tol=1.0e-6, + ) + and train["Train/ActionFlow/Schedule/EffectiveActionVelocityWeight"] == 1.0 + ), "joint smoke step did not enable both delayed objectives", ) for suffix in ("", f"/{SOURCE_LABEL}"): expected_total = ( - flow_weight * train[f"Train/ActionFlow/FlowMatchingLoss{suffix}"] - + reconstruction_weight - * train[f"Train/ActionFlow/ReconstructionLoss{suffix}"] - + train[f"Train/ActionFlow/ActionVelocityLoss{suffix}"] + ( + train[f"Train/ActionFlow/InteriorBridgeNLL{suffix}"] + + train[f"Train/ActionFlow/BoundaryNLL{suffix}"] + ) + if method == LIKELIHOOD_METHOD + else ( + flow_weight * train[f"Train/ActionFlow/FlowMatchingLoss{suffix}"] + + reconstruction_weight + * train[f"Train/ActionFlow/ReconstructionLoss{suffix}"] + + train[f"Train/ActionFlow/ActionVelocityLoss{suffix}"] + ) ) _require( math.isclose( @@ -1161,18 +1367,14 @@ def _validate_history( continue _require( concrete["Train/ActionFlow/Schedule/EffectiveFlowWeight"] == 0.0 - and concrete[ - "Train/ActionFlow/Schedule/EffectiveActionVelocityWeight" - ] + and concrete["Train/ActionFlow/Schedule/EffectiveActionVelocityWeight"] == 0.0, "warmup smoke step enabled a delayed objective", ) for suffix in ("", f"/{SOURCE_LABEL}"): expected_total = ( reconstruction_weight - * concrete[ - f"Train/ActionFlow/ReconstructionLoss{suffix}" - ] + * concrete[f"Train/ActionFlow/ReconstructionLoss{suffix}"] ) _require( math.isclose( @@ -1196,16 +1398,7 @@ def _validate_history( "Valid/EnergyScoreDiversity@32", ): validity.extend((base, f"{base}/{SOURCE_LABEL}")) - validity.extend( - f"Valid/ActionFlow/{name}" - for name in ( - "TotalLoss", - "FlowMatchingLoss", - "ReconstructionLoss", - "ReconstructionL1", - "ActionVelocityLoss", - ) - ) + validity.extend(f"Valid/ActionFlow/{name}" for name in component_names) diagnostics = ( "Valid/ActionFlow/CleanReconstructionMSE", "Valid/ActionFlow/CleanReconstructionNativeMSE", @@ -1220,6 +1413,10 @@ def _validate_history( "Valid/ActionFlow/Alignment/CKA/encoder_00__field_00/t0500", "Valid/ActionFlow/Alignment/CKNNA/encoder_00__field_00/t0500", ) + if method == LIKELIHOOD_METHOD: + diagnostics = () + elif method == GRAPH_METHOD: + diagnostics = tuple(name for name in diagnostics if "/Alignment/CK" not in name) diagnostics = ( *diagnostics, *( @@ -1235,9 +1432,16 @@ def _validate_history( label="scheduled validation", ) expected_valid_total = ( - flow_weight * valid["Valid/ActionFlow/FlowMatchingLoss"] - + reconstruction_weight * valid["Valid/ActionFlow/ReconstructionLoss"] - + valid["Valid/ActionFlow/ActionVelocityLoss"] + ( + valid["Valid/ActionFlow/InteriorBridgeNLL"] + + valid["Valid/ActionFlow/BoundaryNLL"] + ) + if method == LIKELIHOOD_METHOD + else ( + flow_weight * valid["Valid/ActionFlow/FlowMatchingLoss"] + + reconstruction_weight * valid["Valid/ActionFlow/ReconstructionLoss"] + + valid["Valid/ActionFlow/ActionVelocityLoss"] + ) ) _require( math.isclose( @@ -1285,10 +1489,16 @@ def _step_two_artifact(root: Path, *, label: str) -> tuple[Path, Mapping[str, An and (path.suffix in {".tmp", ".temporary"} or ".temporary" in path.name) ] _require(not leftovers, f"unfinished {label} artifacts: {leftovers}") - candidates = sorted([ - *root.glob("epoch-*-step-2/rank-0-batch-*.pt"), - *([] if execution is not None else root.glob("job-*-restart-*/epoch-*-step-2/rank-0-batch-*.pt")), - ]) + candidates = sorted( + [ + *root.glob("epoch-*-step-2/rank-0-batch-*.pt"), + *( + [] + if execution is not None + else root.glob("job-*-restart-*/epoch-*-step-2/rank-0-batch-*.pt") + ), + ] + ) _require( len(candidates) == 1, f"expected one step-2 {label} artifact: {candidates}" ) @@ -1298,12 +1508,17 @@ def _step_two_artifact(root: Path, *, label: str) -> tuple[Path, Mapping[str, An namespace = candidates[0].parent.parent.name if namespace.startswith("job-"): recorded = payload.get("execution") - _require(isinstance(recorded, Mapping), f"{label} artifact lacks execution identity") + _require( + isinstance(recorded, Mapping), f"{label} artifact lacks execution identity" + ) expected_namespace = ( f"job-{recorded.get('slurm_job_id')}" f"-restart-{recorded.get('slurm_restart_count')}" ) - _require(namespace == expected_namespace, f"{label} artifact execution identity mismatch") + _require( + namespace == expected_namespace, + f"{label} artifact execution identity mismatch", + ) _finite_tree(payload, f"{label} artifact") return candidates[0], payload @@ -1491,6 +1706,30 @@ def _validate_artifacts( "typed EnergyScore checkpoint binding differs", ) + if action_flow_method(config) == LIKELIHOOD_METHOD: + # No clean-latent/ODE artifact exists for the discrete likelihood model. + # Its actual learned-reference gradient routes are checkpoint-bound. + _require( + checkpoint.get("likelihood_contract"), + "likelihood scientific checkpoint receipt missing", + ) + return { + "energy_score": { + "path": str(energy_path), + "sha256": _sha256(energy_path), + "checkpoint_global_step": checkpoint["global_step"], + "checkpoint_sha256": checkpoint["checkpoint_sha256"], + "identity_sha256": energy["identity_sha256"], + }, + "method_specific": { + "likelihood_contract": checkpoint["likelihood_contract"], + "gradient_route_manifest_sha256": checkpoint["gradient_routes"][ + "manifest_sha256" + ], + "ode_diagnostics": "not_applicable_discrete_gaussian_reverse_chain", + }, + } + diagnostic_root = _artifact_root( config, "evaluator.action_flow_diagnostics.artifact_root", run_dir ) @@ -1738,6 +1977,12 @@ def verify_smoke( ) approved_reconstruction_weight = APPROVED_EXPERIMENTS[experiment][1] approved_flow_weight = APPROVED_EXPERIMENTS[experiment][2] + method = action_flow_method(config, experiment) + if method == LIKELIHOOD_METHOD: + _require( + expected_reconstruction_weight is None and expected_flow_weight is None, + "likelihood smoke does not accept FM/reconstruction weight arguments", + ) if expected_reconstruction_weight is not None: _require( math.isclose( @@ -1773,6 +2018,8 @@ def verify_smoke( run_dir, reconstruction_weight=approved_reconstruction_weight, flow_weight=approved_flow_weight, + method=method, + config=config, ) expect_reconstruction_warmup = "warmup10k" in experiment if expect_reconstruction_warmup: @@ -1800,6 +2047,7 @@ def verify_smoke( reconstruction_weight=APPROVED_EXPERIMENTS[experiment][1], flow_weight=APPROVED_EXPERIMENTS[experiment][2], expect_reconstruction_warmup=expect_reconstruction_warmup, + method=method, ) artifacts = _validate_artifacts( config=config, run_dir=run_dir, identities=identities, checkpoint=checkpoint @@ -1809,6 +2057,7 @@ def verify_smoke( "artifacts": artifacts, "checkpoint": checkpoint, "experiment": experiment, + "action_flow_method": method, "gpu_probes": gpu_probes, "identities": identities, "metrics": metrics, diff --git a/tests/test_action_flow_candidate_codec.py b/tests/test_action_flow_candidate_codec.py new file mode 100644 index 000000000..a82faeff5 --- /dev/null +++ b/tests/test_action_flow_candidate_codec.py @@ -0,0 +1,139 @@ +"""Exact-section and JVP behavior of the restricted graph diagnostic.""" + +import pytest +import torch +import torch.nn as nn +from torch.func import jvp + +from egomimic.models.action_flow_codec import GraphSectionSequenceCodec +from egomimic.pipeline.core import Pipeline +from egomimic.pipeline.stages_action_flow import ( + ActionFlowObjectiveStage, + ConditionalVelocityStage, + ContentDecoderStage, + ContentEncoderStage, + LatentBridgeStage, +) + + +class _EncoderView(nn.Module): + def __init__(self, codec): + super().__init__() + self.graph = codec.graph + + def forward(self, content): + return torch.cat((content, self.graph(content)), dim=-1) + + +class _Field(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(8, 8) + self.condition = nn.Linear(3, 8) + + def forward(self, value, time, condition, *, condition_drop_mask): + condition = condition.masked_fill(condition_drop_mask[:, None], 0) + return self.linear(value).tanh() + self.condition(condition)[:, None] + + +def _codec(): + return GraphSectionSequenceCodec(action_dim=4, latent_dim=8, horizon=4) + + +def test_graph_section_identity_survives_joint_optimizer_updates_and_shared_ownership(): + torch.manual_seed(46) + codec = _codec().double() + encoder = ContentEncoderStage(_EncoderView(codec)) + decoder = ContentDecoderStage(codec) + field = ConditionalVelocityStage(_Field().double()) + objective = ActionFlowObjectiveStage( + reconstruction_weight=0.0, residual_key="action_flow/fm_velocity_residual" + ) + pipeline = Pipeline([ + encoder, LatentBridgeStage(samples_per_content=2), field, decoder, objective + ]) + parameters = list(pipeline.parameters()) + assert len({id(parameter) for parameter in parameters}) == len(parameters) + assert encoder.encoder.graph is decoder.decoder.graph + optimizer = torch.optim.AdamW(parameters, lr=1e-3) + assert len(optimizer.param_groups[0]["params"]) == len(parameters) + snapshots = { + name: [parameter.detach().clone() for parameter in module.parameters()] + for name, module in (("graph", codec.graph), ("residual", codec.residual)) + } + target = torch.randn(2, 4, 4, dtype=torch.float64) + for _ in range(2): + latent = codec.encode(target) + torch.testing.assert_close(latent[..., :4], target, rtol=0, atol=0) + torch.testing.assert_close(codec(latent), target, rtol=0, atol=1e-12) + batch = pipeline({ + "target": target, + "sampler/noise": torch.randn(2, 4, 8, dtype=torch.float64), + "condition": torch.randn(2, 3, dtype=torch.float64), + }) + torch.testing.assert_close( + batch["loss/action_flow"], + batch["log/action_flow_fm"] + batch["log/action_flow_action_velocity"], + rtol=0, atol=0, + ) + optimizer.zero_grad() + batch["loss/action_flow"].backward() + for module in (codec.graph, codec.residual, field.field): + gradients = [p.grad for p in module.parameters() if p.grad is not None] + assert gradients and all(torch.isfinite(g).all() for g in gradients) + assert sum(g.abs().sum() for g in gradients) > 0 + optimizer.step() + torch.testing.assert_close(codec(codec.encode(target)), target, rtol=0, atol=1e-12) + for name, module in (("graph", codec.graph), ("residual", codec.residual)): + assert any(not torch.equal(old, new) for old, new in zip(snapshots[name], module.parameters())) + + +def test_graph_section_jvp_matches_finite_difference_with_trainable_graph(): + torch.manual_seed(7) + codec = _codec().double() + latent = torch.randn(2, 4, 8, dtype=torch.float64, requires_grad=True) + tangent = torch.randn_like(latent, requires_grad=True) + _, predicted = jvp(codec, (latent,), (tangent,)) + delta = 1e-6 + expected = (codec(latent + delta*tangent) - codec(latent - delta*tangent))/(2*delta) + torch.testing.assert_close(predicted, expected, rtol=2e-5, atol=2e-6) + gradients = torch.autograd.grad(predicted.square().mean(), (latent, tangent, *codec.graph.parameters())) + assert all(torch.isfinite(gradient).all() for gradient in gradients) + assert sum(gradient.abs().sum() for gradient in gradients[2:]) > 0 + + +def test_graph_section_context_free_shape_and_small_private_networks(): + codec = GraphSectionSequenceCodec(action_dim=4, latent_dim=8, horizon=16) + for module in (codec.graph, codec.residual): + assert module.depth == 2 + assert sum(parameter.numel() for parameter in module.parameters()) < 12_000 + with pytest.raises(ValueError, match="expected sequence shape"): + codec(torch.randn(2, 15, 8)) + with pytest.raises(ValueError, match="expected sequence shape"): + codec.encode(torch.randn(2, 16, 3)) + + +def test_graph_section_mixed_precision_identity_and_jvp_backward_are_finite(): + torch.manual_seed(19) + codec = _codec() + content = torch.randn(2, 4, 4) + latent = torch.randn(2, 4, 8, requires_grad=True) + tangent = torch.randn_like(latent, requires_grad=True) + with torch.autocast("cpu", dtype=torch.bfloat16): + reconstructed = codec(codec.encode(content)) + _, decoded_tangent = jvp(codec, (latent,), (tangent,)) + loss = decoded_tangent.square().mean() + torch.testing.assert_close(reconstructed, content, rtol=0, atol=0) + assert torch.isfinite(decoded_tangent).all() + loss.backward() + for module in (codec.graph, codec.residual): + gradients = [p.grad for p in module.parameters() if p.grad is not None] + assert gradients and all(torch.isfinite(g).all() for g in gradients) + assert sum(g.abs().sum() for g in gradients) > 0 + + +@pytest.mark.parametrize("kwargs", [{"action_dim":8}, {"latent_dim":3}, {"dropout":0.1}]) +def test_graph_section_rejects_invalid_or_stochastic_section(kwargs): + options = {"action_dim":4, "latent_dim":8, "horizon":4, **kwargs} + with pytest.raises(ValueError): + GraphSectionSequenceCodec(**options) diff --git a/tests/test_action_flow_candidate_configs.py b/tests/test_action_flow_candidate_configs.py new file mode 100644 index 000000000..edb2743cb --- /dev/null +++ b/tests/test_action_flow_candidate_configs.py @@ -0,0 +1,166 @@ +"""Exercise the maintained Hydra rows and their candidate-specific telemetry.""" + +from pathlib import Path + +import hydra +import pytest +import torch +from hydra import compose, initialize_config_dir +from omegaconf import OmegaConf + +from egomimic.eval.action_flow_diagnostics import ActionFlowDiagnostics +from egomimic.pl_utils.pl_model_action_flow import ActionFlowModelWrapper + + +CONFIG_DIR = Path(__file__).parents[1] / "egomimic/hydra_configs" +ROWS = [ + ("latent_fm_sg_recon1", "latent_fm_stopgrad", 1.0), + ("graph_section", "graph_section_diagnostic", 0.0), +] + + +def _compose(row): + with initialize_config_dir(version_base=None, config_dir=str(CONFIG_DIR.resolve())): + return compose( + config_name="train_zarr_cartesian", + overrides=[ + f"+experiment=pusht/action_flow_bc_usocket_{row}_s42", + "++paths.root_dir=.", + ], + ) + + +@pytest.mark.parametrize("row,method,reconstruction_weight", ROWS) +def test_candidate_row_instantiates_exact_dimensions_and_shared_codec( + row, method, reconstruction_weight +): + cfg = _compose(row) + assert cfg.model.action_flow_method == cfg.train.action_flow_method == method + assert cfg.run_provenance.objective.method == method + assert cfg.model.action_horizon == 16 + assert cfg.model.latent_dim == 8 + assert cfg.model.action_dim == 4 + assert cfg.model.reconstruction_weight == reconstruction_weight + assert cfg.run_provenance.objective.decoded_noise_scale_weight == 0 + assert cfg.model.flow_samples_per_content == 14 + assert cfg.model.condition_dropout_probability == 0.3 + assert cfg.model.optimizer.lr == 3e-5 + + algo = hydra.utils.instantiate(cfg.model.pipeline, device="cpu") + stages = algo.pipeline.stages + encoder, field, decoder, objective = stages[3], stages[5], stages[6], stages[7] + assert field.field.horizon == 16 + assert field.field.input_dim == field.field.output_dim == 8 + assert field.field.condition_dim == 67 + assert field.field.depth == 12 + assert 39_000_000 <= sum(p.numel() for p in field.parameters()) <= 41_000_000 + assert objective.reconstruction_weight == reconstruction_weight + assert stages[1].num_tokens == 16 and stages[1].latent_dim == 8 + parameters = list(algo.nets.parameters()) + assert len(parameters) == len({id(parameter) for parameter in parameters}) + if method == "latent_fm_stopgrad": + assert field.flow_clean_gradient_mode == "all_stopgrad" + assert objective.residual_key == field.flow_residual_key + assert decoder.residual_key == field.residual_key + assert field.flow_residual_key != field.residual_key + assert encoder.encoder.depth == decoder.decoder.depth == 2 + else: + assert field.flow_clean_gradient_mode == "full" + assert encoder.encoder.graph is decoder.decoder.graph + encoder_parameters = {id(parameter) for parameter in encoder.parameters()} + residual_parameters = {id(parameter) for parameter in decoder.decoder.residual.parameters()} + assert encoder_parameters.isdisjoint(residual_parameters) + assert decoder.decoder.latent_dim - decoder.decoder.action_dim == 4 + assert not cfg.run_provenance.objective.reconstruction_is_optimizer_objective + assert cfg.run_provenance.requirements.r6 == "FAIL_restricted_diagnostic" + assert not cfg.evaluator.action_flow_diagnostics.capture_activations + target = torch.randn(2, 16, 4) + torch.testing.assert_close( + decoder.decoder(encoder.encoder(target)), target, rtol=0, atol=1e-6 + ) + + +@pytest.mark.parametrize("row,method,reconstruction_weight", ROWS) +def test_candidate_wrapper_gradient_routes_and_compute_contract( + monkeypatch, row, method, reconstruction_weight +): + cfg = _compose(row) + # Preserve the real Hydra codec/factory/stage wiring and all sequence + # dimensions; replace only the expensive shared field width/depth. Inputs + # are already encoded observations, so image/data nodes are not exercised. + pipeline_cfg = OmegaConf.create( + OmegaConf.to_container(cfg.model.pipeline, resolve=True) + ) + pipeline_cfg.stages = list(pipeline_cfg.stages)[3:] + field_config = pipeline_cfg.stages[2].field + field_config.hidden_dim = 32 + field_config.depth = 2 + field_config.num_heads = 4 + field_config.feedforward_dim = 64 + field_config.time_embedding_dim = 32 + algo = hydra.utils.instantiate(pipeline_cfg, device="cpu") + wrapper = ActionFlowModelWrapper(pipeline=algo, gradient_telemetry_cadence=1) + logged = {} + monkeypatch.setattr( + wrapper, "log", lambda key, value, **kwargs: logged.update({key: value}) + ) + torch.manual_seed(51) + loss = wrapper.training_step({"fixture": { + "target": torch.randn(2, 16, 4), + "condition": torch.randn(2, 67), + "sampler/noise": torch.randn(2, 16, 8), + }}, 0) + assert torch.isfinite(loss) + loss.backward() + assert wrapper._gradient_route_manifest is not None + routes = wrapper._gradient_route_manifest["routes"] + intersections = wrapper._gradient_route_manifest["intersections"] + assert routes["FM"] and routes["ActionVelocity"] + assert intersections["FM__ActionVelocity"] + assert all(torch.isfinite(p.grad).all() for p in wrapper.parameters() if p.grad is not None) + calls = 2 if method == "latent_fm_stopgrad" else 1 + assert logged["Train/ActionFlow/Compute/FieldForwardCallsPerStep"] == calls + assert logged["Train/ActionFlow/Compute/FieldSampleEquivalentsPerStep"] == calls * 14 + if method == "latent_fm_stopgrad": + assert intersections["FM__Reconstruction"] == [] + assert logged["Train/ActionFlow/GradientCosineDefined/FM__Reconstruction"] == 0 + assert logged["Train/ActionFlow/GradientIntersectionParameterCount/FM__Reconstruction"] == 0 + assert routes["Reconstruction"] + else: + assert "Reconstruction" not in routes + assert "Train/ActionFlow/GradientNorm/Reconstruction" not in logged + torch.testing.assert_close( + loss, + logged["Train/ActionFlow/FlowMatchingLoss"] + + logged["Train/ActionFlow/ActionVelocityLoss"], + ) + + +def test_graph_factory_rejects_independent_encoder_configuration(): + cfg = _compose("graph_section") + pipeline_cfg = OmegaConf.create( + OmegaConf.to_container(cfg.model.pipeline, resolve=True) + ) + pipeline_cfg.stages = [pipeline_cfg.stages[3], pipeline_cfg.stages[6]] + pipeline_cfg.stages[0].encoder = {"_target_": "torch.nn.Identity"} + with pytest.raises(hydra.errors.InstantiationException, match="second codec"): + hydra.utils.instantiate(pipeline_cfg, device="cpu") + + +@pytest.mark.parametrize("row,method,reconstruction_weight", ROWS) +def test_candidate_diagnostic_config_constructs_without_stale_activation_settings( + tmp_path, row, method, reconstruction_weight +): + cfg = _compose(row) + diagnostics = cfg.evaluator.action_flow_diagnostics + diagnostics.noise_seed_bank_path = str( + CONFIG_DIR / "evaluator/energy_score_seed_bank_k32_v1.json" + ) + diagnostics.artifact_root = str(tmp_path / "diagnostics") + config = OmegaConf.to_container(diagnostics, resolve=True) + instance = ActionFlowDiagnostics(config) + assert instance.jacobian_samples == 2 + assert instance.capture_activations == (method == "latent_fm_stopgrad") + if method == "graph_section_diagnostic": + assert instance.activation_layer_map == () + assert instance.cknna_k == 0 diff --git a/tests/test_action_flow_candidate_gates.py b/tests/test_action_flow_candidate_gates.py new file mode 100644 index 000000000..aeba30bca --- /dev/null +++ b/tests/test_action_flow_candidate_gates.py @@ -0,0 +1,234 @@ +"""Method-specific gates: real configs plus tiny, honest telemetry receipts.""" + +import hashlib +import re +import subprocess + +import pytest +import torch +from hydra.core.hydra_config import HydraConfig +from omegaconf import OmegaConf, open_dict + +from tools.validate_action_flow_config import ( + CANDIDATE_METHODS, + GRAPH_METHOD, + LIKELIHOOD_METHOD, + STOPGRAD_METHOD, + PreflightError, + compose_experiment, + validate_experiment, + validate_method_contract, +) +from test_verify_action_flow_training_smoke import MODULE, HEAD, _history_row + + +@pytest.mark.parametrize("experiment", CANDIDATE_METHODS) +def test_real_candidate_config_and_parameter_manifest(experiment): + report, _ = validate_experiment(experiment) + assert report["status"] == "PASS" + assert report["action_flow_method"] == CANDIDATE_METHODS[experiment] + assert report["topology"]["shared_field_instance"] + assert report["topology"]["shared_decoder_instance"] + assert len(report["topology"]["train_order"]) == 8 + assert ( + sum( + v["total"] for k, v in report["parameters"].items() if k != "pipeline_total" + ) + == report["parameters"]["pipeline_total"]["total"] + ) + + +@pytest.mark.parametrize("experiment", CANDIDATE_METHODS) +def test_candidate_two_optimizer_update_config_gate(experiment, tmp_path, monkeypatch): + cfg = compose_experiment(experiment) + normalization = tmp_path / "norm_stats.json" + normalization.write_text("{}") + norm_hash = hashlib.sha256(normalization.read_bytes()).hexdigest() + with open_dict(cfg): + cfg.trainer.max_steps = 2 + cfg.trainer.val_check_interval = 1 + cfg.trainer.limit_val_batches = 1 + cfg.callbacks.model_checkpoint.every_n_train_steps = 1 + cfg.model.gradient_telemetry_cadence = 2 + cfg.norm_stats.precomputed_norm_path = str(normalization) + cfg.run_provenance.source_commit = HEAD + cfg.run_provenance.normalization_sha256 = norm_hash + cfg.evaluator.artifact_root = str(tmp_path / "energy") + if CANDIDATE_METHODS[experiment] != LIKELIHOOD_METHOD: + cfg.evaluator.action_flow_diagnostics.artifact_root = str( + tmp_path / "diagnostics" + ) + cfg.hydra.runtime.cwd = str(MODULE.REPOSITORY_ROOT) + cfg.hydra.runtime.output_dir = str(tmp_path) + HydraConfig.instance().set_config(cfg) + selected = OmegaConf.masked_copy(cfg, [key for key in cfg if key != "hydra"]) + path = tmp_path / "config.yaml" + OmegaConf.save( + OmegaConf.create(OmegaConf.to_container(selected, resolve=True)), path + ) + monkeypatch.setattr(MODULE, "_git_head", lambda: HEAD) + MODULE._validate_config( + config_path=path, + experiment=experiment, + run_dir=tmp_path, + expected_head=HEAD, + expected_config_sha256=None, + expected_split_sha256=None, + expected_normalization_sha256=norm_hash, + ) + + +@pytest.mark.parametrize( + "experiment,path,value", + [ + ( + "pusht/action_flow_bc_usocket_latent_fm_sg_recon1_s42", + "model.pipeline.stages.5.flow_clean_gradient_mode", + "attached", + ), + ( + "pusht/action_flow_bc_usocket_graph_section_s42", + "model.reconstruction_weight", + 1.0, + ), + ( + "pusht/action_flow_bc_usocket_bridge_likelihood_s42", + "run_provenance.inference.sampler", + "reverse_euler", + ), + ( + "pusht/action_flow_bc_usocket_bridge_likelihood_s42", + "model.pipeline.stages.7.tau", + 0.01, + ), + ], +) +def test_method_contract_rejects_scientific_identity_drift(experiment, path, value): + cfg = compose_experiment(experiment) + OmegaConf.update(cfg, path, value) + with pytest.raises(PreflightError): + validate_method_contract(cfg, experiment) + + +def test_stopgrad_telemetry_requires_empty_fm_reconstruction_intersection(): + row = _history_row() + for key in list(row): + if key.endswith("/FM__Reconstruction"): + row[key] = 0.0 + row["Train/ActionFlow/Compute/FieldForwardCallsPerStep"] = 2.0 + row["Train/ActionFlow/Compute/FieldSampleEquivalentsPerStep"] = 28.0 + MODULE._validate_history({1: row}, method=STOPGRAD_METHOD) + row["Train/ActionFlow/GradientIntersectionParameterCount/FM__Reconstruction"] = 1.0 + with pytest.raises(MODULE.SmokeVerificationError, match="reachability"): + MODULE._validate_history({1: row}, method=STOPGRAD_METHOD) + + +def test_graph_telemetry_does_not_invent_reconstruction_gradients(): + row = { + key: value + for key, value in _history_row().items() + if not ("/Gradient" in key and "Reconstruction" in key) + and "/Alignment/CK" not in key + } + row["Train/ActionFlow/TotalLoss_step"] = 2.5 + row[f"Train/ActionFlow/TotalLoss/{MODULE.SOURCE_LABEL}_step"] = 2.5 + row["Valid/ActionFlow/TotalLoss"] = 3.0 + MODULE._validate_history({1: row}, method=GRAPH_METHOD, reconstruction_weight=0.0) + + +def test_likelihood_history_requires_real_nll_and_common_validation_only(): + row = { + key: value + for key, value in _history_row().items() + if "/ActionFlow/" not in key or "/Compute/" in key + } + for component in ("InteriorBridgeNLL", "BoundaryNLL", "TotalLoss"): + value = 2.5 if component == "TotalLoss" else 1.25 + row[f"Train/ActionFlow/{component}_step"] = value + row[f"Train/ActionFlow/{component}/{MODULE.SOURCE_LABEL}_step"] = value + row[f"Valid/ActionFlow/{component}"] = value + for label in ("InteriorBridgeNLL", "BoundaryNLL"): + row[f"Train/ActionFlow/GradientNorm/{label}"] = 0.5 + row[f"Train/ActionFlow/GradientParameterCount/{label}"] = 100.0 + pair = "InteriorBridgeNLL__BoundaryNLL" + row[f"Train/ActionFlow/GradientCosine/{pair}"] = 0.25 + row[f"Train/ActionFlow/GradientCosineDefined/{pair}"] = 1.0 + row[f"Train/ActionFlow/GradientIntersectionParameterCount/{pair}"] = 50.0 + row["Train/ActionFlow/Compute/FieldSampleEquivalentsPerStep"] = 15.0 + row["Train/ActionFlow/Compute/DecoderJVPCallsPerStep"] = 0.0 + MODULE._validate_history({1: row}, method=LIKELIHOOD_METHOD) + del row["Valid/EnergyScore@32"] + with pytest.raises(MODULE.SmokeVerificationError): + MODULE._validate_history({1: row}, method=LIKELIHOOD_METHOD) + + +@pytest.mark.parametrize("method", [STOPGRAD_METHOD, GRAPH_METHOD, LIKELIHOOD_METHOD]) +def test_checkpoint_gradient_receipt_uses_exact_method_routes(method): + named = [ + (f"nets.pipeline.stages.{stage}.weight", torch.nn.Parameter(torch.ones(2))) + for stage in (0, 3, 5, 6) + ] + groups = { + "FM": [0, 2] if method == STOPGRAD_METHOD else [0, 1, 2], + "Reconstruction": [1, 3], + "ActionVelocity": [0, 1, 2, 3], + } + if method == GRAPH_METHOD: + del groups["Reconstruction"] + elif method == LIKELIHOOD_METHOD: + groups = {"InteriorBridgeNLL": [0, 1, 2], "BoundaryNLL": [0, 1, 2, 3]} + routes = { + label: [ + {"name": named[i][0], "shape": [2], "numel": 2, "dtype": "torch.float32"} + for i in indices + ] + for label, indices in groups.items() + } + labels = list(groups) + core = { + "schema_version": 1, + "routes": routes, + "route_sha256": { + k: MODULE._canonical_json_sha256(v) for k, v in routes.items() + }, + "intersections": { + f"{left}__{right}": [ + named[i][0] for i in groups[left] if i in groups[right] + ] + for index, left in enumerate(labels) + for right in labels[index + 1 :] + }, + } + manifest = {**core, "manifest_sha256": MODULE._canonical_json_sha256(core)} + MODULE._validate_gradient_route_manifest(manifest, named, method=method) + + +@pytest.mark.parametrize( + "requested,selected,accepted", + [ + ("gpu-h100,gpu-h200", "gpu-h100", True), + ("gpu-h100,gpu-h200", "gpu-h200", True), + ("ice-gpu", "ice-gpu", True), + ("gpu-h100,gpu-h200", "gpu", False), + ("gpu-h100", "gpu-h200", False), + ("gpu-h100,", "gpu-h100", False), + ], +) +def test_portable_launcher_accepts_only_explicit_partition_members( + requested, selected, accepted +): + source = ( + MODULE.REPOSITORY_ROOT / "scripts/train/launch_action_flow_usocket.sbatch" + ).read_text() + function = re.search(r"partition_allowed\(\) \{.*?\n\}", source, re.S).group() + result = subprocess.run( + [ + "bash", + "-c", + function + '\npartition_allowed "$1" "$2"', + "test", + requested, + selected, + ] + ) + assert (result.returncode == 0) is accepted diff --git a/tests/test_action_flow_candidate_stages.py b/tests/test_action_flow_candidate_stages.py new file mode 100644 index 000000000..fe154ada8 --- /dev/null +++ b/tests/test_action_flow_candidate_stages.py @@ -0,0 +1,136 @@ +"""Gradient routes for the separately approved latent-FM-only SG candidate.""" + +import copy + +import pytest +import torch +import torch.nn as nn + +from egomimic.pipeline.stages_action_flow import ( + ActionFlowObjectiveStage, + ConditionalVelocityStage, + ContentDecoderStage, + ContentEncoderStage, + LatentBridgeStage, +) + + +class _Field(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.tensor(0.4, dtype=torch.float64)) + self.condition = nn.Linear(3, 8).double() + self.calls = [] + + def forward(self, state, time, condition, *, condition_drop_mask): + self.calls.append((state, time, condition, condition_drop_mask)) + masked = condition.masked_fill(condition_drop_mask[:, None], 0) + return ( + self.weight * state.tanh() + + self.condition(masked)[:, None] + + time[:, None, None] + ) + + +def _forward(mode, encoder, decoder, field, target, noise, condition): + stages = [ + ContentEncoderStage(encoder), + LatentBridgeStage(samples_per_content=3, condition_dropout_probability=0.3), + ConditionalVelocityStage(field, flow_clean_gradient_mode=mode), + ContentDecoderStage(decoder), + ActionFlowObjectiveStage(residual_key="action_flow/fm_velocity_residual"), + ] + batch = {"target": target, "sampler/noise": noise, "condition": condition} + torch.manual_seed(72) + for stage in stages: + batch = stage(batch) + return batch + + +def _models_and_inputs(): + torch.manual_seed(41) + return ( + nn.Sequential(nn.Linear(4, 8), nn.Tanh()).double(), + nn.Sequential(nn.Linear(8, 6), nn.SiLU(), nn.Linear(6, 4)).double(), + _Field(), + torch.randn(3, 4, 4, dtype=torch.float64), + torch.randn(3, 4, 8, dtype=torch.float64), + torch.randn(3, 3, dtype=torch.float64, requires_grad=True), + ) + + +def test_latent_only_stopgrad_blocks_both_clean_fm_routes_not_shared_condition(): + encoder, decoder, field, target, noise, condition = _models_and_inputs() + batch = _forward("all_stopgrad", encoder, decoder, field, target, noise, condition) + fm = batch["log/action_flow_fm"] + gradients = torch.autograd.grad( + fm, tuple(encoder.parameters()), allow_unused=True, retain_graph=True + ) + assert all(gradient is None for gradient in gradients) + field_gradients = torch.autograd.grad( + fm, tuple(field.parameters()), retain_graph=True + ) + assert all(torch.isfinite(gradient).all() for gradient in field_gradients) + assert sum(gradient.abs().sum() for gradient in field_gradients) > 0 + condition_gradient = torch.autograd.grad(fm, condition)[0] + assert torch.isfinite(condition_gradient).all() + assert condition_gradient.abs().sum() > 0 + + attached, detached = field.calls + assert attached[0].requires_grad and not detached[0].requires_grad + torch.testing.assert_close(attached[0], detached[0], rtol=0, atol=0) + # No new bridge noise, time, observation encoding, or dropout mask draw. + for index in (1, 2, 3): + assert attached[index] is detached[index] + torch.testing.assert_close( + batch["action_flow/fm_velocity_residual"], + batch["action_flow/velocity_residual"], rtol=0, atol=0, + ) + + +def test_latent_only_sg_preserves_all_action_flow_gradients_and_diagnostics(): + originals = _models_and_inputs() + results = [] + for mode in ("full", "all_stopgrad"): + encoder, decoder, field = [copy.deepcopy(module) for module in originals[:3]] + target, noise = originals[3:5] + condition = originals[5].detach().clone().requires_grad_() + batch = _forward(mode, encoder, decoder, field, target, noise, condition) + params = (*encoder.parameters(), *decoder.parameters(), *field.parameters(), condition) + gradients = torch.autograd.grad( + batch["log/action_flow_action_velocity"], params, + allow_unused=True, retain_graph=True, + ) + # The encoder must still receive generative supervision through the + # bridge state, target tangent, and decoder Jacobian evaluation point. + encoder_gradients = gradients[:len(tuple(encoder.parameters()))] + assert all(gradient is not None for gradient in encoder_gradients) + assert sum(gradient.abs().sum() for gradient in encoder_gradients) > 0 + if mode == "full": + fm_gradients = torch.autograd.grad( + batch["log/action_flow_fm"] , tuple(encoder.parameters()) + ) + assert sum(gradient.abs().sum() for gradient in fm_gradients) > 0 + assert len(field.calls) == 1 + assert batch["action_flow/fm_velocity_residual"] is batch["action_flow/velocity_residual"] + results.append((batch, gradients)) + for key in ("action_flow/state", "action_flow/target_velocity", + "action_flow/predicted_velocity", "action_flow/velocity_residual", + "action_flow/decoded_velocity_residual"): + torch.testing.assert_close(results[0][0][key], results[1][0][key], rtol=0, atol=0) + for full, isolated in zip(results[0][1], results[1][1]): + if full is None: + assert isolated is None + else: + torch.testing.assert_close(full, isolated, rtol=0, atol=0) + + +@pytest.mark.parametrize("mode", ["target_stopgrad", "all", None]) +def test_candidate_rejects_unsupported_gradient_modes(mode): + with pytest.raises(ValueError, match="flow_clean_gradient_mode"): + ConditionalVelocityStage(_Field(), flow_clean_gradient_mode=mode) + + +def test_fm_key_cannot_clobber_action_flow_residual(): + with pytest.raises(ValueError, match="must be distinct"): + ConditionalVelocityStage(_Field(), flow_residual_key="action_flow/velocity_residual") diff --git a/tests/test_action_flow_likelihood.py b/tests/test_action_flow_likelihood.py new file mode 100644 index 000000000..a0eac32e8 --- /dev/null +++ b/tests/test_action_flow_likelihood.py @@ -0,0 +1,435 @@ +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +import hydra +import pytest +import torch +from hydra import compose, initialize_config_dir +from torch import nn + +from egomimic.models.action_flow_codec import ContextFreeSequenceDecoder +from egomimic.models.action_flow_likelihood import ( + GaussianBridgeSchedule, + TimeDependentSequenceMean, +) +from egomimic.pipeline.algo import PipelineAlgo +from egomimic.pipeline.core import Pipeline +from egomimic.pipeline.stages_action_flow_likelihood import ( + ConditionalReverseMeanStage, + GaussianBridgeNoisingStage, + GaussianBridgeObjectiveStage, + LikelihoodDecoderStage, + LikelihoodReferenceStage, +) +from egomimic.pl_utils.pl_model_action_flow_likelihood import ( + ActionFlowLikelihoodModelWrapper, +) + + +class TinyField(nn.Module): + def __init__(self, latent_dim=3, condition_dim=2): + super().__init__() + self.linear = nn.Linear(latent_dim + condition_dim + 1, latent_dim) + self.calls = [] + + def forward(self, state, time, condition, *, condition_drop_mask): + self.calls.append(time.detach().clone()) + condition = torch.where(condition_drop_mask[:, None], 0, condition) + features = torch.cat( + ( + state, + time[:, None, None].expand(-1, state.shape[1], 1), + condition[:, None, :].expand(-1, state.shape[1], -1), + ), + -1, + ) + return self.linear(features) + + +def stages(num_levels=4, samples=3, dropout=0.0): + return [ + LikelihoodReferenceStage( + TimeDependentSequenceMean( + input_dim=2, + latent_dim=3, + horizon=2, + hidden_dim=8, + depth=1, + num_heads=2, + feedforward_dim=16, + ), + num_levels=num_levels, + interior_samples_per_content=samples, + ), + GaussianBridgeNoisingStage( + num_levels=num_levels, condition_dropout_probability=dropout + ), + ConditionalReverseMeanStage(TinyField(), num_levels=num_levels), + LikelihoodDecoderStage( + ContextFreeSequenceDecoder( + latent_dim=3, + output_dim=2, + horizon=2, + hidden_dim=8, + depth=1, + num_heads=2, + feedforward_dim=16, + ) + ), + GaussianBridgeObjectiveStage(num_levels=num_levels), + ] + + +def batch(): + return { + "target": torch.randn(2, 2, 2), + "condition": torch.randn(2, 2, requires_grad=True), + "sampler/noise": torch.randn(2, 2, 3), + } + + +def test_terminal_mean_is_exact_zero_and_codec_is_context_free(): + mean = TimeDependentSequenceMean() + action = torch.randn(2, 16, 4) + assert torch.count_nonzero(mean(action, torch.ones(2))) == 0 + assert mean(action, torch.full((2,), 1 / 32)).shape == (2, 16, 8) + with pytest.raises(TypeError): + mean(action, torch.ones(2), condition=torch.ones(2, 67)) + assert sum(p.numel() for p in mean.parameters()) < 15000 + + +def test_reference_noising_exact_targets_independent_samples_and_shared_drop_mask(): + torch.manual_seed(7) + modules = stages(samples=14, dropout=0.3) + result = modules[1](modules[0](batch())) + p = "likelihood/" + levels, noise = result[p + "levels"], result[p + "interior_noise"] + schedule = modules[1].schedule + assert levels.shape == (28,) and bool(((levels >= 2) & (levels <= 4)).all()) + assert not torch.equal(noise[0], noise[1]) + expected_state = ( + result[p + "current_mean"] + schedule.sigmas[levels - 1, None, None] * noise + ) + expected_target = ( + result[p + "previous_mean"] + + 0.95 * schedule.sigmas[levels - 2, None, None] * noise + ) + torch.testing.assert_close(result[p + "state"][:28], expected_state) + torch.testing.assert_close(result[p + "posterior_target"], expected_target) + torch.testing.assert_close( + result[p + "posterior_variance"], + (1 - 0.95**2) * schedule.sigmas[levels - 2].square(), + ) + masks = result[p + "condition_drop_mask"] + assert torch.equal(masks[:28].view(2, 14), masks[28:, None].expand(2, 14)) + assert result[p + "posterior_target"].requires_grad + + +def test_likelihood_reduction_sums_full_chunk_not_toy_coefficient(): + stage = GaussianBridgeObjectiveStage(num_levels=32, tau=0.02) + values = { + "likelihood/interior_prediction": torch.ones(28, 16, 8), + "likelihood/posterior_target": torch.zeros(28, 16, 8), + "likelihood/posterior_variance": torch.ones(28), + "likelihood/boundary_prediction": torch.ones(2, 16, 4), + "target": torch.zeros(2, 16, 4), + } + result = stage(values) + torch.testing.assert_close( + result["log/ActionFlow/InteriorBridgeNLL"], torch.tensor(31 * 128 / 2) + ) + torch.testing.assert_close( + result["log/ActionFlow/BoundaryNLL"], torch.tensor(80000.0) + ) + torch.testing.assert_close(result["log/MSE"], torch.tensor(1.0)) + assert result["loss/likelihood"] == result["log/ActionFlow/TotalLoss"] + assert not any("recon" in k.lower() or "fm" in k.lower() for k in result) + + +def test_both_likelihood_terms_update_private_means_and_shared_conditioned_field(): + torch.manual_seed(5) + modules = stages() + pipeline = Pipeline(modules) + value = batch() + result = pipeline(value) + mean = modules[0].mean_encoder.network.output_projection.weight + field = modules[2].field.linear.weight + decoder = modules[3].decoder.output_projection.weight + targets = (mean, field, decoder, value["condition"]) + interior = torch.autograd.grad( + result["log/ActionFlow/InteriorBridgeNLL"], + targets, + retain_graph=True, + allow_unused=True, + ) + boundary = torch.autograd.grad( + result["log/ActionFlow/BoundaryNLL"], + targets, + retain_graph=True, + allow_unused=True, + ) + assert interior[2] is None + assert all( + g is not None and torch.count_nonzero(g) > 0 + for i, g in enumerate(interior) + if i != 2 + ) + assert all(g is not None and torch.count_nonzero(g) > 0 for g in boundary) + assert len(modules[2].field.calls) == 1 + assert modules[2].field.calls[0].shape == (2 * (3 + 1),) + optimizer = torch.optim.AdamW(pipeline.parameters(), lr=1e-4) + optimizer.zero_grad() + result["loss/likelihood"].backward() + optimizer.step() + assert all(torch.isfinite(p).all() for p in pipeline.parameters()) + + +def test_inference_uses_all_discrete_calls_innovations_and_output_noise(monkeypatch): + modules = stages(num_levels=32) + pipeline = Pipeline(modules).eval() + shapes = [] + + def noise_like(value, **kwargs): + shapes.append(tuple(value.shape)) + return torch.ones_like(value, **kwargs) + + monkeypatch.setattr(torch, "randn_like", noise_like) + value = batch() + result = pipeline.execute( + {"condition": value["condition"], "sampler/noise": value["sampler/noise"]}, + mode="inference", + ) + calls = modules[2].field.calls + assert len(calls) == 32 + torch.testing.assert_close( + torch.stack(calls)[:, 0], torch.arange(32, 0, -1).float() / 32 + ) + assert shapes == [(2, 2, 3)] * 31 + [(2, 2, 2)] + assert result["likelihood/trajectory"].shape == (33, 2, 2, 3) + torch.testing.assert_close( + result["pred_action"], result["likelihood/action_mean"] + 0.02 + ) + assert result["log/ActionFlow/SamplerCalls"] == 32 + assert result["log/ActionFlow/LatentInnovations"] == 31 + assert "loss/likelihood" not in result and "likelihood/boundary_mean" not in result + + +def test_seeded_actual_sampler_replays_including_output_noise(): + pipeline = Pipeline(stages()).eval() + value = batch() + torch.manual_seed(42) + first = pipeline.execute(value, mode="inference")["pred_action"] + torch.manual_seed(42) + second = pipeline.execute(value, mode="inference")["pred_action"] + torch.testing.assert_close(first, second, rtol=0, atol=0) + torch.manual_seed(43) + assert not torch.equal( + first, pipeline.execute(value, mode="inference")["pred_action"] + ) + + +def test_wrapper_uses_likelihood_labels_and_records_observed_gradient_routes( + monkeypatch, +): + torch.manual_seed(3) + algo = PipelineAlgo(stages(), device="cpu") + model = ActionFlowLikelihoodModelWrapper( + pipeline=algo, enable_grad_norm=False, gradient_telemetry_cadence=1 + ) + recorded = {} + monkeypatch.setattr( + model, "log", lambda name, value, **kwargs: recorded.__setitem__(name, value) + ) + loss = model.training_step({"example": batch()}, 0) + assert torch.isfinite(loss) and loss.requires_grad + assert { + "Train/MSE", + "Train/ActionFlow/InteriorBridgeNLL", + "Train/ActionFlow/BoundaryNLL", + "Train/ActionFlow/TotalLoss", + } <= set(recorded) + assert not any( + "FlowMatchingLoss" in name or "ReconstructionLoss" in name for name in recorded + ) + checkpoint = {} + model.on_save_checkpoint(checkpoint) + manifest = checkpoint["action_flow_gradient_route_manifest"] + assert set(manifest["routes"]) == {"InteriorBridgeNLL", "BoundaryNLL"} + interior = [x["name"] for x in manifest["routes"]["InteriorBridgeNLL"]] + boundary = [x["name"] for x in manifest["routes"]["BoundaryNLL"]] + assert any("mean_encoder" in name for name in interior) + assert any("field" in name for name in interior) + assert not any("decoder" in name for name in interior) + assert any("decoder" in name for name in boundary) + core = {k: v for k, v in manifest.items() if k != "manifest_sha256"} + assert ( + hashlib.sha256( + json.dumps( + core, sort_keys=True, separators=(",", ":"), ensure_ascii=True + ).encode() + ).hexdigest() + == manifest["manifest_sha256"] + ) + assert checkpoint["action_flow_likelihood_contract"]["latent_fm_objective"] is False + + +def test_two_update_smoke_captures_gradient_routes_on_second_update(monkeypatch): + model = ActionFlowLikelihoodModelWrapper( + pipeline=PipelineAlgo(stages(), device="cpu"), + enable_grad_norm=False, + gradient_telemetry_cadence=2, + ) + recorded = {} + monkeypatch.setattr( + model, "log", lambda name, value, **kwargs: recorded.__setitem__(name, value) + ) + current_step = [0] + monkeypatch.setattr( + ActionFlowLikelihoodModelWrapper, + "global_step", + property(lambda self: current_step[0]), + ) + model.training_step({"example": batch()}, 0) + assert model._gradient_route_manifest is None + current_step[0] = 1 + loss = model.training_step({"example": batch()}, 1) + loss.backward() + model.on_after_backward() + model.on_before_optimizer_step(None) + assert model._gradient_route_manifest is not None + assert "Train/ActionFlow/GradientNorm/TotalPreclip" in recorded + assert recorded["Train/ActionFlow/Compute/PeakAllocatedBytes"] == 0 + + +def test_validation_step_emits_real_likelihood_and_preserves_evaluator_rng(monkeypatch): + algo = PipelineAlgo(stages(dropout=0.3), device="cpu") + calls, predictions, logged = [], [], {} + + class Evaluator: + def on_validation_start(self): + calls.append("start") + + def on_validation_step(self, processed, index, loader_index): + calls.append((index, loader_index)) + # Exercise the actual stochastic inference path delegated to the + # evaluator, which owns normalized/native MSE and EnergyScore32. + self.generated = algo.forward_eval(processed) + + def on_validation_end(self): + calls.append("end") + + evaluator = Evaluator() + wrapper = ActionFlowLikelihoodModelWrapper( + pipeline=algo, evaluator=evaluator, enable_grad_norm=False + ).eval() + real_forward = algo.forward_training + + def capture_forward(processed): + assert not torch.is_grad_enabled() and not wrapper.training + result = real_forward(processed) + predictions.append(result) + return result + + monkeypatch.setattr(algo, "forward_training", capture_forward) + monkeypatch.setattr( + wrapper, "log", lambda name, value, **kw: logged.update({name: (value, kw)}) + ) + value = {"first": batch(), "second": batch()} + with torch.random.fork_rng(), torch.inference_mode(): + torch.manual_seed(17) + expected = algo.forward_eval(value) + expected_rng = torch.random.get_rng_state() + torch.manual_seed(17) + wrapper.on_validation_start() + with torch.inference_mode(): + wrapper.validation_step({**value, "inactive": None}, 3, 1) + wrapper.on_validation_end() + assert calls == ["start", (3, 1), "end"] + assert len(predictions) == 1 + assert torch.equal(torch.random.get_rng_state(), expected_rng) + for source in value: + torch.testing.assert_close( + evaluator.generated[source]["pred_action"], + expected[source]["pred_action"], + rtol=0, + atol=0, + ) + assert "target" in value[source] # loss graph cannot consume caller input + assert set(logged) == { + "Valid/ActionFlow/InteriorBridgeNLL", + "Valid/ActionFlow/BoundaryNLL", + "Valid/ActionFlow/TotalLoss", + } + for name, (actual, options) in logged.items(): + key = "log/" + name.removeprefix("Valid/") + expected_loss = torch.stack([r[key] for r in predictions[0].values()]).mean() + torch.testing.assert_close(actual, expected_loss) + assert torch.isfinite(actual) and not actual.requires_grad + assert options == dict( + on_step=False, + on_epoch=True, + batch_size=4, + sync_dist=True, + add_dataloader_idx=False, + ) + assert all(parameter.grad is None for parameter in wrapper.parameters()) + + +def test_validation_step_rejects_nonfinite_likelihood_component(monkeypatch): + algo = PipelineAlgo(stages(), device="cpu") + wrapper = ActionFlowLikelihoodModelWrapper( + pipeline=algo, enable_grad_norm=False + ).eval() + real_forward = algo.forward_training + + def nonfinite_forward(value): + result = real_forward(value) + result["example"]["log/ActionFlow/BoundaryNLL"] = torch.tensor(float("nan")) + return result + + monkeypatch.setattr(algo, "forward_training", nonfinite_forward) + with pytest.raises(RuntimeError, match="Non-finite pipeline metric"): + wrapper.validation_step({"example": batch()}, 0) + + +def test_full_hydra_likelihood_graph_has_unblocked_inference_plan(): + config_dir = Path(__file__).parents[1] / "egomimic/hydra_configs" + with initialize_config_dir(version_base=None, config_dir=str(config_dir.resolve())): + cfg = compose( + config_name="train_zarr_cartesian", + overrides=[ + "+experiment=pusht/action_flow_bc_usocket_bridge_likelihood_s42", + "++paths.root_dir=.", + ], + ) + assert cfg.train.action_flow_method == cfg.model.action_flow_method + assert not cfg.evaluator.action_flow_diagnostics.enabled + algo = hydra.utils.instantiate(cfg.model.pipeline, device="cpu") + pipeline = algo.pipeline + runnable, excluded = pipeline.plan( + ["front_img_1", "state_agent_obj"], mode="inference" + ) + assert [type(s).__name__ for s in runnable] == [ + "FusedObsEncoder", + "GaussianLatentNoise", + "ConditionalReverseMeanStage", + "LikelihoodDecoderStage", + ] + assert all(missing == [""] for _, missing in excluded) + assert sum(p.numel() for p in pipeline.stages[0].parameters()) == 11197088 + assert sum(p.numel() for p in pipeline.stages[3].parameters()) == 10768 + assert sum(p.numel() for p in pipeline.stages[5].parameters()) == 39506641 + assert sum(p.numel() for p in pipeline.stages[6].parameters()) == 10744 + # PipelineAlgo is the orchestration adapter; registered parameters live in nets. + assert sum(p.numel() for p in algo.nets.parameters()) == 50725241 + + +@pytest.mark.parametrize( + "kwargs", [{"num_levels": 1}, {"rho": 1}, {"sigma_min": 0}, {"sigma_max": 2}] +) +def test_invalid_reference_schedule_fails_closed(kwargs): + with pytest.raises(ValueError): + GaussianBridgeSchedule(**kwargs) diff --git a/tests/test_action_flow_smoke_schedule.py b/tests/test_action_flow_smoke_schedule.py new file mode 100644 index 000000000..7a39dc1b7 --- /dev/null +++ b/tests/test_action_flow_smoke_schedule.py @@ -0,0 +1,140 @@ +"""Regressions for the actual Phoenix 12935484/12935488 verifier failures.""" + +from types import SimpleNamespace + +import pytest +import torch +from omegaconf import OmegaConf + +from test_verify_action_flow_training_smoke import MODULE, _checkpoint_payload +from tools.validate_action_flow_config import ( + GRAPH_METHOD, + STOPGRAD_METHOD, + compose_experiment, +) + +# Exact six-key schedules printed by the failed final verifiers. The jobs' +# training steps exited 0; both experiments intentionally train jointly at 0. +OBSERVED_SCHEDULES = ( + ( + "pusht/action_flow_bc_usocket_latent_fm_sg_recon1_s42", + STOPGRAD_METHOD, + { + "joint_objective_begins_at_global_step": 0, + "reconstruction_only_optimizer_steps": 0, + "joint_flow_weight": 1.0, + "joint_reconstruction_weight": 1.0, + "joint_action_velocity_weight": 1.0, + "schema_version": 1, + }, + ), + ( + "pusht/action_flow_bc_usocket_graph_section_s42", + GRAPH_METHOD, + { + "joint_objective_begins_at_global_step": 0, + "reconstruction_only_optimizer_steps": 0, + "joint_flow_weight": 1.0, + "joint_reconstruction_weight": 0.0, + "joint_action_velocity_weight": 1.0, + "schema_version": 1, + }, + ), +) + + +@pytest.mark.parametrize("experiment,method,schedule", OBSERVED_SCHEDULES) +@pytest.mark.parametrize("restored_warmup", (0, 1)) +def test_observed_joint_schedule_and_strict_reload( + tmp_path, monkeypatch, experiment, method, schedule, restored_warmup +): + config = compose_experiment(experiment) + checkpoint_dir = tmp_path / "checkpoints" + checkpoint_dir.mkdir() + payload = _checkpoint_payload() + payload["action_flow_loss_schedule"] = schedule + immutable = checkpoint_dir / "epoch-0-step-2.ckpt" + torch.save(payload, immutable) + (checkpoint_dir / "last.ckpt").symlink_to(immutable.name) + + class Wrapper: + reconstruction_only_warmup_steps = restored_warmup + encoder_e, field_v, decoder_g = object(), object(), object() + model = SimpleNamespace(pipeline=SimpleNamespace(stages=[])) + nets = SimpleNamespace(named_parameters=lambda **kwargs: []) + + @classmethod + def load_from_checkpoint(cls, *args, **kwargs): + assert kwargs["strict"] is True + return cls() + + def parameters(self): + return [SimpleNamespace(numel=lambda: MODULE.EXPECTED_PARAMETER_COUNT)] + + monkeypatch.setattr(MODULE, "ActionFlowModelWrapper", Wrapper) + monkeypatch.setattr(MODULE, "_validate_dimensions_and_modules", lambda *args: None) + monkeypatch.setattr(MODULE, "_validate_gradient_route_manifest", lambda *args: {}) + kwargs = dict( + reconstruction_weight=schedule["joint_reconstruction_weight"], + flow_weight=1.0, + method=method, + config=config, + ) + if restored_warmup: + with pytest.raises(MODULE.SmokeVerificationError, match="strict reload lost"): + MODULE._validate_checkpoint(tmp_path, **kwargs) + else: + result = MODULE._validate_checkpoint(tmp_path, **kwargs) + assert result["loss_schedule"] == schedule + assert result["strict_checkpoint_reload"] == "passed" + + +@pytest.mark.parametrize("key", tuple(OBSERVED_SCHEDULES[0][2])) +def test_every_schedule_field_remains_checked(key): + experiment, _, schedule = OBSERVED_SCHEDULES[0] + changed = {**schedule, key: schedule[key] + 1} + with pytest.raises(MODULE.SmokeVerificationError, match="unexpected.*schedule"): + MODULE._validate_checkpoint_loss_schedule( + changed, + compose_experiment(experiment), + reconstruction_weight=1, + flow_weight=1, + ) + + +@pytest.mark.parametrize("flow_weight", (1.0, 0.01)) +def test_actual_warmup_smoke_still_requires_configured_step_one(flow_weight): + suffix = "" if flow_weight == 1.0 else "_flow001" + config = compose_experiment( + f"pusht/action_flow_bc_usocket_recon10_warmup10k{suffix}_s42" + ) + # The maintained launcher reduces this experiment's full 10k warmup to one + # update for the real two-update warmup smoke. + OmegaConf.update(config, "model.reconstruction_only_warmup_steps", 1) + schedule = { + "joint_objective_begins_at_global_step": 1, + "reconstruction_only_optimizer_steps": 1, + "joint_flow_weight": flow_weight, + "joint_reconstruction_weight": 10.0, + "joint_action_velocity_weight": 1.0, + "schema_version": 1, + } + assert ( + MODULE._validate_checkpoint_loss_schedule( + schedule, config, reconstruction_weight=10, flow_weight=flow_weight + ) + == 1 + ) + schedule["joint_objective_begins_at_global_step"] = 0 + schedule["reconstruction_only_optimizer_steps"] = 0 + with pytest.raises(MODULE.SmokeVerificationError, match="unexpected.*schedule"): + MODULE._validate_checkpoint_loss_schedule( + schedule, config, reconstruction_weight=10, flow_weight=flow_weight + ) + + +def test_schedule_cannot_be_validated_without_resolved_config(): + with pytest.raises(MODULE.SmokeVerificationError, match="exact config"): + MODULE._validate_checkpoint_loss_schedule( + OBSERVED_SCHEDULES[0][2], None, reconstruction_weight=1, flow_weight=1 + ) diff --git a/tests/test_action_flow_stages.py b/tests/test_action_flow_stages.py index 0b79215f7..d36933d71 100644 --- a/tests/test_action_flow_stages.py +++ b/tests/test_action_flow_stages.py @@ -341,7 +341,6 @@ def test_stage_source_keeps_the_pipeline_boundary_generic(): "action_key", "egomimic.models", "unite", - ".detach(", "monotonic", "scale_loss", "noisy_reconstruction", diff --git a/tests/test_batch_signal_relay.py b/tests/test_batch_signal_relay.py new file mode 100644 index 000000000..9534a0689 --- /dev/null +++ b/tests/test_batch_signal_relay.py @@ -0,0 +1,155 @@ +"""Real bash/child regression for Slurm batch-shell-only signal delivery.""" + +from __future__ import annotations + +import json +import os +import shlex +import signal +import subprocess +import sys +import time +from pathlib import Path + +import pytest + +ROOT = Path(__file__).parents[1] +HELPER = ROOT / "scripts/ice/relay_batch_signals.sh" +CHILD = r""" +import json, os, pathlib, signal, sys, time +root = pathlib.Path(sys.argv[1]) +root.joinpath("started").write_text(str(os.getpid())) +if sys.argv[2] == "early": + raise SystemExit(int(sys.argv[3])) +time.sleep(float(sys.argv[4])) +def handle(received, frame): + root.joinpath("received").write_text(signal.Signals(received).name) + if received == signal.SIGUSR1 and sys.argv[2] != "boundary": + root.joinpath("COMPLETE.json").write_text("complete") + time.sleep(.05) + raise SystemExit(0 if received == signal.SIGUSR1 else 128 + received) +for received in (signal.SIGUSR1, signal.SIGTERM, signal.SIGINT): + signal.signal(received, handle) +pathlib.Path(os.environ["ICE_RUNNER_SIGNAL_READY_FILE"]).write_text(str(os.getpid())) +if sys.argv[2] == "normal": + raise SystemExit(int(sys.argv[3])) +time.sleep(5) +""" + + +def wait_file(path): + deadline = time.monotonic() + 3 + while time.monotonic() < deadline: + if path.exists() and path.stat().st_size: + return + time.sleep(0.01) + raise AssertionError(f"child readiness timeout: {path}") + + +def launch(root, mode="signal", exit_code=0, delay=0): + argv = [sys.executable, "-c", CHILD, str(root), mode, str(exit_code), str(delay)] + command = ( + f"set -euo pipefail; source {shlex.quote(str(HELPER))}; " + f"ice_run_with_signal_relay {shlex.quote(str(root / 'ready'))} " + f"{shlex.join(argv)}; " + f'if test "$ICE_BATCH_BOUNDARY_FORWARDED" = 1 && test ! -s {shlex.quote(str(root / "COMPLETE.json"))}; ' + "then printf 'REQUEUE_ACCEPTED\\n'; exit 0; fi; " + "printf 'POST_RUNNER_VERIFIER\\n'" + ) + return subprocess.Popen( + ["bash", "-c", command], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + +def assert_reaped(root): + pid = int((root / "started").read_text()) + with pytest.raises(ProcessLookupError): + os.kill(pid, 0) + + +@pytest.mark.parametrize("received", [signal.SIGUSR1, signal.SIGTERM, signal.SIGINT]) +@pytest.mark.parametrize("during_startup", [False, True]) +def test_exact_signal_reaches_child_and_wait_preserves_status( + tmp_path, received, during_startup +): + child = launch(tmp_path, delay=0.2 if during_startup else 0) + try: + wait_file(tmp_path / ("started" if during_startup else "ready")) + child.send_signal(received) + output, errors = child.communicate(timeout=4) + expected = 0 if received == signal.SIGUSR1 else 128 + received + assert child.returncode == expected, (output, errors) + assert (tmp_path / "received").read_text() == received.name + assert ("POST_RUNNER_VERIFIER" in output) == (expected == 0) + assert_reaped(tmp_path) + finally: + if child.poll() is None: + child.kill() + child.wait(timeout=3) + + +@pytest.mark.parametrize("mode", ["normal", "early"]) +@pytest.mark.parametrize("exit_code", [0, 7]) +def test_early_and_normal_exit_status_is_not_replaced_by_wait( + tmp_path, mode, exit_code +): + child = launch(tmp_path, mode=mode, exit_code=exit_code) + output, errors = child.communicate(timeout=4) + assert child.returncode == exit_code, (output, errors) + assert ("POST_RUNNER_VERIFIER" in output) == (exit_code == 0) + assert_reaped(tmp_path) + + +def test_existing_readiness_marker_refuses_before_launch(tmp_path): + (tmp_path / "ready").write_text("stale") + child = launch(tmp_path) + output, errors = child.communicate(timeout=3) + assert child.returncode == 73, (output, errors) + assert not (tmp_path / "started").exists() + assert "POST_RUNNER_VERIFIER" not in output + + +def test_successful_requeue_without_completion_does_not_run_final_verifier(tmp_path): + child = launch(tmp_path, mode="boundary") + wait_file(tmp_path / "ready") + child.send_signal(signal.SIGUSR1) + output, errors = child.communicate(timeout=4) + assert child.returncode == 0, (output, errors) + assert "REQUEUE_ACCEPTED" in output + assert "POST_RUNNER_VERIFIER" not in output + assert_reaped(tmp_path) + + +def test_actual_runner_publishes_readiness_before_command_validation(tmp_path): + ready = tmp_path / "ready.json" + environment = {**os.environ, "ICE_RUNNER_SIGNAL_READY_FILE": str(ready)} + result = subprocess.run( + [ + sys.executable, + str(ROOT / "scripts/ice/ice_requeue_runner.py"), + "--state-dir", + str(tmp_path / "state"), + "--checkpoint-glob", + str(tmp_path / "*.ckpt"), + "--checkpoint-signal", + "USR2", + "--checkpoint-forwarding", + "slurm-steps", + "--requeue-owner", + "runner", + ], + env=environment, + capture_output=True, + text=True, + timeout=3, + ) + assert result.returncode != 0 + payload = json.loads(ready.read_text()) + assert payload["signal_handlers_ready"] is True + assert payload["pid"] > 0 + launcher = (ROOT / "scripts/train/launch_action_flow_usocket.sbatch").read_text() + assert 'source "$AF_REPO/scripts/ice/relay_batch_signals.sh"' in launcher + assert 'ice_run_with_signal_relay "$ATTEMPT/runner-signal-ready.json"' in launcher diff --git a/tools/validate_action_flow_config.py b/tools/validate_action_flow_config.py index 4f816a4f3..6adce8f51 100644 --- a/tools/validate_action_flow_config.py +++ b/tools/validate_action_flow_config.py @@ -86,6 +86,195 @@ ConditionalVelocityStage, ContentDecoderStage, ) +LEGACY_METHOD = "action_flow_joint" +LIKELIHOOD_METHOD = "gaussian_bridge_likelihood" +GRAPH_METHOD = "graph_section_diagnostic" +STOPGRAD_METHOD = "latent_fm_stopgrad" +CANDIDATE_METHODS = { + "pusht/action_flow_bc_usocket_latent_fm_sg_recon1_s42": STOPGRAD_METHOD, + "pusht/action_flow_bc_usocket_bridge_likelihood_s42": LIKELIHOOD_METHOD, + "pusht/action_flow_bc_usocket_graph_section_s42": GRAPH_METHOD, +} +LIKELIHOOD_STAGE_TARGETS = EXPECTED_STAGE_TARGETS[:3] + tuple( + f"egomimic.pipeline.stages_action_flow_likelihood.{name}" + for name in ( + "LikelihoodReferenceStage", + "GaussianBridgeNoisingStage", + "ConditionalReverseMeanStage", + "LikelihoodDecoderStage", + "GaussianBridgeObjectiveStage", + ) +) + + +def action_flow_method(config: DictConfig, experiment: str | None = None) -> str: + """Fail closed on typed method drift; old named configs keep their contract.""" + method = str( + OmegaConf.select(config, "train.action_flow_method", default=LEGACY_METHOD) + ) + model_method = str( + OmegaConf.select(config, "model.action_flow_method", default=LEGACY_METHOD) + ) + _exact(model_method, method, "model/train action_flow_method") + _require( + method in {LEGACY_METHOD, *CANDIDATE_METHODS.values()}, + "unknown action_flow_method", + ) + if experiment is not None: + _exact( + method, + CANDIDATE_METHODS.get(experiment, LEGACY_METHOD), + "experiment action_flow_method", + ) + return method + + +def method_stage_targets(method: str) -> tuple[str, ...]: + return ( + LIKELIHOOD_STAGE_TARGETS + if method == LIKELIHOOD_METHOD + else EXPECTED_STAGE_TARGETS + ) + + +def method_wrapper_target(method: str) -> str: + if method == LIKELIHOOD_METHOD: + return "egomimic.pl_utils.pl_model_action_flow_likelihood.ActionFlowLikelihoodModelWrapper" + return "egomimic.pl_utils.pl_model_action_flow.ActionFlowModelWrapper" + + +def validate_method_contract(config: DictConfig, experiment: str | None = None) -> str: + """Small shared scientific gate used by CPU preflight and the real launcher.""" + method = action_flow_method(config, experiment) + _exact(str(config.model._target_), method_wrapper_target(method), "model wrapper") + stages = config.model.pipeline.stages + _exact( + tuple(str(stage._target_) for stage in stages), + method_stage_targets(method), + "stage topology", + ) + if method == STOPGRAD_METHOD: + _exact( + str(stages[5].flow_clean_gradient_mode), + "all_stopgrad", + "FM-only reference detachment", + ) + _float( + config.model.reconstruction_weight, 1.0, "candidate reconstruction weight" + ) + elif method == GRAPH_METHOD: + _exact( + str(config.model.pipeline._target_), + "egomimic.models.action_flow_graph.build_graph_section_pipeline", + "shared graph codec factory", + ) + _exact( + bool(config.model.pipeline._recursive_), + False, + "graph codec factory recursion", + ) + _float( + config.model.reconstruction_weight, + 0.0, + "graph diagnostic reconstruction weight", + ) + _exact( + bool(config.evaluator.action_flow_diagnostics.capture_activations), + False, + "graph diagnostic activation capture", + ) + elif method == LIKELIHOOD_METHOD: + for stage_index, key, expected in ( + (4, "num_levels", 32), + (3, "interior_samples_per_content", 14), + (4, "sigma_min", 0.1), + (4, "sigma_max", 1.0), + (4, "rho", 0.95), + (6, "tau", 0.02), + ): + _float( + stages[stage_index][key], + expected, + f"likelihood stage {stage_index}.{key}", + ) + _float(config.model[key], expected, f"likelihood model {key}") + _float( + config.run_provenance.objective[key], + expected, + f"likelihood provenance {key}", + ) + for index in (3, 5, 7): + _exact( + int(stages[index].num_levels), 32, f"likelihood stage {index} levels" + ) + for key, expected in (("sigma_min", 0.1), ("sigma_max", 1.0), ("rho", 0.95)): + _float(stages[5][key], expected, f"reverse-chain {key}") + _float(stages[7].tau, 0.02, "boundary likelihood output noise") + _exact(int(config.model.num_inference_steps), 32, "likelihood inference steps") + objective = config.run_provenance.objective + _exact( + str(objective.target_mean_gradients), + "attached", + "learned reference target gradients", + ) + _exact( + str(objective.reduction), + "sum_all_horizon_coordinates_then_batch_mean", + "likelihood coordinate reduction", + ) + _exact(int(objective.interior_level_multiplier), 31, "interior multiplier") + _float( + objective.boundary_per_coordinate_mse_multiplier, + 80000.0, + "boundary MSE multiplier", + ) + _exact( + str(config.run_provenance.inference.sampler), + "gaussian_bridge_reverse_chain", + "likelihood sampler", + ) + _exact( + int(config.run_provenance.inference.steps), 32, "likelihood sampler steps" + ) + _exact( + int(config.run_provenance.inference.latent_innovations), + 31, + "reverse-chain innovations", + ) + _float( + config.run_provenance.inference.output_noise_std, + 0.02, + "Gaussian action output noise", + ) + for key in ( + "flow_weight", + "reconstruction_weight", + "action_velocity_weight", + "reconstruction_only_warmup_steps", + ): + _require(key not in config.model, f"likelihood model must not carry {key}") + _exact( + bool(config.evaluator.action_flow_diagnostics.enabled), + False, + "ODE diagnostics are not likelihood diagnostics", + ) + if method != LEGACY_METHOD and method != LIKELIHOOD_METHOD: + _float(config.model.flow_weight, 1.0, "candidate flow weight") + _float( + stages[7].action_velocity_weight, 1.0, "candidate action velocity weight" + ) + _exact( + int( + OmegaConf.select( + config, "model.reconstruction_only_warmup_steps", default=0 + ) + ), + 0, + "candidate joint-from-step-zero objective", + ) + return method + + FORBIDDEN_PIPELINE_KEYS = frozenset( { "ac_keys", @@ -274,6 +463,7 @@ def _validate_dimensions_and_modules( _exact(int(config.model.latent_dim), 8, "model latent dimension") _exact(int(config.model.condition_dim), 67, "model condition dimension") + method = action_flow_method(config) observation = stages[0] noise = stages[1] encoder_stage = stages[3] @@ -281,12 +471,36 @@ def _validate_dimensions_and_modules( field_stage = stages[5] decoder_stage = stages[6] objective = stages[7] - _require(isinstance(encoder_stage.encoder, ContextFreeSequenceEncoder), "wrong E") - _require(isinstance(field_stage.field, AdaLNSequenceField), "wrong field v") - _require(isinstance(decoder_stage.decoder, ContextFreeSequenceDecoder), "wrong g") + if method == GRAPH_METHOD: + from egomimic.models.action_flow_codec import GraphSectionSequenceCodec - encoder = encoder_stage.encoder - decoder = decoder_stage.decoder + _require( + isinstance(decoder_stage.decoder, GraphSectionSequenceCodec), + "wrong graph codec", + ) + _require( + encoder_stage.encoder.graph is decoder_stage.decoder.graph, + "graph encoder/decoder must share the identical graph map", + ) + encoder, decoder = decoder_stage.decoder.graph, decoder_stage.decoder.residual + elif method == LIKELIHOOD_METHOD: + from egomimic.models.action_flow_likelihood import TimeDependentSequenceMean + + _require( + isinstance(encoder_stage.mean_encoder, TimeDependentSequenceMean), + "wrong reference mean", + ) + _exact( + tuple(inspect.signature(encoder_stage.mean_encoder.forward).parameters), + ("action", "time"), + "context-free learned reference signature", + ) + encoder, decoder = encoder_stage.mean_encoder.network, decoder_stage.decoder + else: + encoder, decoder = encoder_stage.encoder, decoder_stage.decoder + _require(isinstance(encoder, ContextFreeSequenceEncoder), "wrong E") + _require(isinstance(field_stage.field, AdaLNSequenceField), "wrong field v") + _require(isinstance(decoder, ContextFreeSequenceDecoder), "wrong g") field = field_stage.field codec_expected = { "horizon": 16, @@ -299,8 +513,16 @@ def _validate_dimensions_and_modules( for attribute, expected in codec_expected.items(): _exact(getattr(codec, attribute), expected, f"{label} {attribute}") _float(codec.dropout, 0.0, f"{label} dropout") - _exact(encoder.input_dim, 4, "encoder E input dimension") - _exact(encoder.output_dim, 8, "encoder E output dimension") + _exact( + encoder.input_dim, + 5 if method == LIKELIHOOD_METHOD else 4, + "encoder E input dimension", + ) + _exact( + encoder.output_dim, + 4 if method == GRAPH_METHOD else 8, + "encoder E output dimension", + ) _exact(decoder.input_dim, 8, "decoder g input dimension") _exact(decoder.output_dim, 4, "decoder g output dimension") _exact( @@ -356,22 +578,34 @@ def _validate_dimensions_and_modules( _exact(image_width, 64, "observation image-feature width") _exact(low_dim_width + image_width, 67, "observation condition width") - _exact(bridge.samples_per_content, 14, "bridge samples per content") + if method == LIKELIHOOD_METHOD: + _exact( + encoder_stage.interior_samples_per_content, + 14, + "interior samples per content", + ) + _exact( + int(config.model.num_inference_steps), 32, "reverse-chain field evaluations" + ) + else: + _exact(bridge.samples_per_content, 14, "bridge samples per content") + _exact(field_stage.num_inference_steps, 16, "inference field evaluations") _float(bridge.condition_dropout_probability, 0.3, "bridge condition dropout") - _exact(field_stage.num_inference_steps, 16, "inference field evaluations") - flow_weight = float(config.model.flow_weight) - _require(flow_weight in {0.01, 1.0}, "unsupported FM weight") - _float(objective.flow_weight, flow_weight, "FM weight") - _float(objective.action_velocity_weight, 1.0, "action-velocity weight") - _float( - objective.reconstruction_weight, - float(config.model.reconstruction_weight), - "reconstruction weight", - ) - _require( - float(config.model.reconstruction_weight) in {1.0, 10.0, 100.0}, - "reconstruction weight must be exactly 1, 10, or 100", - ) + if method != LIKELIHOOD_METHOD: + flow_weight = float(config.model.flow_weight) + _require(flow_weight in {0.01, 1.0}, "unsupported FM weight") + _float(objective.flow_weight, flow_weight, "FM weight") + _float(objective.action_velocity_weight, 1.0, "action-velocity weight") + _float( + objective.reconstruction_weight, + float(config.model.reconstruction_weight), + "reconstruction weight", + ) + _require( + float(config.model.reconstruction_weight) + in ({0.0} if method == GRAPH_METHOD else {1.0, 10.0, 100.0}), + "unsupported reconstruction weight", + ) parameters = { "observation_encoder": _parameter_manifest(observation), @@ -405,11 +639,12 @@ def _validate_topology( ) -> dict[str, Any]: stage_configs = tuple(config.model.pipeline.stages) stage_targets = tuple(str(stage._target_) for stage in stage_configs) - _exact(stage_targets, EXPECTED_STAGE_TARGETS, "configured stage topology") + expected_targets = method_stage_targets(action_flow_method(config)) + _exact(stage_targets, expected_targets, "configured stage topology") stages = tuple(pipeline_algo.pipeline.stages) _exact( - tuple(type(stage) for stage in stages), - EXPECTED_STAGE_TYPES, + tuple(f"{type(stage).__module__}.{type(stage).__name__}" for stage in stages), + expected_targets, "instantiated stage topology", ) @@ -431,24 +666,16 @@ def _validate_topology( not blocked_inference, f"inference graph has blocked stages: {blocked_inference}", ) - _exact(tuple(type(stage) for stage in train), EXPECTED_STAGE_TYPES, "train plan") + _exact(tuple(train), tuple(stages), "train plan") _exact( - tuple(type(stage) for stage in inference), - EXPECTED_INFERENCE_TYPES, + tuple(inference), + tuple(stages[index] for index in (0, 1, 5, 6)), "inference plan", ) - train_field = next( - stage for stage in train if isinstance(stage, ConditionalVelocityStage) - ) - inference_field = next( - stage for stage in inference if isinstance(stage, ConditionalVelocityStage) - ) - train_decoder = next( - stage for stage in train if isinstance(stage, ContentDecoderStage) - ) - inference_decoder = next( - stage for stage in inference if isinstance(stage, ContentDecoderStage) - ) + train_field = next(stage for stage in train if stage is stages[5]) + inference_field = next(stage for stage in inference if stage is stages[5]) + train_decoder = next(stage for stage in train if stage is stages[6]) + inference_decoder = next(stage for stage in inference if stage is stages[6]) _require(train_field is inference_field, "train/inference field stage was copied") _require( train_decoder is inference_decoder, "train/inference decoder stage was copied" @@ -711,38 +938,71 @@ def _validate_data_and_launch( "evaluator dataset aggregate content hash provenance", ) + method = action_flow_method(config) objective = provenance.objective - _float( - objective.flow_weight, - float(config.model.flow_weight), - "provenance FM weight", - ) - _float(objective.action_velocity_weight, 1.0, "provenance action weight") - _float( - objective.reconstruction_weight, - float(config.model.reconstruction_weight), - "provenance reconstruction weight", - ) - warmup_steps = int( - OmegaConf.select( - config, "model.reconstruction_only_warmup_steps", default=0 + if method == LIKELIHOOD_METHOD: + for key in ( + "num_levels", + "interior_samples_per_content", + "sigma_min", + "sigma_max", + "rho", + "tau", + ): + _float(objective[key], config.model[key], f"likelihood provenance {key}") + _exact( + str(objective.reduction), + "sum_all_horizon_coordinates_then_batch_mean", + "likelihood coordinate reduction", ) - ) - _require( - warmup_steps in {0, 10_000}, - "reconstruction-only warmup must be exactly 0 or 10000 steps", - ) - recorded_warmup = OmegaConf.select( - config, - "run_provenance.objective.reconstruction_only_warmup_steps", - default=0, - ) - _exact(int(recorded_warmup), warmup_steps, "provenance objective warmup") - _exact(int(objective.flow_samples_per_content), 14, "provenance bridge samples") - _float(objective.decoded_noise_scale_weight, 0.0, "decoded-noise scale weight") - _float(objective.monotonic_weight, 0.0, "monotonicity weight") - _exact(str(provenance.inference.sampler), "reverse_euler", "inference sampler") - _exact(int(provenance.inference.steps), 16, "inference sampler steps") + _exact( + int(objective.interior_level_multiplier), 31, "interior level multiplier" + ) + _float( + objective.boundary_per_coordinate_mse_multiplier, + 80000.0, + "boundary NLL coordinate multiplier", + ) + _exact( + str(provenance.inference.sampler), + "gaussian_bridge_reverse_chain", + "inference sampler", + ) + _exact(int(provenance.inference.steps), 32, "reverse-chain levels") + _exact(int(provenance.inference.latent_innovations), 31, "latent innovations") + _float(provenance.inference.output_noise_std, 0.02, "Gaussian output noise") + else: + _float( + objective.flow_weight, + float(config.model.flow_weight), + "provenance FM weight", + ) + _float(objective.action_velocity_weight, 1.0, "provenance action weight") + _float( + objective.reconstruction_weight, + float(config.model.reconstruction_weight), + "provenance reconstruction weight", + ) + warmup_steps = int( + OmegaConf.select( + config, "model.reconstruction_only_warmup_steps", default=0 + ) + ) + _require( + warmup_steps in {0, 10_000}, + "reconstruction-only warmup must be exactly 0 or 10000 steps", + ) + recorded_warmup = OmegaConf.select( + config, + "run_provenance.objective.reconstruction_only_warmup_steps", + default=0, + ) + _exact(int(recorded_warmup), warmup_steps, "provenance objective warmup") + _exact(int(objective.flow_samples_per_content), 14, "provenance bridge samples") + _float(objective.decoded_noise_scale_weight, 0.0, "decoded-noise scale weight") + _float(objective.monotonic_weight, 0.0, "monotonicity weight") + _exact(str(provenance.inference.sampler), "reverse_euler", "inference sampler") + _exact(int(provenance.inference.steps), 16, "inference sampler steps") _exact( bool(provenance.inference.classifier_free_guidance), False, @@ -802,103 +1062,114 @@ def _validate_data_and_launch( ) _exact(evaluator_distance, distance, "evaluator EnergyScore distance provenance") - diagnostics = config.evaluator.action_flow_diagnostics - _exact(bool(diagnostics.enabled), True, "Action Flow diagnostics enabled") - _exact( - [float(value) for value in diagnostics.raw_noise_levels], - [0.0, 0.25, 0.5, 0.75, 1.0], - "Action Flow diagnostic noise levels", - ) - _exact(int(diagnostics.max_batches_per_rank), 1, "diagnostic batch limit") - _exact(int(diagnostics.max_samples), 16, "diagnostic sample limit") - _exact(int(diagnostics.jacobian_samples), 2, "diagnostic Jacobian sample limit") - _exact(bool(diagnostics.capture_activations), True, "activation capture") - _exact( - { - int(key): int(value) - for key, value in diagnostics.activation_layer_map.items() - }, - {0: 0, 1: 11}, - "diagnostic activation layer map", - ) - _exact(int(diagnostics.cknna_k), 10, "diagnostic CKNNA k") - try: - native_error = normalize_usocket_native_error_config( - OmegaConf.to_container(diagnostics.native_error, resolve=True) + if action_flow_method(config) != LIKELIHOOD_METHOD: + diagnostics = config.evaluator.action_flow_diagnostics + _exact(bool(diagnostics.enabled), True, "Action Flow diagnostics enabled") + _exact( + [float(value) for value in diagnostics.raw_noise_levels], + [0.0, 0.25, 0.5, 0.75, 1.0], + "Action Flow diagnostic noise levels", + ) + _exact(int(diagnostics.max_batches_per_rank), 1, "diagnostic batch limit") + _exact(int(diagnostics.max_samples), 16, "diagnostic sample limit") + _exact(int(diagnostics.jacobian_samples), 2, "diagnostic Jacobian sample limit") + _exact( + bool(diagnostics.capture_activations), + action_flow_method(config) != GRAPH_METHOD, + "activation capture", + ) + _exact( + { + int(key): int(value) + for key, value in diagnostics.activation_layer_map.items() + }, + {} if action_flow_method(config) == GRAPH_METHOD else {0: 0, 1: 11}, + "diagnostic activation layer map", + ) + _exact( + int(diagnostics.cknna_k), + 0 if action_flow_method(config) == GRAPH_METHOD else 10, + "diagnostic CKNNA k", + ) + try: + native_error = normalize_usocket_native_error_config( + OmegaConf.to_container(diagnostics.native_error, resolve=True) + ) + except (TypeError, ValueError) as error: + raise PreflightError(str(error)) from error + _exact(native_error, USOCKET_NATIVE_ERROR_CONFIG, "diagnostic native error") + _exact( + diagnostics.provenance.source_commit, + provenance.source_commit, + "diagnostic source commit provenance", + ) + _exact( + diagnostics.provenance.normalization_sha256, + provenance.normalization_sha256, + "diagnostic normalization provenance", + ) + _exact( + str(diagnostics.provenance.split_manifest_sha256), + str(provenance.split_manifest_sha256), + "diagnostic split provenance", + ) + diagnostic_content = OmegaConf.to_container( + diagnostics.provenance.dataset_content, + resolve=True, + ) + _require( + isinstance(diagnostic_content, Mapping), + "diagnostic dataset-content provenance must be a mapping", + ) + _exact( + set(diagnostic_content), + {"aggregate_sha256", "manifest_sha256"}, + "diagnostic dataset-content provenance keys", + ) + _exact( + str(diagnostic_content["manifest_sha256"]).lower(), + configured_content_manifest_sha256, + "diagnostic dataset-content manifest provenance", + ) + _exact( + str(diagnostic_content["aggregate_sha256"]).lower(), + configured_aggregate_sha256, + "diagnostic dataset aggregate content provenance", + ) + _exact( + str(diagnostics.validation_view.split_manifest_sha256), + str(provenance.split_manifest_sha256), + "diagnostic validation-view split provenance", + ) + _exact( + int(diagnostics.validation_view.per_rank_batch_size), + 16, + "diagnostic validation batch size", + ) + _exact( + int(diagnostics.validation_view.world_size), + 1, + "diagnostic validation world size", + ) + _exact( + str(diagnostics.noise_seed_bank_sha256), + str(provenance.energy_score_contract.seed_bank_sha256), + "diagnostic seed-bank hash", + ) + seed_bank = ( + Path(config_root) / "evaluator" / "energy_score_seed_bank_k32_v1.json" + ) + _require(seed_bank.is_file(), f"diagnostic seed bank missing: {seed_bank}") + _exact( + _sha256(seed_bank), + str(diagnostics.noise_seed_bank_sha256), + "diagnostic seed-bank file hash", + ) + _exact( + str(diagnostics.provenance.decoder_jacobian_evaluation), + "declared_bridge_state_at_each_fixed_noise_level", + "diagnostic Jacobian evaluation point", ) - except (TypeError, ValueError) as error: - raise PreflightError(str(error)) from error - _exact(native_error, USOCKET_NATIVE_ERROR_CONFIG, "diagnostic native error") - _exact( - diagnostics.provenance.source_commit, - provenance.source_commit, - "diagnostic source commit provenance", - ) - _exact( - diagnostics.provenance.normalization_sha256, - provenance.normalization_sha256, - "diagnostic normalization provenance", - ) - _exact( - str(diagnostics.provenance.split_manifest_sha256), - str(provenance.split_manifest_sha256), - "diagnostic split provenance", - ) - diagnostic_content = OmegaConf.to_container( - diagnostics.provenance.dataset_content, - resolve=True, - ) - _require( - isinstance(diagnostic_content, Mapping), - "diagnostic dataset-content provenance must be a mapping", - ) - _exact( - set(diagnostic_content), - {"aggregate_sha256", "manifest_sha256"}, - "diagnostic dataset-content provenance keys", - ) - _exact( - str(diagnostic_content["manifest_sha256"]).lower(), - configured_content_manifest_sha256, - "diagnostic dataset-content manifest provenance", - ) - _exact( - str(diagnostic_content["aggregate_sha256"]).lower(), - configured_aggregate_sha256, - "diagnostic dataset aggregate content provenance", - ) - _exact( - str(diagnostics.validation_view.split_manifest_sha256), - str(provenance.split_manifest_sha256), - "diagnostic validation-view split provenance", - ) - _exact( - int(diagnostics.validation_view.per_rank_batch_size), - 16, - "diagnostic validation batch size", - ) - _exact( - int(diagnostics.validation_view.world_size), - 1, - "diagnostic validation world size", - ) - _exact( - str(diagnostics.noise_seed_bank_sha256), - str(provenance.energy_score_contract.seed_bank_sha256), - "diagnostic seed-bank hash", - ) - seed_bank = Path(config_root) / "evaluator" / "energy_score_seed_bank_k32_v1.json" - _require(seed_bank.is_file(), f"diagnostic seed bank missing: {seed_bank}") - _exact( - _sha256(seed_bank), - str(diagnostics.noise_seed_bank_sha256), - "diagnostic seed-bank file hash", - ) - _exact( - str(diagnostics.provenance.decoder_jacobian_evaluation), - "declared_bridge_state_at_each_fixed_noise_level", - "diagnostic Jacobian evaluation point", - ) launch = { "accumulate_grad_batches": 1, @@ -936,15 +1207,11 @@ def validate_config( """Validate one composed config and return report plus resolved payload.""" resolved, resolved_hash = resolved_config_payload(config) - _exact( - str(config.model._target_), - "egomimic.pl_utils.pl_model_action_flow.ActionFlowModelWrapper", - "model wrapper", - ) + method = validate_method_contract(config, experiment) configured_targets = tuple( str(stage._target_) for stage in config.model.pipeline.stages ) - _exact(configured_targets, EXPECTED_STAGE_TARGETS, "stage topology") + _exact(configured_targets, method_stage_targets(method), "stage topology") _require( "time_scale" in config.model.pipeline.stages[5].field, "field time_scale must be explicit", @@ -967,26 +1234,29 @@ def validate_config( ) _exact(accounted, parameters["pipeline_total"]["total"], "parameter accounting") - reconstruction_weight = float(config.model.reconstruction_weight) - report = { - "config_name": str(config.name), - "dimensions": dimensions, - "experiment": str(experiment), - "launch": launch, - "objective": { + objective_report = OmegaConf.to_container( + config.run_provenance.objective, resolve=True + ) + if method != LIKELIHOOD_METHOD: + objective_report = { "action_velocity_weight": 1.0, "condition_dropout_probability": 0.3, "flow_samples_per_content": 14, "flow_weight": float(config.model.flow_weight), - "reconstruction_weight": reconstruction_weight, + "reconstruction_weight": float(config.model.reconstruction_weight), "reconstruction_only_warmup_steps": int( OmegaConf.select( - config, - "model.reconstruction_only_warmup_steps", - default=0, + config, "model.reconstruction_only_warmup_steps", default=0 ) ), - }, + } + report = { + "config_name": str(config.name), + "dimensions": dimensions, + "experiment": str(experiment), + "launch": launch, + "action_flow_method": method, + "objective": objective_report, "optimization": optimization, "parameters": parameters, "resolved_config_sha256": resolved_hash,