From c5371d0691d08cd3e3da591848ed865abc9bea24 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Fri, 17 Jul 2026 14:10:32 +0200 Subject: [PATCH] refactor(processors): share policy pipeline builders (#4016) * refactor(processors): share policy pipeline builders * Apply suggestions from code review Co-authored-by: Martino Russi <77496684+nepyope@users.noreply.github.com> Signed-off-by: Steven Palma * fix(processor): solve style after commit suggestions --------- Signed-off-by: Steven Palma Co-authored-by: Martino Russi <77496684+nepyope@users.noreply.github.com> --- src/lerobot/policies/act/processor_act.py | 41 +------ .../policies/diffusion/processor_diffusion.py | 39 +----- src/lerobot/policies/eo1/processor_eo1.py | 49 ++------ .../policies/fastwam/processor_fastwam.py | 49 ++------ .../processor_gaussian_actor.py | 40 +----- .../lingbot_va/processor_lingbot_va.py | 40 ++---- .../processor_multi_task_dit.py | 46 ++----- src/lerobot/policies/pi0/processor_pi0.py | 43 ++----- src/lerobot/policies/pi05/processor_pi05.py | 48 ++------ .../policies/pi0_fast/processor_pi0_fast.py | 48 ++------ .../policies/smolvla/processor_smolvla.py | 43 ++----- src/lerobot/policies/tdmpc/processor_tdmpc.py | 39 +----- .../policies/vla_jepa/processor_vla_jepa.py | 38 ++---- src/lerobot/policies/vqbet/processor_vqbet.py | 39 +----- .../policies/wall_x/processor_wall_x.py | 43 ++----- src/lerobot/policies/xvla/processor_xvla.py | 45 ++----- src/lerobot/processor/__init__.py | 8 ++ src/lerobot/processor/factory.py | 116 +++++++++++++++++- 18 files changed, 255 insertions(+), 559 deletions(-) diff --git a/src/lerobot/policies/act/processor_act.py b/src/lerobot/policies/act/processor_act.py index d87ade900..d3109aeb3 100644 --- a/src/lerobot/policies/act/processor_act.py +++ b/src/lerobot/policies/act/processor_act.py @@ -18,17 +18,10 @@ from typing import Any import torch from lerobot.processor import ( - AddBatchDimensionProcessorStep, - DeviceProcessorStep, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, - RenameObservationsProcessorStep, - UnnormalizerProcessorStep, - policy_action_to_transition, - transition_to_policy_action, + make_default_pre_post_processors, ) -from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME from .configuration_act import ACTConfig @@ -54,34 +47,4 @@ def make_act_pre_post_processors( tuple[PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], PolicyProcessorPipeline[PolicyAction, PolicyAction]]: A tuple containing the pre-processor pipeline and the post-processor pipeline. """ - - input_steps = [ - RenameObservationsProcessorStep(rename_map={}), - AddBatchDimensionProcessorStep(), - DeviceProcessorStep(device=config.device), - NormalizerProcessorStep( - features={**config.input_features, **config.output_features}, - norm_map=config.normalization_mapping, - stats=dataset_stats, - device=config.device, - ), - ] - output_steps = [ - UnnormalizerProcessorStep( - features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats - ), - DeviceProcessorStep(device="cpu"), - ] - - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + return make_default_pre_post_processors(config, dataset_stats, normalizer_device=config.device) diff --git a/src/lerobot/policies/diffusion/processor_diffusion.py b/src/lerobot/policies/diffusion/processor_diffusion.py index c4bc17680..28d77dc83 100644 --- a/src/lerobot/policies/diffusion/processor_diffusion.py +++ b/src/lerobot/policies/diffusion/processor_diffusion.py @@ -19,17 +19,10 @@ from typing import Any import torch from lerobot.processor import ( - AddBatchDimensionProcessorStep, - DeviceProcessorStep, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, - RenameObservationsProcessorStep, - UnnormalizerProcessorStep, - policy_action_to_transition, - transition_to_policy_action, + make_default_pre_post_processors, ) -from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME from .configuration_diffusion import DiffusionConfig @@ -63,32 +56,4 @@ def make_diffusion_pre_post_processors( Returns: A tuple containing the configured pre-processor and post-processor pipelines. """ - - input_steps = [ - RenameObservationsProcessorStep(rename_map={}), - AddBatchDimensionProcessorStep(), - DeviceProcessorStep(device=config.device), - NormalizerProcessorStep( - features={**config.input_features, **config.output_features}, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), - ] - output_steps = [ - UnnormalizerProcessorStep( - features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats - ), - DeviceProcessorStep(device="cpu"), - ] - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + return make_default_pre_post_processors(config, dataset_stats) diff --git a/src/lerobot/policies/eo1/processor_eo1.py b/src/lerobot/policies/eo1/processor_eo1.py index b1f32756a..1d3dc0515 100644 --- a/src/lerobot/policies/eo1/processor_eo1.py +++ b/src/lerobot/policies/eo1/processor_eo1.py @@ -23,24 +23,16 @@ import torch from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.processor import ( - AddBatchDimensionProcessorStep, ComplementaryDataProcessorStep, - DeviceProcessorStep, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry, - RenameObservationsProcessorStep, - UnnormalizerProcessorStep, + make_default_policy_processor_steps, + make_policy_processor_pipelines, ) -from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action from lerobot.types import TransitionKey -from lerobot.utils.constants import ( - OBS_STATE, - POLICY_POSTPROCESSOR_DEFAULT_NAME, - POLICY_PREPROCESSOR_DEFAULT_NAME, -) +from lerobot.utils.constants import OBS_STATE from lerobot.utils.import_utils import _transformers_available, require_package from .configuration_eo1 import EO1Config @@ -242,14 +234,12 @@ def make_eo1_pre_post_processors( ]: """Build pre/post processor pipelines for EO1.""" + steps = make_default_policy_processor_steps(config, dataset_stats) + input_steps: list[ProcessorStep] = [ - RenameObservationsProcessorStep(rename_map={}), - AddBatchDimensionProcessorStep(), - NormalizerProcessorStep( - features={**config.input_features, **config.output_features}, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), + steps.rename_observations, + steps.add_batch_dim, + steps.normalize, EO1ConversationTemplateStep(input_features=config.input_features, chunk_size=config.chunk_size), EO1QwenProcessorStep( processor_name=config.vlm_base, @@ -257,27 +247,12 @@ def make_eo1_pre_post_processors( image_max_pixels=config.image_max_pixels, use_fast_processor=config.use_fast_processor, ), - DeviceProcessorStep(device=config.device), + steps.to_device, ] output_steps: list[ProcessorStep] = [ - UnnormalizerProcessorStep( - features=config.output_features, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), - DeviceProcessorStep(device="cpu"), + steps.unnormalize, + steps.to_cpu, ] - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps) diff --git a/src/lerobot/policies/fastwam/processor_fastwam.py b/src/lerobot/policies/fastwam/processor_fastwam.py index 31f3b9277..318c7ff97 100644 --- a/src/lerobot/policies/fastwam/processor_fastwam.py +++ b/src/lerobot/policies/fastwam/processor_fastwam.py @@ -22,20 +22,11 @@ import torch from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.processor import ( ActionProcessorStep, - AddBatchDimensionProcessorStep, - DeviceProcessorStep, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, ProcessorStepRegistry, - RenameObservationsProcessorStep, - UnnormalizerProcessorStep, - policy_action_to_transition, - transition_to_policy_action, -) -from lerobot.utils.constants import ( - POLICY_POSTPROCESSOR_DEFAULT_NAME, - POLICY_PREPROCESSOR_DEFAULT_NAME, + make_default_policy_processor_steps, + make_policy_processor_pipelines, ) from .configuration_fastwam import FastWAMConfig @@ -105,38 +96,20 @@ def make_fastwam_pre_post_processors( # anyway) and unsafe across fine-tuning: its `resize_size` would be inherited from the base # checkpoint's camera geometry, not this dataset's, making the concatenation N_cameras x too wide. + steps = make_default_policy_processor_steps(config, normalization_stats, normalizer_device=config.device) + input_steps = [ - RenameObservationsProcessorStep(rename_map={}), - AddBatchDimensionProcessorStep(), - DeviceProcessorStep(device=config.device), - NormalizerProcessorStep( - features={**config.input_features, **config.output_features}, - norm_map=config.normalization_mapping, - stats=normalization_stats, - device=config.device, - ), + steps.rename_observations, + steps.add_batch_dim, + steps.to_device, + steps.normalize, ] output_steps = [ - UnnormalizerProcessorStep( - features=config.output_features, - norm_map=config.normalization_mapping, - stats=normalization_stats, - ), + steps.unnormalize, ] if config.toggle_action_dimensions: output_steps.append( FastWAMActionToggleProcessorStep(toggle_dimensions=config.toggle_action_dimensions) ) - output_steps.append(DeviceProcessorStep(device="cpu")) - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + output_steps.append(steps.to_cpu) + return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps) diff --git a/src/lerobot/policies/gaussian_actor/processor_gaussian_actor.py b/src/lerobot/policies/gaussian_actor/processor_gaussian_actor.py index 1e930d178..fede487ff 100644 --- a/src/lerobot/policies/gaussian_actor/processor_gaussian_actor.py +++ b/src/lerobot/policies/gaussian_actor/processor_gaussian_actor.py @@ -20,17 +20,10 @@ from typing import Any import torch from lerobot.processor import ( - AddBatchDimensionProcessorStep, - DeviceProcessorStep, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, - RenameObservationsProcessorStep, - UnnormalizerProcessorStep, - policy_action_to_transition, - transition_to_policy_action, + make_default_pre_post_processors, ) -from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME from .configuration_gaussian_actor import GaussianActorConfig @@ -62,33 +55,4 @@ def make_gaussian_actor_pre_post_processors( Returns: A tuple containing the configured pre-processor and post-processor pipelines. """ - - # Add remaining processors - input_steps = [ - RenameObservationsProcessorStep(rename_map={}), - AddBatchDimensionProcessorStep(), - DeviceProcessorStep(device=config.device), - NormalizerProcessorStep( - features={**config.input_features, **config.output_features}, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), - ] - output_steps = [ - UnnormalizerProcessorStep( - features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats - ), - DeviceProcessorStep(device="cpu"), - ] - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + return make_default_pre_post_processors(config, dataset_stats) diff --git a/src/lerobot/policies/lingbot_va/processor_lingbot_va.py b/src/lerobot/policies/lingbot_va/processor_lingbot_va.py index 119f77d5b..d183b8c34 100644 --- a/src/lerobot/policies/lingbot_va/processor_lingbot_va.py +++ b/src/lerobot/policies/lingbot_va/processor_lingbot_va.py @@ -25,19 +25,12 @@ import torch from lerobot.configs.types import FeatureType, NormalizationMode from lerobot.processor import ( - AddBatchDimensionProcessorStep, - DeviceProcessorStep, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, ProcessorStep, - RenameObservationsProcessorStep, UnnormalizerProcessorStep, -) -from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action -from lerobot.utils.constants import ( - POLICY_POSTPROCESSOR_DEFAULT_NAME, - POLICY_PREPROCESSOR_DEFAULT_NAME, + make_default_policy_processor_steps, + make_policy_processor_pipelines, ) from .configuration_lingbot_va import LingBotVAConfig @@ -52,15 +45,13 @@ def make_lingbot_va_pre_post_processors( ]: """Build the pre/post processor pipelines for LingBot-VA.""" + steps = make_default_policy_processor_steps(config, dataset_stats) + input_steps: list[ProcessorStep] = [ - RenameObservationsProcessorStep(rename_map={}), - AddBatchDimensionProcessorStep(), - NormalizerProcessorStep( - features={**config.input_features, **config.output_features}, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), - DeviceProcessorStep(device=config.device), + steps.rename_observations, + steps.add_batch_dim, + steps.normalize, + steps.to_device, ] # Unnormalize actions from [-1, 1] to physical units (QUANTILES) using q01/q99 restored from the checkpoint. @@ -70,18 +61,7 @@ def make_lingbot_va_pre_post_processors( norm_map={FeatureType.ACTION: NormalizationMode.QUANTILES}, stats=dataset_stats, ), - DeviceProcessorStep(device="cpu"), + steps.to_cpu, ] - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps) diff --git a/src/lerobot/policies/multi_task_dit/processor_multi_task_dit.py b/src/lerobot/policies/multi_task_dit/processor_multi_task_dit.py index 5f5b9994e..7cb49d4ea 100644 --- a/src/lerobot/policies/multi_task_dit/processor_multi_task_dit.py +++ b/src/lerobot/policies/multi_task_dit/processor_multi_task_dit.py @@ -19,18 +19,12 @@ from typing import Any import torch from lerobot.processor import ( - AddBatchDimensionProcessorStep, - DeviceProcessorStep, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, - RenameObservationsProcessorStep, TokenizerProcessorStep, - UnnormalizerProcessorStep, - policy_action_to_transition, - transition_to_policy_action, + make_default_policy_processor_steps, + make_policy_processor_pipelines, ) -from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME from .configuration_multi_task_dit import MultiTaskDiTConfig @@ -66,9 +60,11 @@ def make_multi_task_dit_pre_post_processors( A tuple containing the configured pre-processor and post-processor pipelines. """ + steps = make_default_policy_processor_steps(config, dataset_stats, normalizer_device=config.device) + input_steps = [ - RenameObservationsProcessorStep(rename_map={}), - AddBatchDimensionProcessorStep(), + steps.rename_observations, + steps.add_batch_dim, TokenizerProcessorStep( tokenizer_name=config.text_encoder_name, padding=config.tokenizer_padding, @@ -76,32 +72,12 @@ def make_multi_task_dit_pre_post_processors( max_length=config.tokenizer_max_length, truncation=config.tokenizer_truncation, ), - DeviceProcessorStep(device=config.device), - NormalizerProcessorStep( - features={**config.input_features, **config.output_features}, - norm_map=config.normalization_mapping, - stats=dataset_stats, - device=config.device, - ), + steps.to_device, + steps.normalize, ] output_steps = [ - UnnormalizerProcessorStep( - features=config.output_features, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), - DeviceProcessorStep(device="cpu"), + steps.unnormalize, + steps.to_cpu, ] - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps) diff --git a/src/lerobot/policies/pi0/processor_pi0.py b/src/lerobot/policies/pi0/processor_pi0.py index ad861f85b..ea6d662bd 100644 --- a/src/lerobot/policies/pi0/processor_pi0.py +++ b/src/lerobot/policies/pi0/processor_pi0.py @@ -21,22 +21,16 @@ import torch from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.processor import ( AbsoluteActionsProcessorStep, - AddBatchDimensionProcessorStep, ComplementaryDataProcessorStep, - DeviceProcessorStep, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry, RelativeActionsProcessorStep, - RenameObservationsProcessorStep, TokenizerProcessorStep, - UnnormalizerProcessorStep, - policy_action_to_transition, - transition_to_policy_action, + make_default_policy_processor_steps, + make_policy_processor_pipelines, ) -from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME from .configuration_pi0 import PI0Config @@ -136,10 +130,12 @@ def make_pi0_pre_post_processors( action_names=getattr(config, "action_feature_names", None), ) + steps = make_default_policy_processor_steps(config, dataset_stats) + # OpenPI order: raw → relative → normalize → model → unnormalize → absolute input_steps: list[ProcessorStep] = [ - RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one - AddBatchDimensionProcessorStep(), + steps.rename_observations, # To mimic the same processor as pretrained one + steps.add_batch_dim, Pi0NewLineProcessor(), # Add newlines before tokenization for PaliGemma TokenizerProcessorStep( tokenizer_name="google/paligemma-3b-pt-224", @@ -147,32 +143,15 @@ def make_pi0_pre_post_processors( padding_side="right", padding="max_length", ), - DeviceProcessorStep(device=config.device), + steps.to_device, relative_step, - NormalizerProcessorStep( - features={**config.input_features, **config.output_features}, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), + steps.normalize, ] output_steps: list[ProcessorStep] = [ - UnnormalizerProcessorStep( - features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats - ), + steps.unnormalize, AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step), - DeviceProcessorStep(device="cpu"), + steps.to_cpu, ] - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps) diff --git a/src/lerobot/policies/pi05/processor_pi05.py b/src/lerobot/policies/pi05/processor_pi05.py index 2d015b24f..ac16269a0 100644 --- a/src/lerobot/policies/pi05/processor_pi05.py +++ b/src/lerobot/policies/pi05/processor_pi05.py @@ -24,26 +24,17 @@ import torch from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.processor import ( AbsoluteActionsProcessorStep, - AddBatchDimensionProcessorStep, - DeviceProcessorStep, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry, RelativeActionsProcessorStep, - RenameObservationsProcessorStep, TokenizerProcessorStep, - UnnormalizerProcessorStep, - policy_action_to_transition, - transition_to_policy_action, + make_default_policy_processor_steps, + make_policy_processor_pipelines, ) from lerobot.types import EnvTransition, TransitionKey -from lerobot.utils.constants import ( - OBS_STATE, - POLICY_POSTPROCESSOR_DEFAULT_NAME, - POLICY_PREPROCESSOR_DEFAULT_NAME, -) +from lerobot.utils.constants import OBS_STATE from .configuration_pi05 import PI05Config @@ -135,18 +126,16 @@ def make_pi05_pre_post_processors( action_names=getattr(config, "action_feature_names", None), ) + steps = make_default_policy_processor_steps(config, dataset_stats) + # OpenPI order: raw → relative → normalize → model → unnormalize → absolute input_steps: list[ProcessorStep] = [ - RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one - AddBatchDimensionProcessorStep(), + steps.rename_observations, # To mimic the same processor as pretrained one + steps.add_batch_dim, relative_step, # NOTE: NormalizerProcessorStep MUST come before Pi05PrepareStateTokenizerProcessorStep # because the tokenizer step expects normalized state in [-1, 1] range for discretization - NormalizerProcessorStep( - features={**config.input_features, **config.output_features}, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), + steps.normalize, Pi05PrepareStateTokenizerProcessorStep(max_state_dim=config.max_state_dim), TokenizerProcessorStep( tokenizer_name="google/paligemma-3b-pt-224", @@ -154,26 +143,13 @@ def make_pi05_pre_post_processors( padding_side="right", padding="max_length", ), - DeviceProcessorStep(device=config.device), + steps.to_device, ] output_steps: list[ProcessorStep] = [ - UnnormalizerProcessorStep( - features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats - ), + steps.unnormalize, AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step), - DeviceProcessorStep(device="cpu"), + steps.to_cpu, ] - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps) diff --git a/src/lerobot/policies/pi0_fast/processor_pi0_fast.py b/src/lerobot/policies/pi0_fast/processor_pi0_fast.py index 60a519786..fdf759e29 100644 --- a/src/lerobot/policies/pi0_fast/processor_pi0_fast.py +++ b/src/lerobot/policies/pi0_fast/processor_pi0_fast.py @@ -25,26 +25,17 @@ from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.processor import ( AbsoluteActionsProcessorStep, ActionTokenizerProcessorStep, - AddBatchDimensionProcessorStep, - DeviceProcessorStep, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry, RelativeActionsProcessorStep, - RenameObservationsProcessorStep, TokenizerProcessorStep, - UnnormalizerProcessorStep, - policy_action_to_transition, - transition_to_policy_action, + make_default_policy_processor_steps, + make_policy_processor_pipelines, ) from lerobot.types import EnvTransition, TransitionKey -from lerobot.utils.constants import ( - OBS_STATE, - POLICY_POSTPROCESSOR_DEFAULT_NAME, - POLICY_PREPROCESSOR_DEFAULT_NAME, -) +from lerobot.utils.constants import OBS_STATE from .configuration_pi0_fast import PI0FastConfig @@ -135,6 +126,8 @@ def make_pi0_fast_pre_post_processors( action_names=getattr(config, "action_feature_names", None), ) + steps = make_default_policy_processor_steps(config, dataset_stats) + # Pi0Fast order: relative → normalize → tokenize → model → unnormalize → absolute # This matches pi0/pi0.5: RelativeActionsProcessorStep runs first on raw absolute actions, # caching the raw state. NormalizerProcessorStep then normalizes the raw relative actions, @@ -144,14 +137,10 @@ def make_pi0_fast_pre_post_processors( # before Pi0FastPrepareStateAndLanguageTokenizerProcessorStep, so the state tokenizer # continues to receive normalized state in [-1, 1] as expected. input_steps: list[ProcessorStep] = [ - RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one - AddBatchDimensionProcessorStep(), + steps.rename_observations, # To mimic the same processor as pretrained one + steps.add_batch_dim, relative_step, - NormalizerProcessorStep( - features={**config.input_features, **config.output_features}, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), + steps.normalize, Pi0FastPrepareStateAndLanguageTokenizerProcessorStep(max_state_dim=config.max_state_dim), TokenizerProcessorStep( tokenizer_name=config.text_tokenizer_name, @@ -165,26 +154,13 @@ def make_pi0_fast_pre_post_processors( fast_skip_tokens=config.fast_skip_tokens, paligemma_tokenizer_name=config.text_tokenizer_name, ), - DeviceProcessorStep(device=config.device), + steps.to_device, ] output_steps: list[ProcessorStep] = [ - UnnormalizerProcessorStep( - features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats - ), + steps.unnormalize, AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step), - DeviceProcessorStep(device="cpu"), + steps.to_cpu, ] - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps) diff --git a/src/lerobot/policies/smolvla/processor_smolvla.py b/src/lerobot/policies/smolvla/processor_smolvla.py index 8d6c8aca4..e924780cd 100644 --- a/src/lerobot/policies/smolvla/processor_smolvla.py +++ b/src/lerobot/policies/smolvla/processor_smolvla.py @@ -19,19 +19,13 @@ from typing import Any import torch from lerobot.processor import ( - AddBatchDimensionProcessorStep, - DeviceProcessorStep, NewLineTaskProcessorStep, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, - RenameObservationsProcessorStep, TokenizerProcessorStep, - UnnormalizerProcessorStep, - policy_action_to_transition, - transition_to_policy_action, + make_default_policy_processor_steps, + make_policy_processor_pipelines, ) -from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME from .configuration_smolvla import SmolVLAConfig @@ -66,9 +60,11 @@ def make_smolvla_pre_post_processors( A tuple containing the configured pre-processor and post-processor pipelines. """ + steps = make_default_policy_processor_steps(config, dataset_stats) + input_steps = [ - RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one - AddBatchDimensionProcessorStep(), + steps.rename_observations, # To mimic the same processor as pretrained one + steps.add_batch_dim, NewLineTaskProcessorStep(), TokenizerProcessorStep( tokenizer_name=config.vlm_model_name, @@ -76,28 +72,11 @@ def make_smolvla_pre_post_processors( padding_side="right", max_length=config.tokenizer_max_length, ), - DeviceProcessorStep(device=config.device), - NormalizerProcessorStep( - features={**config.input_features, **config.output_features}, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), + steps.to_device, + steps.normalize, ] output_steps = [ - UnnormalizerProcessorStep( - features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats - ), - DeviceProcessorStep(device="cpu"), + steps.unnormalize, + steps.to_cpu, ] - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps) diff --git a/src/lerobot/policies/tdmpc/processor_tdmpc.py b/src/lerobot/policies/tdmpc/processor_tdmpc.py index 7afe956dc..08756e819 100644 --- a/src/lerobot/policies/tdmpc/processor_tdmpc.py +++ b/src/lerobot/policies/tdmpc/processor_tdmpc.py @@ -19,17 +19,10 @@ from typing import Any import torch from lerobot.processor import ( - AddBatchDimensionProcessorStep, - DeviceProcessorStep, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, - RenameObservationsProcessorStep, - UnnormalizerProcessorStep, - policy_action_to_transition, - transition_to_policy_action, + make_default_pre_post_processors, ) -from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME from .configuration_tdmpc import TDMPCConfig @@ -61,32 +54,4 @@ def make_tdmpc_pre_post_processors( Returns: A tuple containing the configured pre-processor and post-processor pipelines. """ - - input_steps = [ - RenameObservationsProcessorStep(rename_map={}), - AddBatchDimensionProcessorStep(), - DeviceProcessorStep(device=config.device), - NormalizerProcessorStep( - features={**config.input_features, **config.output_features}, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), - ] - output_steps = [ - UnnormalizerProcessorStep( - features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats - ), - DeviceProcessorStep(device="cpu"), - ] - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + return make_default_pre_post_processors(config, dataset_stats) diff --git a/src/lerobot/policies/vla_jepa/processor_vla_jepa.py b/src/lerobot/policies/vla_jepa/processor_vla_jepa.py index b59cc0e90..aed2b754b 100644 --- a/src/lerobot/policies/vla_jepa/processor_vla_jepa.py +++ b/src/lerobot/policies/vla_jepa/processor_vla_jepa.py @@ -20,20 +20,16 @@ import torch from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig from lerobot.processor import ( - AddBatchDimensionProcessorStep, - DeviceProcessorStep, EnvTransition, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry, - RenameObservationsProcessorStep, TransitionKey, UnnormalizerProcessorStep, + make_default_policy_processor_steps, + make_policy_processor_pipelines, ) -from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action -from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME @ProcessorStepRegistry.register(name="vla_jepa_clip_actions") @@ -112,15 +108,12 @@ def make_vla_jepa_pre_post_processors( PolicyProcessorPipeline[PolicyAction, PolicyAction], ]: features = {**config.input_features, **config.output_features} + steps = make_default_policy_processor_steps(config, dataset_stats) input_steps = [ - RenameObservationsProcessorStep(rename_map={}), - AddBatchDimensionProcessorStep(), - DeviceProcessorStep(device=config.device), - NormalizerProcessorStep( - features=features, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), + steps.rename_observations, + steps.add_batch_dim, + steps.to_device, + steps.normalize, ] output_steps: list[ProcessorStep] = [] if config.clip_normalized_actions: @@ -129,6 +122,8 @@ def make_vla_jepa_pre_post_processors( output_steps.append( PreSnapGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold) ) + # NOTE: unlike the default policy unnormalizer (output features only), VLA-JEPA + # unnormalizes over BOTH input and output features. output_steps.append( UnnormalizerProcessorStep( features=features, @@ -140,16 +135,5 @@ def make_vla_jepa_pre_post_processors( output_steps.append( BinarizeGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold) ) - output_steps.append(DeviceProcessorStep(device="cpu")) - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + output_steps.append(steps.to_cpu) + return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps) diff --git a/src/lerobot/policies/vqbet/processor_vqbet.py b/src/lerobot/policies/vqbet/processor_vqbet.py index f7b6a061e..68182db3f 100644 --- a/src/lerobot/policies/vqbet/processor_vqbet.py +++ b/src/lerobot/policies/vqbet/processor_vqbet.py @@ -20,17 +20,10 @@ from typing import Any import torch from lerobot.processor import ( - AddBatchDimensionProcessorStep, - DeviceProcessorStep, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, - RenameObservationsProcessorStep, - UnnormalizerProcessorStep, - policy_action_to_transition, - transition_to_policy_action, + make_default_pre_post_processors, ) -from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME from .configuration_vqbet import VQBeTConfig @@ -62,32 +55,4 @@ def make_vqbet_pre_post_processors( Returns: A tuple containing the configured pre-processor and post-processor pipelines. """ - - input_steps = [ - RenameObservationsProcessorStep(rename_map={}), # Let the possibility to the user to rename the keys - AddBatchDimensionProcessorStep(), - DeviceProcessorStep(device=config.device), - NormalizerProcessorStep( - features={**config.input_features, **config.output_features}, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), - ] - output_steps = [ - UnnormalizerProcessorStep( - features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats - ), - DeviceProcessorStep(device="cpu"), - ] - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + return make_default_pre_post_processors(config, dataset_stats) diff --git a/src/lerobot/policies/wall_x/processor_wall_x.py b/src/lerobot/policies/wall_x/processor_wall_x.py index 069cef5d6..e8cec6ccb 100644 --- a/src/lerobot/policies/wall_x/processor_wall_x.py +++ b/src/lerobot/policies/wall_x/processor_wall_x.py @@ -20,19 +20,13 @@ import torch from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.processor import ( - AddBatchDimensionProcessorStep, ComplementaryDataProcessorStep, - DeviceProcessorStep, - NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, ProcessorStepRegistry, - RenameObservationsProcessorStep, - UnnormalizerProcessorStep, - policy_action_to_transition, - transition_to_policy_action, + make_default_policy_processor_steps, + make_policy_processor_pipelines, ) -from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME from .configuration_wall_x import WallXConfig @@ -65,37 +59,22 @@ def make_wall_x_pre_post_processors( A tuple containing the configured pre-processor and post-processor pipelines """ + steps = make_default_policy_processor_steps(config, dataset_stats) + input_steps = [ - RenameObservationsProcessorStep(rename_map={}), - AddBatchDimensionProcessorStep(), + steps.rename_observations, + steps.add_batch_dim, WallXTaskProcessor(), # Process task description - NormalizerProcessorStep( - features={**config.input_features, **config.output_features}, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), - DeviceProcessorStep(device=config.device), + steps.normalize, + steps.to_device, ] output_steps = [ - UnnormalizerProcessorStep( - features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats - ), - DeviceProcessorStep(device="cpu"), + steps.unnormalize, + steps.to_cpu, ] - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps) @ProcessorStepRegistry.register(name="wall_x_task_processor") diff --git a/src/lerobot/policies/xvla/processor_xvla.py b/src/lerobot/policies/xvla/processor_xvla.py index 0336ec722..02c54f9c9 100644 --- a/src/lerobot/policies/xvla/processor_xvla.py +++ b/src/lerobot/policies/xvla/processor_xvla.py @@ -22,19 +22,14 @@ import torch from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.processor import ( - AddBatchDimensionProcessorStep, - DeviceProcessorStep, - NormalizerProcessorStep, ObservationProcessorStep, PolicyAction, PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry, - RenameObservationsProcessorStep, TokenizerProcessorStep, - UnnormalizerProcessorStep, - policy_action_to_transition, - transition_to_policy_action, + make_default_policy_processor_steps, + make_policy_processor_pipelines, ) from lerobot.types import EnvTransition, TransitionKey from lerobot.utils.constants import ( @@ -42,8 +37,6 @@ from lerobot.utils.constants import ( OBS_IMAGES, OBS_PREFIX, OBS_STATE, - POLICY_POSTPROCESSOR_DEFAULT_NAME, - POLICY_PREPROCESSOR_DEFAULT_NAME, ) from .configuration_xvla import XVLAConfig @@ -61,10 +54,11 @@ def make_xvla_pre_post_processors( Build the LeRobot processor pipelines for XVLA. """ - features = {**config.input_features, **config.output_features} + steps = make_default_policy_processor_steps(config, dataset_stats) + input_steps = [ - RenameObservationsProcessorStep(rename_map={}), - AddBatchDimensionProcessorStep(), + steps.rename_observations, + steps.add_batch_dim, TokenizerProcessorStep( tokenizer_name=config.tokenizer_name, max_length=config.tokenizer_max_length, @@ -74,32 +68,15 @@ def make_xvla_pre_post_processors( XVLAImageToFloatProcessorStep(), XVLAImageNetNormalizeProcessorStep(), XVLAAddDomainIdProcessorStep(), - DeviceProcessorStep(device=config.device), - NormalizerProcessorStep( - features=features, norm_map=config.normalization_mapping, stats=dataset_stats - ), + steps.to_device, + steps.normalize, ] output_steps = [ - UnnormalizerProcessorStep( - features=config.output_features, - norm_map=config.normalization_mapping, - stats=dataset_stats, - ), - DeviceProcessorStep(device="cpu"), + steps.unnormalize, + steps.to_cpu, ] - return ( - PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( - steps=input_steps, - name=POLICY_PREPROCESSOR_DEFAULT_NAME, - ), - PolicyProcessorPipeline[PolicyAction, PolicyAction]( - steps=output_steps, - name=POLICY_POSTPROCESSOR_DEFAULT_NAME, - to_transition=policy_action_to_transition, - to_output=transition_to_policy_action, - ), - ) + return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps) # Custom XVLA processor steps diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index fe35af4b4..79bb44c85 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -42,10 +42,14 @@ from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorTo from .device_processor import DeviceProcessorStep from .env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep from .factory import ( + DefaultPolicyProcessorSteps, + make_default_policy_processor_steps, + make_default_pre_post_processors, make_default_processors, make_default_robot_action_processor, make_default_robot_observation_processor, make_default_teleop_action_processor, + make_policy_processor_pipelines, ) from .gym_action_processor import ( Numpy2TorchActionProcessorStep, @@ -129,10 +133,14 @@ __all__ = [ "ImageCropResizeProcessorStep", "InfoProcessorStep", "InterventionActionProcessorStep", + "DefaultPolicyProcessorSteps", + "make_default_policy_processor_steps", + "make_default_pre_post_processors", "make_default_processors", "make_default_teleop_action_processor", "make_default_robot_action_processor", "make_default_robot_observation_processor", + "make_policy_processor_pipelines", "AbsoluteActionsProcessorStep", "RelativeActionsProcessorStep", "MapDeltaActionToRobotActionStep", diff --git a/src/lerobot/processor/factory.py b/src/lerobot/processor/factory.py index 5028122f1..ec929f365 100644 --- a/src/lerobot/processor/factory.py +++ b/src/lerobot/processor/factory.py @@ -14,15 +14,33 @@ # See the License for the specific language governing permissions and # limitations under the License. -from lerobot.types import RobotAction, RobotObservation +from dataclasses import dataclass +from typing import Any +import torch + +from lerobot.configs.policies import PreTrainedConfig +from lerobot.types import PolicyAction, RobotAction, RobotObservation +from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME + +from .batch_processor import AddBatchDimensionProcessorStep from .converters import ( observation_to_transition, + policy_action_to_transition, robot_action_observation_to_transition, transition_to_observation, + transition_to_policy_action, transition_to_robot_action, ) -from .pipeline import IdentityProcessorStep, RobotProcessorPipeline +from .device_processor import DeviceProcessorStep +from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep +from .pipeline import ( + IdentityProcessorStep, + PolicyProcessorPipeline, + ProcessorStep, + RobotProcessorPipeline, +) +from .rename_processor import RenameObservationsProcessorStep def make_default_teleop_action_processor() -> RobotProcessorPipeline[ @@ -61,3 +79,97 @@ def make_default_processors(): robot_action_processor = make_default_robot_action_processor() robot_observation_processor = make_default_robot_observation_processor() return (teleop_action_processor, robot_action_processor, robot_observation_processor) + + +@dataclass +class DefaultPolicyProcessorSteps: + """The canonical processor steps shared by most policies' pre/post pipelines. + + Policies compose these in their own order (step ORDER is a Hub-serialized contract + and intentionally stays explicit per policy) and interleave their custom steps. + """ + + rename_observations: RenameObservationsProcessorStep + add_batch_dim: AddBatchDimensionProcessorStep + to_device: DeviceProcessorStep + normalize: NormalizerProcessorStep + unnormalize: UnnormalizerProcessorStep + to_cpu: DeviceProcessorStep + + +def make_default_policy_processor_steps( + config: PreTrainedConfig, + dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, + *, + normalizer_device: torch.device | str | None = None, +) -> DefaultPolicyProcessorSteps: + """Construct the canonical policy processor steps from a policy config. + + Args: + config: A `PreTrainedConfig` providing `device`, `input_features`, + `output_features` and `normalization_mapping`. + dataset_stats: Dataset statistics used for (un)normalization. + normalizer_device: Device passed to `NormalizerProcessorStep` (some policies pin + their normalization stats to the policy device; most leave it unset). + """ + return DefaultPolicyProcessorSteps( + rename_observations=RenameObservationsProcessorStep(rename_map={}), + add_batch_dim=AddBatchDimensionProcessorStep(), + to_device=DeviceProcessorStep(device=config.device), + normalize=NormalizerProcessorStep( + features={**config.input_features, **config.output_features}, + norm_map=config.normalization_mapping, + stats=dataset_stats, + device=normalizer_device, + ), + unnormalize=UnnormalizerProcessorStep( + features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats + ), + to_cpu=DeviceProcessorStep(device="cpu"), + ) + + +def make_policy_processor_pipelines( + input_steps: list[ProcessorStep], + output_steps: list[ProcessorStep], +) -> tuple[ + PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], + PolicyProcessorPipeline[PolicyAction, PolicyAction], +]: + """Wrap pre/post step lists into the canonical policy pipeline pair. + + Uses the standard pipeline names (which determine the serialized JSON filenames on + the Hub) and the standard policy-action converters on the postprocessor. + """ + return ( + PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( + steps=input_steps, + name=POLICY_PREPROCESSOR_DEFAULT_NAME, + ), + PolicyProcessorPipeline[PolicyAction, PolicyAction]( + steps=output_steps, + name=POLICY_POSTPROCESSOR_DEFAULT_NAME, + to_transition=policy_action_to_transition, + to_output=transition_to_policy_action, + ), + ) + + +def make_default_pre_post_processors( + config: PreTrainedConfig, + dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, + *, + normalizer_device: torch.device | str | None = None, +) -> tuple[ + PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], + PolicyProcessorPipeline[PolicyAction, PolicyAction], +]: + """The pure-scaffold policy pipeline pair: Rename -> Batch -> Device -> Normalize, + and Unnormalize -> Device(cpu). Policies with custom steps or a different step order + compose `make_default_policy_processor_steps` themselves instead. + """ + s = make_default_policy_processor_steps(config, dataset_stats, normalizer_device=normalizer_device) + return make_policy_processor_pipelines( + input_steps=[s.rename_observations, s.add_batch_dim, s.to_device, s.normalize], + output_steps=[s.unnormalize, s.to_cpu], + )