diff --git a/src/lerobot/policies/evo1/processor_evo1.py b/src/lerobot/policies/evo1/processor_evo1.py index adff8b66a..ba818f9db 100644 --- a/src/lerobot/policies/evo1/processor_evo1.py +++ b/src/lerobot/policies/evo1/processor_evo1.py @@ -302,6 +302,33 @@ def _pad_evo1_stats( return padded_stats +def _refresh_evo1_normalization_steps( + config: Evo1Config, + preprocessor: PolicyProcessorPipeline, + postprocessor: PolicyProcessorPipeline, +) -> None: + """Re-pad checkpoint-loaded (un)normalizer stats/features to EVO1's fixed widths. + + Loading a checkpoint injects the raw dataset stats (unpadded to max_state_dim/max_action_dim) + into the (un)normalizer via the generic override path in make_pre_post_processors. Those stats + and their declared features must be re-padded/reshaped to EVO1's fixed widths, otherwise + normalization fails against the padded state/action tensors (e.g. state padded to 24 vs. 8-dim + LIBERO stats). Padding is a no-op when stats are already at the target width. + """ + normalization_features = _evo1_normalization_features(config) + action_features = _evo1_action_features(config) + for step in preprocessor.steps: + if isinstance(step, NormalizerProcessorStep): + step.features = normalization_features + step.stats = _pad_evo1_stats(config, step.stats) + step.to(device=step.device, dtype=step.dtype) + for step in postprocessor.steps: + if isinstance(step, UnnormalizerProcessorStep): + step.features = action_features + step.stats = _pad_evo1_stats(config, step.stats) + step.to(device=step.device, dtype=step.dtype) + + def reconcile_evo1_processors( config: Evo1Config, preprocessor: PolicyProcessorPipeline, @@ -309,16 +336,19 @@ def reconcile_evo1_processors( ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: """Reconcile checkpoint-loaded pipelines with the current EVO1 config. - Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter - (converters are plain functions and are never serialized), and eval-time CLI overrides of the - action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This - restores the converter and rebuilds the action step from the current config so those overrides - take effect. + Three things cannot be restored from a serialized pipeline alone: the EVO1 batch converter + (converters are plain functions and are never serialized), eval-time CLI overrides of the + action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the + (un)normalizer stats/features when the generic override path injects raw, unpadded dataset + stats. This restores the converter, re-pads the normalization stats to EVO1's fixed widths, and + rebuilds the action step from the current config so those overrides take effect. """ # Pipelines reloaded from a checkpoint come back with the default batch converter, which drops # non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1. preprocessor.to_transition = evo1_batch_to_transition + _refresh_evo1_normalization_steps(config, preprocessor, postprocessor) + action_step = Evo1ActionProcessorStep( action_dim=_evo1_action_dim(config), binarize_gripper=config.binarize_gripper, diff --git a/tests/policies/evo1/test_evo1.py b/tests/policies/evo1/test_evo1.py index e9b9faf7d..4110657d5 100644 --- a/tests/policies/evo1/test_evo1.py +++ b/tests/policies/evo1/test_evo1.py @@ -496,6 +496,60 @@ def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path): assert "embodiment_id" in processed +def test_reconcile_evo1_processors_repads_overridden_stats(tmp_path): + """Loading a checkpoint and injecting raw (unpadded) dataset stats must be re-padded. + + Regression test: lerobot-train passes the raw dataset stats as normalizer/unnormalizer + overrides when resuming from a checkpoint (e.g. stage2 from a stage1 checkpoint). Those stats + are at the dataset dims (e.g. LIBERO state=8/action=7), but EVO1 pads state/action to + max_state_dim/max_action_dim before normalization, so reconcile_evo1_processors must re-pad the + stats or normalization crashes with a shape mismatch. + """ + config = make_config() + preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=make_stats()) + preprocessor.save_pretrained(tmp_path) + postprocessor.save_pretrained(tmp_path) + + # Reload with the generic override path injecting raw, unpadded dataset stats. + raw_stats = make_stats() + loaded_pre = PolicyProcessorPipeline.from_pretrained( + tmp_path, + config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json", + overrides={"normalizer_processor": {"stats": raw_stats}}, + to_transition=batch_to_transition, + to_output=transition_to_batch, + ) + loaded_post = PolicyProcessorPipeline.from_pretrained( + tmp_path, + config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json", + overrides={"unnormalizer_processor": {"stats": raw_stats}}, + to_transition=policy_action_to_transition, + to_output=transition_to_policy_action, + ) + + # Sanity: the override really injected unpadded stats before reconciliation. + normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep)) + assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (STATE_DIM,) + + loaded_pre, loaded_post = reconcile_evo1_processors(config, loaded_pre, loaded_post) + + normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep)) + unnormalizer = next(step for step in loaded_post.steps if isinstance(step, UnnormalizerProcessorStep)) + assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,) + assert normalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,) + assert unnormalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,) + + # Normalizing a padded state must not raise (this is the exact runtime path that crashed). + processed = loaded_pre( + { + "task": "pick the block", + OBS_STATE: torch.zeros(STATE_DIM), + f"{OBS_IMAGES}.front": torch.rand(3, 16, 16), + } + ) + assert processed[OBS_STATE].shape == (1, MAX_STATE_DIM) + + def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch): monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model) policy = modeling_evo1.Evo1Policy(make_config())