fix(g05): unnormalize queued actions by timestep

src/lerobot/policies/g05/processor_g05.py

@@ -469,6 +469,49 @@ class G05NormalizationClampStep(ProcessorStep):
         return {"minimum": self.minimum, "maximum": self.maximum}
 
 
+@dataclass
+@ProcessorStepRegistry.register(name="g05_stepwise_unnormalizer")
+class G05StepwiseUnnormalizerStep(UnnormalizerProcessorStep):
+    """Apply the checkpoint's action-timestep statistics to queued single actions."""
+
+    n_action_steps: int = 1
+    _action_index: int = field(default=0, init=False, repr=False)
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        action = transition.get(TransitionKey.ACTION)
+        stats = self._tensor_stats.get(ACTION)
+        if not isinstance(action, torch.Tensor) or not stats:
+            return super().__call__(transition)
+
+        first_stat = next(iter(stats.values()))
+        if first_stat.device != action.device or first_stat.dtype != action.dtype:
+            self.to(device=action.device, dtype=action.dtype)
+            stats = self._tensor_stats[ACTION]
+            first_stat = next(iter(stats.values()))
+
+        stats_steps = first_stat.shape[-2] if first_stat.ndim >= 2 else 1
+        is_single_queued_action = action.ndim <= 2 and action.shape[-2] != stats_steps
+        if not is_single_queued_action or stats_steps == 1:
+            return super().__call__(transition)
+
+        index = self._action_index % min(self.n_action_steps, stats_steps)
+        self._tensor_stats[ACTION] = {
+            name: value[index] if value.ndim >= 2 and value.shape[-2] == stats_steps else value
+            for name, value in stats.items()
+        }
+        try:
+            return super().__call__(transition)
+        finally:
+            self._tensor_stats[ACTION] = stats
+            self._action_index = (self._action_index + 1) % self.n_action_steps
+
+    def reset(self) -> None:
+        self._action_index = 0
+
+    def get_config(self) -> dict[str, Any]:
+        return {**super().get_config(), "n_action_steps": self.n_action_steps}
+
+
 @dataclass
 @ProcessorStepRegistry.register(name="g05_inverse_action_projection")
 class G05InverseActionProjectionStep(ProcessorStep):
@@ -708,11 +751,19 @@ def make_g05_pre_post_processors(
         to_transition=batch_to_transition,
         to_output=transition_to_batch,
     )
+    unnormalizer_cls = (
+        G05StepwiseUnnormalizerStep if config.use_stepwise_action_norm else UnnormalizerProcessorStep
+    )
+    unnormalizer_kwargs = {
+        "features": {ACTION: policy_features[ACTION]},
+        "norm_map": {FeatureType.ACTION: mode},
+        "stats": projected_stats,
+    }
+    if config.use_stepwise_action_norm:
+        unnormalizer_kwargs["n_action_steps"] = config.n_action_steps
     output_steps: list[ProcessorStep] = [
-        UnnormalizerProcessorStep(
+        unnormalizer_cls(
-            features={ACTION: policy_features[ACTION]},
+            **unnormalizer_kwargs,
-            norm_map={FeatureType.ACTION: mode},
-            stats=projected_stats,
         )
     ]
     if config.normalization_mode == "z_score_tail_mixed":

src/lerobot/scripts/lerobot_eval.py

@@ -220,8 +220,10 @@ def rollout(
     """
     assert isinstance(policy, nn.Module), "Policy must be a PyTorch nn module."
 
-    # Reset the policy and environments.
+    # Reset the policy, stateful processors, and environments.
     policy.reset()
+    for processor in (env_preprocessor, env_postprocessor, preprocessor, postprocessor):
+        processor.reset()
     observation, info = env.reset(seed=seeds)
     if render_callback is not None:
         render_callback(env)

tests/policies/g05/test_g05.py

@@ -606,7 +606,11 @@ def test_checkpoint_normalization_clips_to_author_finite_range():
 
 
 def test_stepwise_quantiles_constant_dimension_are_finite_and_serializable(tmp_path: Path):
-    config = _config(normalization_mode="q01_q99")
+    config = _config(
+        normalization_mode="q01_q99",
+        use_stepwise_action_norm=True,
+        n_action_steps=2,
+    )
     q01_action = torch.zeros(4, 7)
     q99_action = torch.ones(4, 7)
     q99_action[:, 2] = 0
@@ -627,6 +631,19 @@ def test_stepwise_quantiles_constant_dimension_are_finite_and_serializable(tmp_p
     assert torch.isfinite(processed[ACTION]).all()
     torch.testing.assert_close(postprocessor(processed[ACTION]), torch.zeros(4, 7))
 
+    step_q01 = torch.arange(4, dtype=torch.float32).view(4, 1).expand(4, 7)
+    step_stats = {
+        OBS_STATE: {"q01": torch.zeros(7), "q99": torch.ones(7)},
+        ACTION: {"q01": step_q01, "q99": step_q01 + 2},
+    }
+    _, stepwise_postprocessor = make_pre_post_processors(config, dataset_stats=step_stats)
+    normalized_action = torch.zeros(1, config.policy_action_dim)
+    torch.testing.assert_close(stepwise_postprocessor(normalized_action), torch.ones(1, 7))
+    torch.testing.assert_close(stepwise_postprocessor(normalized_action), torch.full((1, 7), 2.0))
+    torch.testing.assert_close(stepwise_postprocessor(normalized_action), torch.ones(1, 7))
+    stepwise_postprocessor.reset()
+    torch.testing.assert_close(stepwise_postprocessor(normalized_action), torch.ones(1, 7))
+
     preprocessor.save_pretrained(tmp_path)
     loaded = PolicyProcessorPipeline.from_pretrained(
         tmp_path, config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"