From e04e3399b92d6eebd84f22270c065298f3893e9d Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Wed, 25 Mar 2026 19:26:41 +0100 Subject: [PATCH] fix normalizatiom --- src/lerobot/processor/normalize_processor.py | 13 +++++++++++++ src/lerobot/rl/actor.py | 10 +++++++++- src/lerobot/rl/learner.py | 12 ++++++++++++ 3 files changed, 34 insertions(+), 1 deletion(-) diff --git a/src/lerobot/processor/normalize_processor.py b/src/lerobot/processor/normalize_processor.py index 8a7a1176a..341ebfd45 100644 --- a/src/lerobot/processor/normalize_processor.py +++ b/src/lerobot/processor/normalize_processor.py @@ -131,6 +131,15 @@ class _NormalizationMixin: if self.dtype is None: self.dtype = torch.float32 self._tensor_stats = to_tensor(self.stats, device=self.device, dtype=self.dtype) + self._reshape_visual_stats() + + def _reshape_visual_stats(self) -> None: + """Reshape visual stats from ``[C]`` to ``[C, 1, 1]`` for image broadcasting.""" + for key, feature in self.features.items(): + if feature.type == FeatureType.VISUAL and key in self._tensor_stats: + for stat_name, stat_tensor in self._tensor_stats[key].items(): + if isinstance(stat_tensor, Tensor) and stat_tensor.ndim == 1: + self._tensor_stats[key][stat_name] = stat_tensor.reshape(-1, 1, 1) def to( self, device: torch.device | str | None = None, dtype: torch.dtype | None = None @@ -149,6 +158,7 @@ class _NormalizationMixin: if dtype is not None: self.dtype = dtype self._tensor_stats = to_tensor(self.stats, device=self.device, dtype=self.dtype) + self._reshape_visual_stats() return self def state_dict(self) -> dict[str, Tensor]: @@ -198,6 +208,7 @@ class _NormalizationMixin: # Don't load from state_dict, keep the explicitly provided stats # But ensure _tensor_stats is properly initialized self._tensor_stats = to_tensor(self.stats, device=self.device, dtype=self.dtype) # type: ignore[assignment] + self._reshape_visual_stats() return # Normal behavior: load stats from state_dict @@ -209,6 +220,8 @@ class _NormalizationMixin: dtype=torch.float32, device=self.device ) + self._reshape_visual_stats() + # Reconstruct the original stats dict from tensor stats for compatibility with to() method # and other functions that rely on self.stats self.stats = {} diff --git a/src/lerobot/rl/actor.py b/src/lerobot/rl/actor.py index 18c0ca1ea..9e4e69493 100644 --- a/src/lerobot/rl/actor.py +++ b/src/lerobot/rl/actor.py @@ -62,6 +62,7 @@ from lerobot.configs import parser from lerobot.configs.train import TrainRLServerPipelineConfig from lerobot.policies.factory import make_policy from lerobot.policies.sac.modeling_sac import SACPolicy +from lerobot.policies.sac.processor_sac import make_sac_pre_post_processors from lerobot.rl.process import ProcessSignalHandler from lerobot.rl.queue import get_last_item_from_queue from lerobot.robots import so_follower # noqa: F401 @@ -258,6 +259,11 @@ def act_with_policy( policy = policy.eval() assert isinstance(policy, nn.Module) + preprocessor, postprocessor = make_sac_pre_post_processors( + config=cfg.policy, + dataset_stats=cfg.policy.dataset_stats, + ) + obs, info = online_env.reset() env_processor.reset() action_processor.reset() @@ -289,7 +295,9 @@ def act_with_policy( # Time policy inference and check if it meets FPS requirement with policy_timer: # Extract observation from transition for policy - action = policy.select_action(batch=observation) + normalized_observation = preprocessor.process_observation(observation) + action = policy.select_action(batch=normalized_observation) + # action = postprocessor.process_action(action) policy_fps = policy_timer.fps_last log_policy_frequency_issue(policy_fps=policy_fps, cfg=cfg, interaction_step=interaction_step) diff --git a/src/lerobot/rl/learner.py b/src/lerobot/rl/learner.py index 2853fbcb3..96557b35a 100644 --- a/src/lerobot/rl/learner.py +++ b/src/lerobot/rl/learner.py @@ -66,6 +66,7 @@ from lerobot.datasets.factory import make_dataset from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.policies.factory import make_policy from lerobot.policies.sac.modeling_sac import SACPolicy +from lerobot.policies.sac.processor_sac import make_sac_pre_post_processors from lerobot.rl.buffer import ReplayBuffer, concatenate_batch_transitions from lerobot.rl.process import ProcessSignalHandler from lerobot.rl.wandb_utils import WandBLogger @@ -313,6 +314,11 @@ def add_actor_information_and_train( assert isinstance(policy, nn.Module) + preprocessor, _ = make_sac_pre_post_processors( + config=cfg.policy, + dataset_stats=cfg.policy.dataset_stats, + ) + policy.train() push_actor_policy_to_queue(parameters_queue=parameters_queue, policy=policy) @@ -408,6 +414,9 @@ def add_actor_information_and_train( done = batch["done"] check_nan_in_transition(observations=observations, actions=actions, next_state=next_observations) + observations = preprocessor.process_observation(observations) + next_observations = preprocessor.process_observation(next_observations) + observation_features, next_observation_features = get_observation_features( policy=policy, observations=observations, next_observations=next_observations ) @@ -467,6 +476,9 @@ def add_actor_information_and_train( check_nan_in_transition(observations=observations, actions=actions, next_state=next_observations) + observations = preprocessor.process_observation(observations) + next_observations = preprocessor.process_observation(next_observations) + observation_features, next_observation_features = get_observation_features( policy=policy, observations=observations, next_observations=next_observations )