From ba89c73b67f86cb63b3339f1ff235a8930e0d822 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Sun, 5 Jul 2026 13:10:00 +0000 Subject: [PATCH] fix(evo1): re-pad normalizer stats when loading from checkpoint reconcile_evo1_processors did not re-pad the (un)normalizer stats to max_state_dim/max_action_dim on the checkpoint-load path. When lerobot-train loads a checkpoint (e.g. stage2 from a stage1 checkpoint) it injects the raw dataset stats via processor overrides, so LIBERO's 8-dim state stats normalized a 24-dim padded state and crashed with "size of tensor a (24) must match tensor b (8)". Restore _refresh_evo1_normalization_steps (removed in the "remove legacy codepaths" refactor) and call it from reconcile_evo1_processors so the loaded stats/features are re-padded to EVO1's fixed widths. Padding is a no-op when stats are already at the target width. Co-authored-by: Cursor --- src/lerobot/policies/evo1/processor_evo1.py | 40 ++++++++++++++++++--- 1 file changed, 35 insertions(+), 5 deletions(-) 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,