diff --git a/src/lerobot/policies/g05/processor_g05.py b/src/lerobot/policies/g05/processor_g05.py index 3c19642fe..e994b1d07 100644 --- a/src/lerobot/policies/g05/processor_g05.py +++ b/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( - features={ACTION: policy_features[ACTION]}, - norm_map={FeatureType.ACTION: mode}, - stats=projected_stats, + unnormalizer_cls( + **unnormalizer_kwargs, ) ] if config.normalization_mode == "z_score_tail_mixed": diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index b219f9bb1..08d7cb77d 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/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) diff --git a/tests/policies/g05/test_g05.py b/tests/policies/g05/test_g05.py index 866e92f11..63382e8f6 100644 --- a/tests/policies/g05/test_g05.py +++ b/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"