From 21baa8fa023ec25e9043e310ae62bfda75a7ae99 Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Thu, 24 Jul 2025 10:40:03 +0200 Subject: [PATCH] refactor(factory): Remove unused imports and NaN detection hook from processor creation --- src/lerobot/policies/factory.py | 17 +---------------- 1 file changed, 1 insertion(+), 16 deletions(-) diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 490658a09..ea9e06a8e 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -16,7 +16,6 @@ import logging -import torch from torch import nn from lerobot.configs.policies import PreTrainedConfig @@ -35,7 +34,7 @@ from lerobot.policies.sac.reward_model.configuration_classifier import RewardCla from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig -from lerobot.processor.pipeline import EnvTransition, IdentityProcessor, RobotProcessor, TransitionIndex +from lerobot.processor.pipeline import IdentityProcessor, RobotProcessor def get_policy_class(name: str) -> PreTrainedPolicy: @@ -181,20 +180,6 @@ def make_processor( else: raise NotImplementedError(f"Processor for policy type '{policy_cfg.type}' is not implemented.") - # Helper hook function to detect NaNs in observation - def nan_detection_hook(step_idx: int, transition: EnvTransition) -> None: - observation = transition[TransitionIndex.OBSERVATION] - if observation is not None: - for key, value in observation.items(): - if isinstance(value, torch.Tensor) and torch.isnan(value).any(): - logging.warning(f"NaN detected in observation key '{key}' after step {step_idx}: {value}") - - # Attach the hook to all returned processors - if isinstance(processors, RobotProcessor): - processors = (processors,) # Wrap single processor in tuple for consistency - for processor in processors: - processor.register_after_step_hook(nan_detection_hook) - return processors