admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-23 09:46:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(processors): share policy pipeline builders (#4016)

Browse Source 
* 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 <imstevenpmwork@ieee.org>

* fix(processor): solve style after commit suggestions

---------

Signed-off-by: Steven Palma <imstevenpmwork@ieee.org>
Co-authored-by: Martino Russi <77496684+nepyope@users.noreply.github.com>
This commit is contained in:
Steven Palma
2026-07-17 14:10:32 +02:00
committed by GitHub
parent b2c062c0f4
commit c5371d0691
18 changed files with 255 additions and 559 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/act/processor_act.py
+2 -39
View File
@@ -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)
src/lerobot/policies/diffusion/processor_diffusion.py
+2 -37
View File
@@ -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)
src/lerobot/policies/eo1/processor_eo1.py
+12 -37
View File
@@ -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)
src/lerobot/policies/fastwam/processor_fastwam.py
+11 -38
View File
@@ -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)
src/lerobot/policies/gaussian_actor/processor_gaussian_actor.py
+2 -38
View File
@@ -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)
src/lerobot/policies/lingbot_va/processor_lingbot_va.py
+10 -30
View File
@@ -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)
src/lerobot/policies/multi_task_dit/processor_multi_task_dit.py
+11 -35
View File
@@ -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)
src/lerobot/policies/pi0/processor_pi0.py
+11 -32
View File
@@ -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)
src/lerobot/policies/pi05/processor_pi05.py
+12 -36
View File
@@ -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)
src/lerobot/policies/pi0_fast/processor_pi0_fast.py
+12 -36
View File
@@ -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)
src/lerobot/policies/smolvla/processor_smolvla.py
+11 -32
View File
@@ -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)
src/lerobot/policies/tdmpc/processor_tdmpc.py
+2 -37
View File
@@ -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)
src/lerobot/policies/vla_jepa/processor_vla_jepa.py
+11 -27
View File
@@ -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)
src/lerobot/policies/vqbet/processor_vqbet.py
+2 -37
View File
@@ -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)
src/lerobot/policies/wall_x/processor_wall_x.py
+11 -32
View File
@@ -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")
src/lerobot/policies/xvla/processor_xvla.py
+11 -34
View File
@@ -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
src/lerobot/processor/__init__.py
+8
View File
@@ -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",
src/lerobot/processor/factory.py
+114 -2
View File
@@ -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],
)