chore(processor): rename RobotProcessor -> DataProcessorPipeline (#1850)

This commit is contained in:
Steven Palma
2025-09-03 17:13:16 +02:00
committed by GitHub
parent 4ebe482a7e
commit 8c796b39f5
38 changed files with 326 additions and 298 deletions
+3 -3
View File
@@ -21,7 +21,7 @@ from lerobot.datasets.utils import merge_features
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.policies.act.modeling_act import ACTPolicy from lerobot.policies.act.modeling_act import ACTPolicy
from lerobot.policies.factory import make_pre_post_processors from lerobot.policies.factory import make_pre_post_processors
from lerobot.processor import RobotProcessor from lerobot.processor import DataProcessorPipeline
from lerobot.processor.converters import ( from lerobot.processor.converters import (
to_output_robot_action, to_output_robot_action,
to_transition_robot_observation, to_transition_robot_observation,
@@ -65,7 +65,7 @@ kinematics_solver = RobotKinematics(
) )
# Build pipeline to convert ee pose action to joint action # Build pipeline to convert ee pose action to joint action
robot_ee_to_joints_processor = RobotProcessor( robot_ee_to_joints_processor = DataProcessorPipeline(
steps=[ steps=[
AddRobotObservationAsComplimentaryData(robot=robot), AddRobotObservationAsComplimentaryData(robot=robot),
InverseKinematicsEEToJoints( InverseKinematicsEEToJoints(
@@ -79,7 +79,7 @@ robot_ee_to_joints_processor = RobotProcessor(
) )
# Build pipeline to convert joint observation to ee pose observation # Build pipeline to convert joint observation to ee pose observation
robot_joints_to_ee_pose_processor = RobotProcessor( robot_joints_to_ee_pose_processor = DataProcessorPipeline(
steps=[ steps=[
ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
], ],
+4 -4
View File
@@ -20,7 +20,7 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features
from lerobot.datasets.utils import merge_features from lerobot.datasets.utils import merge_features
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import RobotProcessor from lerobot.processor import DataProcessorPipeline
from lerobot.processor.converters import ( from lerobot.processor.converters import (
to_output_robot_action, to_output_robot_action,
to_transition_robot_observation, to_transition_robot_observation,
@@ -73,7 +73,7 @@ kinematics_solver = RobotKinematics(
) )
# Build pipeline to convert phone action to ee pose action # Build pipeline to convert phone action to ee pose action
phone_to_robot_ee_pose_processor = RobotProcessor( phone_to_robot_ee_pose_processor = DataProcessorPipeline(
steps=[ steps=[
MapPhoneActionToRobotAction(platform=teleop_config.phone_os), MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
AddRobotObservationAsComplimentaryData(robot=robot), AddRobotObservationAsComplimentaryData(robot=robot),
@@ -93,7 +93,7 @@ phone_to_robot_ee_pose_processor = RobotProcessor(
) )
# Build pipeline to convert ee pose action to joint action # Build pipeline to convert ee pose action to joint action
robot_ee_to_joints_processor = RobotProcessor( robot_ee_to_joints_processor = DataProcessorPipeline(
steps=[ steps=[
InverseKinematicsEEToJoints( InverseKinematicsEEToJoints(
kinematics=kinematics_solver, kinematics=kinematics_solver,
@@ -110,7 +110,7 @@ robot_ee_to_joints_processor = RobotProcessor(
) )
# Build pipeline to convert joint observation to ee pose observation # Build pipeline to convert joint observation to ee pose observation
robot_joints_to_ee_pose = RobotProcessor( robot_joints_to_ee_pose = DataProcessorPipeline(
steps=[ steps=[
ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
], ],
+2 -2
View File
@@ -19,7 +19,7 @@ import time
from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import RobotProcessor from lerobot.processor import DataProcessorPipeline
from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.robot_kinematic_processor import ( from lerobot.robots.so100_follower.robot_kinematic_processor import (
@@ -50,7 +50,7 @@ kinematics_solver = RobotKinematics(
) )
# Build pipeline to convert ee pose action to joint action # Build pipeline to convert ee pose action to joint action
robot_ee_to_joints_processor = RobotProcessor( robot_ee_to_joints_processor = DataProcessorPipeline(
steps=[ steps=[
AddRobotObservationAsComplimentaryData(robot=robot), AddRobotObservationAsComplimentaryData(robot=robot),
InverseKinematicsEEToJoints( InverseKinematicsEEToJoints(
+2 -2
View File
@@ -16,7 +16,7 @@
import time import time
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import RobotProcessor from lerobot.processor import DataProcessorPipeline
from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.robot_kinematic_processor import ( from lerobot.robots.so100_follower.robot_kinematic_processor import (
@@ -49,7 +49,7 @@ kinematics_solver = RobotKinematics(
) )
# Build pipeline to convert phone action to ee pose action to joint action # Build pipeline to convert phone action to ee pose action to joint action
phone_to_robot_joints_processor = RobotProcessor( phone_to_robot_joints_processor = DataProcessorPipeline(
steps=[ steps=[
MapPhoneActionToRobotAction(platform=teleop_config.phone_os), MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
AddRobotObservationAsComplimentaryData(robot=robot), AddRobotObservationAsComplimentaryData(robot=robot),
+2 -2
View File
@@ -17,11 +17,11 @@ from typing import Any
from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE
from lerobot.datasets.utils import hw_to_dataset_features from lerobot.datasets.utils import hw_to_dataset_features
from lerobot.processor import RobotProcessor from lerobot.processor import DataProcessorPipeline
def aggregate_pipeline_dataset_features( def aggregate_pipeline_dataset_features(
pipeline: RobotProcessor, pipeline: DataProcessorPipeline,
initial_features: dict[str, Any], initial_features: dict[str, Any],
*, *,
use_videos: bool = True, use_videos: bool = True,
+4 -4
View File
@@ -18,11 +18,11 @@ import torch
from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME
from lerobot.policies.act.configuration_act import ACTConfig from lerobot.policies.act.configuration_act import ACTConfig
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
NormalizerProcessor, NormalizerProcessor,
ProcessorKwargs, ProcessorKwargs,
RenameProcessor, RenameProcessor,
RobotProcessor,
ToBatchProcessor, ToBatchProcessor,
UnnormalizerProcessor, UnnormalizerProcessor,
) )
@@ -33,7 +33,7 @@ def make_act_pre_post_processors(
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
preprocessor_kwargs: ProcessorKwargs | None = None, preprocessor_kwargs: ProcessorKwargs | None = None,
postprocessor_kwargs: ProcessorKwargs | None = None, postprocessor_kwargs: ProcessorKwargs | None = None,
) -> tuple[RobotProcessor, RobotProcessor]: ) -> tuple[DataProcessorPipeline, DataProcessorPipeline]:
if preprocessor_kwargs is None: if preprocessor_kwargs is None:
preprocessor_kwargs = {} preprocessor_kwargs = {}
if postprocessor_kwargs is None: if postprocessor_kwargs is None:
@@ -57,12 +57,12 @@ def make_act_pre_post_processors(
] ]
return ( return (
RobotProcessor( DataProcessorPipeline(
steps=input_steps, steps=input_steps,
name=PREPROCESSOR_DEFAULT_NAME, name=PREPROCESSOR_DEFAULT_NAME,
**preprocessor_kwargs, **preprocessor_kwargs,
), ),
RobotProcessor( DataProcessorPipeline(
steps=output_steps, steps=output_steps,
name=POSTPROCESSOR_DEFAULT_NAME, name=POSTPROCESSOR_DEFAULT_NAME,
**postprocessor_kwargs, **postprocessor_kwargs,
@@ -19,11 +19,11 @@ import torch
from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME
from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
NormalizerProcessor, NormalizerProcessor,
ProcessorKwargs, ProcessorKwargs,
RenameProcessor, RenameProcessor,
RobotProcessor,
ToBatchProcessor, ToBatchProcessor,
UnnormalizerProcessor, UnnormalizerProcessor,
) )
@@ -34,7 +34,7 @@ def make_diffusion_pre_post_processors(
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
preprocessor_kwargs: ProcessorKwargs | None = None, preprocessor_kwargs: ProcessorKwargs | None = None,
postprocessor_kwargs: ProcessorKwargs | None = None, postprocessor_kwargs: ProcessorKwargs | None = None,
) -> tuple[RobotProcessor, RobotProcessor]: ) -> tuple[DataProcessorPipeline, DataProcessorPipeline]:
if preprocessor_kwargs is None: if preprocessor_kwargs is None:
preprocessor_kwargs = {} preprocessor_kwargs = {}
if postprocessor_kwargs is None: if postprocessor_kwargs is None:
@@ -57,12 +57,12 @@ def make_diffusion_pre_post_processors(
), ),
] ]
return ( return (
RobotProcessor( DataProcessorPipeline(
steps=input_steps, steps=input_steps,
name=PREPROCESSOR_DEFAULT_NAME, name=PREPROCESSOR_DEFAULT_NAME,
**preprocessor_kwargs, **preprocessor_kwargs,
), ),
RobotProcessor( DataProcessorPipeline(
steps=output_steps, steps=output_steps,
name=POSTPROCESSOR_DEFAULT_NAME, name=POSTPROCESSOR_DEFAULT_NAME,
**postprocessor_kwargs, **postprocessor_kwargs,
+4 -4
View File
@@ -38,7 +38,7 @@ from lerobot.policies.sac.reward_model.configuration_classifier import RewardCla
from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
from lerobot.processor import ProcessorKwargs, RobotProcessor from lerobot.processor import DataProcessorPipeline, ProcessorKwargs
def get_policy_class(name: str) -> type[PreTrainedPolicy]: def get_policy_class(name: str) -> type[PreTrainedPolicy]:
@@ -122,7 +122,7 @@ def make_pre_post_processors(
policy_cfg: PreTrainedConfig, policy_cfg: PreTrainedConfig,
pretrained_path: str | None = None, pretrained_path: str | None = None,
**kwargs: Unpack[ProcessorConfigKwargs], **kwargs: Unpack[ProcessorConfigKwargs],
) -> tuple[RobotProcessor, RobotProcessor]: ) -> tuple[DataProcessorPipeline, DataProcessorPipeline]:
"""Make a processor instance for a given policy type. """Make a processor instance for a given policy type.
This function creates the appropriate processor configuration based on the policy type. This function creates the appropriate processor configuration based on the policy type.
@@ -146,14 +146,14 @@ def make_pre_post_processors(
postprocessor_kwargs = kwargs.get("postprocessor_kwargs", {}) postprocessor_kwargs = kwargs.get("postprocessor_kwargs", {})
return ( return (
RobotProcessor.from_pretrained( DataProcessorPipeline.from_pretrained(
pretrained_model_name_or_path=pretrained_path, pretrained_model_name_or_path=pretrained_path,
config_filename=kwargs.get("preprocessor_config_filename", "robot_preprocessor.json"), config_filename=kwargs.get("preprocessor_config_filename", "robot_preprocessor.json"),
overrides=kwargs.get("preprocessor_overrides", {}), overrides=kwargs.get("preprocessor_overrides", {}),
to_transition=preprocessor_kwargs.get("to_transition"), to_transition=preprocessor_kwargs.get("to_transition"),
to_output=preprocessor_kwargs.get("to_output"), to_output=preprocessor_kwargs.get("to_output"),
), ),
RobotProcessor.from_pretrained( DataProcessorPipeline.from_pretrained(
pretrained_model_name_or_path=pretrained_path, pretrained_model_name_or_path=pretrained_path,
config_filename=kwargs.get("postprocessor_config_filename", "robot_postprocessor.json"), config_filename=kwargs.get("postprocessor_config_filename", "robot_postprocessor.json"),
overrides=kwargs.get("postprocessor_overrides", {}), overrides=kwargs.get("postprocessor_overrides", {}),
+4 -4
View File
@@ -22,13 +22,13 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_N
from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.policies.pi0.configuration_pi0 import PI0Config
from lerobot.processor import ( from lerobot.processor import (
ComplementaryDataProcessor, ComplementaryDataProcessor,
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
NormalizerProcessor, NormalizerProcessor,
ProcessorKwargs, ProcessorKwargs,
ProcessorStep, ProcessorStep,
ProcessorStepRegistry, ProcessorStepRegistry,
RenameProcessor, RenameProcessor,
RobotProcessor,
ToBatchProcessor, ToBatchProcessor,
TokenizerProcessor, TokenizerProcessor,
UnnormalizerProcessor, UnnormalizerProcessor,
@@ -72,7 +72,7 @@ def make_pi0_pre_post_processors(
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
preprocessor_kwargs: ProcessorKwargs | None = None, preprocessor_kwargs: ProcessorKwargs | None = None,
postprocessor_kwargs: ProcessorKwargs | None = None, postprocessor_kwargs: ProcessorKwargs | None = None,
) -> tuple[RobotProcessor, RobotProcessor]: ) -> tuple[DataProcessorPipeline, DataProcessorPipeline]:
if preprocessor_kwargs is None: if preprocessor_kwargs is None:
preprocessor_kwargs = {} preprocessor_kwargs = {}
if postprocessor_kwargs is None: if postprocessor_kwargs is None:
@@ -105,12 +105,12 @@ def make_pi0_pre_post_processors(
] ]
return ( return (
RobotProcessor( DataProcessorPipeline(
steps=input_steps, steps=input_steps,
name=PREPROCESSOR_DEFAULT_NAME, name=PREPROCESSOR_DEFAULT_NAME,
**preprocessor_kwargs, **preprocessor_kwargs,
), ),
RobotProcessor( DataProcessorPipeline(
steps=output_steps, steps=output_steps,
name=POSTPROCESSOR_DEFAULT_NAME, name=POSTPROCESSOR_DEFAULT_NAME,
**postprocessor_kwargs, **postprocessor_kwargs,
@@ -19,11 +19,11 @@ import torch
from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME
from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.policies.pi0.configuration_pi0 import PI0Config
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
NormalizerProcessor, NormalizerProcessor,
ProcessorKwargs, ProcessorKwargs,
RenameProcessor, RenameProcessor,
RobotProcessor,
ToBatchProcessor, ToBatchProcessor,
UnnormalizerProcessor, UnnormalizerProcessor,
) )
@@ -34,7 +34,7 @@ def make_pi0fast_pre_post_processors(
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
preprocessor_kwargs: ProcessorKwargs | None = None, preprocessor_kwargs: ProcessorKwargs | None = None,
postprocessor_kwargs: ProcessorKwargs | None = None, postprocessor_kwargs: ProcessorKwargs | None = None,
) -> tuple[RobotProcessor, RobotProcessor]: ) -> tuple[DataProcessorPipeline, DataProcessorPipeline]:
if preprocessor_kwargs is None: if preprocessor_kwargs is None:
preprocessor_kwargs = {} preprocessor_kwargs = {}
if postprocessor_kwargs is None: if postprocessor_kwargs is None:
@@ -57,12 +57,12 @@ def make_pi0fast_pre_post_processors(
), ),
] ]
return ( return (
RobotProcessor( DataProcessorPipeline(
steps=input_steps, steps=input_steps,
name=PREPROCESSOR_DEFAULT_NAME, name=PREPROCESSOR_DEFAULT_NAME,
**preprocessor_kwargs, **preprocessor_kwargs,
), ),
RobotProcessor( DataProcessorPipeline(
steps=output_steps, steps=output_steps,
name=POSTPROCESSOR_DEFAULT_NAME, name=POSTPROCESSOR_DEFAULT_NAME,
**postprocessor_kwargs, **postprocessor_kwargs,
+4 -4
View File
@@ -20,11 +20,11 @@ import torch
from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME
from lerobot.policies.sac.configuration_sac import SACConfig from lerobot.policies.sac.configuration_sac import SACConfig
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
NormalizerProcessor, NormalizerProcessor,
ProcessorKwargs, ProcessorKwargs,
RenameProcessor, RenameProcessor,
RobotProcessor,
ToBatchProcessor, ToBatchProcessor,
UnnormalizerProcessor, UnnormalizerProcessor,
) )
@@ -35,7 +35,7 @@ def make_sac_pre_post_processors(
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
preprocessor_kwargs: ProcessorKwargs | None = None, preprocessor_kwargs: ProcessorKwargs | None = None,
postprocessor_kwargs: ProcessorKwargs | None = None, postprocessor_kwargs: ProcessorKwargs | None = None,
) -> tuple[RobotProcessor, RobotProcessor]: ) -> tuple[DataProcessorPipeline, DataProcessorPipeline]:
if preprocessor_kwargs is None: if preprocessor_kwargs is None:
preprocessor_kwargs = {} preprocessor_kwargs = {}
if postprocessor_kwargs is None: if postprocessor_kwargs is None:
@@ -58,12 +58,12 @@ def make_sac_pre_post_processors(
), ),
] ]
return ( return (
RobotProcessor( DataProcessorPipeline(
steps=input_steps, steps=input_steps,
name=PREPROCESSOR_DEFAULT_NAME, name=PREPROCESSOR_DEFAULT_NAME,
**preprocessor_kwargs, **preprocessor_kwargs,
), ),
RobotProcessor( DataProcessorPipeline(
steps=output_steps, steps=output_steps,
name=POSTPROCESSOR_DEFAULT_NAME, name=POSTPROCESSOR_DEFAULT_NAME,
**postprocessor_kwargs, **postprocessor_kwargs,
@@ -17,11 +17,11 @@ import torch
from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
IdentityProcessor, IdentityProcessor,
NormalizerProcessor, NormalizerProcessor,
ProcessorKwargs, ProcessorKwargs,
RobotProcessor,
) )
@@ -30,7 +30,7 @@ def make_classifier_processor(
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
preprocessor_kwargs: ProcessorKwargs | None = None, preprocessor_kwargs: ProcessorKwargs | None = None,
postprocessor_kwargs: ProcessorKwargs | None = None, postprocessor_kwargs: ProcessorKwargs | None = None,
) -> tuple[RobotProcessor, RobotProcessor]: ) -> tuple[DataProcessorPipeline, DataProcessorPipeline]:
if preprocessor_kwargs is None: if preprocessor_kwargs is None:
preprocessor_kwargs = {} preprocessor_kwargs = {}
if postprocessor_kwargs is None: if postprocessor_kwargs is None:
@@ -48,12 +48,12 @@ def make_classifier_processor(
output_steps = [DeviceProcessor(device="cpu"), IdentityProcessor()] output_steps = [DeviceProcessor(device="cpu"), IdentityProcessor()]
return ( return (
RobotProcessor( DataProcessorPipeline(
steps=input_steps, steps=input_steps,
name="classifier_preprocessor", name="classifier_preprocessor",
**preprocessor_kwargs, **preprocessor_kwargs,
), ),
RobotProcessor( DataProcessorPipeline(
steps=output_steps, steps=output_steps,
name="classifier_postprocessor", name="classifier_postprocessor",
**postprocessor_kwargs, **postprocessor_kwargs,
@@ -21,12 +21,12 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_N
from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
from lerobot.processor import ( from lerobot.processor import (
ComplementaryDataProcessor, ComplementaryDataProcessor,
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
NormalizerProcessor, NormalizerProcessor,
ProcessorKwargs, ProcessorKwargs,
ProcessorStepRegistry, ProcessorStepRegistry,
RenameProcessor, RenameProcessor,
RobotProcessor,
ToBatchProcessor, ToBatchProcessor,
TokenizerProcessor, TokenizerProcessor,
UnnormalizerProcessor, UnnormalizerProcessor,
@@ -38,7 +38,7 @@ def make_smolvla_pre_post_processors(
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
preprocessor_kwargs: ProcessorKwargs | None = None, preprocessor_kwargs: ProcessorKwargs | None = None,
postprocessor_kwargs: ProcessorKwargs | None = None, postprocessor_kwargs: ProcessorKwargs | None = None,
) -> tuple[RobotProcessor, RobotProcessor]: ) -> tuple[DataProcessorPipeline, DataProcessorPipeline]:
if preprocessor_kwargs is None: if preprocessor_kwargs is None:
preprocessor_kwargs = {} preprocessor_kwargs = {}
if postprocessor_kwargs is None: if postprocessor_kwargs is None:
@@ -68,12 +68,12 @@ def make_smolvla_pre_post_processors(
), ),
] ]
return ( return (
RobotProcessor( DataProcessorPipeline(
steps=input_steps, steps=input_steps,
name=PREPROCESSOR_DEFAULT_NAME, name=PREPROCESSOR_DEFAULT_NAME,
**preprocessor_kwargs, **preprocessor_kwargs,
), ),
RobotProcessor( DataProcessorPipeline(
steps=output_steps, steps=output_steps,
name=POSTPROCESSOR_DEFAULT_NAME, name=POSTPROCESSOR_DEFAULT_NAME,
**postprocessor_kwargs, **postprocessor_kwargs,
@@ -19,11 +19,11 @@ import torch
from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME
from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
NormalizerProcessor, NormalizerProcessor,
ProcessorKwargs, ProcessorKwargs,
RenameProcessor, RenameProcessor,
RobotProcessor,
ToBatchProcessor, ToBatchProcessor,
UnnormalizerProcessor, UnnormalizerProcessor,
) )
@@ -34,7 +34,7 @@ def make_tdmpc_pre_post_processors(
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
preprocessor_kwargs: ProcessorKwargs | None = None, preprocessor_kwargs: ProcessorKwargs | None = None,
postprocessor_kwargs: ProcessorKwargs | None = None, postprocessor_kwargs: ProcessorKwargs | None = None,
) -> tuple[RobotProcessor, RobotProcessor]: ) -> tuple[DataProcessorPipeline, DataProcessorPipeline]:
if preprocessor_kwargs is None: if preprocessor_kwargs is None:
preprocessor_kwargs = {} preprocessor_kwargs = {}
if postprocessor_kwargs is None: if postprocessor_kwargs is None:
@@ -57,12 +57,12 @@ def make_tdmpc_pre_post_processors(
), ),
] ]
return ( return (
RobotProcessor( DataProcessorPipeline(
steps=input_steps, steps=input_steps,
name=PREPROCESSOR_DEFAULT_NAME, name=PREPROCESSOR_DEFAULT_NAME,
**preprocessor_kwargs, **preprocessor_kwargs,
), ),
RobotProcessor( DataProcessorPipeline(
steps=output_steps, steps=output_steps,
name=POSTPROCESSOR_DEFAULT_NAME, name=POSTPROCESSOR_DEFAULT_NAME,
**postprocessor_kwargs, **postprocessor_kwargs,
@@ -20,11 +20,11 @@ import torch
from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME
from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
NormalizerProcessor, NormalizerProcessor,
ProcessorKwargs, ProcessorKwargs,
RenameProcessor, RenameProcessor,
RobotProcessor,
ToBatchProcessor, ToBatchProcessor,
UnnormalizerProcessor, UnnormalizerProcessor,
) )
@@ -35,7 +35,7 @@ def make_vqbet_pre_post_processors(
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
preprocessor_kwargs: ProcessorKwargs | None = None, preprocessor_kwargs: ProcessorKwargs | None = None,
postprocessor_kwargs: ProcessorKwargs | None = None, postprocessor_kwargs: ProcessorKwargs | None = None,
) -> tuple[RobotProcessor, RobotProcessor]: ) -> tuple[DataProcessorPipeline, DataProcessorPipeline]:
if preprocessor_kwargs is None: if preprocessor_kwargs is None:
preprocessor_kwargs = {} preprocessor_kwargs = {}
if postprocessor_kwargs is None: if postprocessor_kwargs is None:
@@ -58,12 +58,12 @@ def make_vqbet_pre_post_processors(
), ),
] ]
return ( return (
RobotProcessor( DataProcessorPipeline(
steps=input_steps, steps=input_steps,
name=PREPROCESSOR_DEFAULT_NAME, name=PREPROCESSOR_DEFAULT_NAME,
**preprocessor_kwargs, **preprocessor_kwargs,
), ),
RobotProcessor( DataProcessorPipeline(
steps=output_steps, steps=output_steps,
name=POSTPROCESSOR_DEFAULT_NAME, name=POSTPROCESSOR_DEFAULT_NAME,
**postprocessor_kwargs, **postprocessor_kwargs,
+2 -2
View File
@@ -41,6 +41,7 @@ from .observation_processor import VanillaObservationProcessor
from .pipeline import ( from .pipeline import (
ActionProcessor, ActionProcessor,
ComplementaryDataProcessor, ComplementaryDataProcessor,
DataProcessorPipeline,
DoneProcessor, DoneProcessor,
IdentityProcessor, IdentityProcessor,
InfoProcessor, InfoProcessor,
@@ -49,7 +50,6 @@ from .pipeline import (
ProcessorStep, ProcessorStep,
ProcessorStepRegistry, ProcessorStepRegistry,
RewardProcessor, RewardProcessor,
RobotProcessor,
TruncatedProcessor, TruncatedProcessor,
) )
from .rename_processor import RenameProcessor from .rename_processor import RenameProcessor
@@ -85,7 +85,7 @@ __all__ = [
"RenameProcessor", "RenameProcessor",
"RewardClassifierProcessor", "RewardClassifierProcessor",
"RewardProcessor", "RewardProcessor",
"RobotProcessor", "DataProcessorPipeline",
"TimeLimitProcessor", "TimeLimitProcessor",
"ToBatchProcessor", "ToBatchProcessor",
"TokenizerProcessor", "TokenizerProcessor",
@@ -50,7 +50,7 @@ from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from .batch_processor import ToBatchProcessor from .batch_processor import ToBatchProcessor
from .device_processor import DeviceProcessor from .device_processor import DeviceProcessor
from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor
from .pipeline import RobotProcessor from .pipeline import DataProcessorPipeline
from .rename_processor import RenameProcessor from .rename_processor import RenameProcessor
# Policy type to class mapping # Policy type to class mapping
@@ -422,14 +422,14 @@ def main():
ToBatchProcessor(), ToBatchProcessor(),
DeviceProcessor(device=policy_config.device), DeviceProcessor(device=policy_config.device),
] ]
preprocessor = RobotProcessor(steps=preprocessor_steps, name="robot_preprocessor") preprocessor = DataProcessorPipeline(steps=preprocessor_steps, name="robot_preprocessor")
# Create postprocessor with unnormalizer for outputs only # Create postprocessor with unnormalizer for outputs only
postprocessor_steps = [ postprocessor_steps = [
DeviceProcessor(device="cpu"), DeviceProcessor(device="cpu"),
UnnormalizerProcessor(features=output_features, norm_map=norm_map, stats=stats), UnnormalizerProcessor(features=output_features, norm_map=norm_map, stats=stats),
] ]
postprocessor = RobotProcessor(steps=postprocessor_steps, name="robot_postprocessor") postprocessor = DataProcessorPipeline(steps=postprocessor_steps, name="robot_postprocessor")
# Determine hub repo ID if pushing to hub # Determine hub repo ID if pushing to hub
if args.push_to_hub: if args.push_to_hub:
+4 -2
View File
@@ -12,7 +12,7 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset
from .converters import to_tensor from .converters import to_tensor
from .core import EnvTransition, TransitionKey from .core import EnvTransition, TransitionKey
from .pipeline import ProcessorStep, ProcessorStepRegistry, RobotProcessor from .pipeline import DataProcessorPipeline, ProcessorStep, ProcessorStepRegistry
@dataclass @dataclass
@@ -252,7 +252,9 @@ class UnnormalizerProcessor(_NormalizationMixin, ProcessorStep):
return features return features
def hotswap_stats(robot_processor: RobotProcessor, stats: dict[str, dict[str, Any]]) -> RobotProcessor: def hotswap_stats(
robot_processor: DataProcessorPipeline, stats: dict[str, dict[str, Any]]
) -> DataProcessorPipeline:
""" """
Replaces normalization statistics in a RobotProcessor pipeline. Replaces normalization statistics in a RobotProcessor pipeline.
+21 -19
View File
@@ -117,7 +117,7 @@ class ProcessorStep(ABC):
A step is any callable accepting a full `EnvTransition` dict and A step is any callable accepting a full `EnvTransition` dict and
returning a (possibly modified) dict of the same structure. Implementers returning a (possibly modified) dict of the same structure. Implementers
are encouraged—but not required—to expose the optional helper methods are encouraged—but not required—to expose the optional helper methods
listed below. When present, these hooks let `RobotProcessor` listed below. When present, these hooks let `DataProcessorPipeline`
automatically serialise the step's configuration and learnable state using automatically serialise the step's configuration and learnable state using
a safe-to-share JSON + SafeTensors format. a safe-to-share JSON + SafeTensors format.
@@ -175,14 +175,14 @@ class ProcessorStep(ABC):
class ProcessorKwargs(TypedDict, total=False): class ProcessorKwargs(TypedDict, total=False):
"""Keyword arguments for RobotProcessor constructor.""" """Keyword arguments for DataProcessorPipeline constructor."""
to_transition: Callable[[dict[str, Any]], EnvTransition] | None to_transition: Callable[[dict[str, Any]], EnvTransition] | None
to_output: Callable[[EnvTransition], Any] | None to_output: Callable[[EnvTransition], Any] | None
@dataclass @dataclass
class RobotProcessor(ModelHubMixin, Generic[TOutput]): class DataProcessorPipeline(ModelHubMixin, Generic[TOutput]):
""" """
Composable, debuggable post-processing processor for robot transitions. Composable, debuggable post-processing processor for robot transitions.
@@ -196,7 +196,7 @@ class RobotProcessor(ModelHubMixin, Generic[TOutput]):
Args: Args:
steps: Ordered list of processing steps executed on every call. Defaults to empty list. steps: Ordered list of processing steps executed on every call. Defaults to empty list.
name: Human-readable identifier that is persisted inside the JSON config. name: Human-readable identifier that is persisted inside the JSON config.
Defaults to "RobotProcessor". Defaults to "DataProcessorPipeline".
to_transition: Function to convert batch dict to EnvTransition dict. to_transition: Function to convert batch dict to EnvTransition dict.
Defaults to _default_batch_to_transition. Defaults to _default_batch_to_transition.
to_output: Function to convert EnvTransition dict to the desired output format of type TOutput. to_output: Function to convert EnvTransition dict to the desired output format of type TOutput.
@@ -210,18 +210,20 @@ class RobotProcessor(ModelHubMixin, Generic[TOutput]):
Type Safety Examples: Type Safety Examples:
```python ```python
# Default behavior - returns batch dict # Default behavior - returns batch dict
processor: RobotProcessor[dict[str, Any]] = RobotProcessor(steps=[some_step1, some_step2]) processor: DataProcessorPipeline[dict[str, Any]] = DataProcessorPipeline(
steps=[some_step1, some_step2]
)
result: dict[str, Any] = processor(batch_data) # Type checker knows this is a dict result: dict[str, Any] = processor(batch_data) # Type checker knows this is a dict
# For EnvTransition output, explicitly specify identity function # For EnvTransition output, explicitly specify identity function
transition_processor: RobotProcessor[EnvTransition] = RobotProcessor( transition_processor: DataProcessorPipeline[EnvTransition] = DataProcessorPipeline(
steps=[some_step1, some_step2], steps=[some_step1, some_step2],
to_output=lambda x: x, # Identity function to_output=lambda x: x, # Identity function
) )
result: EnvTransition = transition_processor(batch_data) # Type checker knows this is EnvTransition result: EnvTransition = transition_processor(batch_data) # Type checker knows this is EnvTransition
# For custom output types # For custom output types
processor: RobotProcessor[str] = RobotProcessor( processor: DataProcessorPipeline[str] = DataProcessorPipeline(
steps=[custom_step], to_output=lambda t: f"Processed {len(t)} keys" steps=[custom_step], to_output=lambda t: f"Processed {len(t)} keys"
) )
result: str = processor(batch_data) # Type checker knows this is str result: str = processor(batch_data) # Type checker knows this is str
@@ -243,7 +245,7 @@ class RobotProcessor(ModelHubMixin, Generic[TOutput]):
""" """
steps: Sequence[ProcessorStep] = field(default_factory=list) steps: Sequence[ProcessorStep] = field(default_factory=list)
name: str = "RobotProcessor" name: str = "DataProcessorPipeline"
to_transition: Callable[[dict[str, Any]], EnvTransition] = field(default=batch_to_transition, repr=False) to_transition: Callable[[dict[str, Any]], EnvTransition] = field(default=batch_to_transition, repr=False)
to_output: Callable[[EnvTransition], TOutput] = field( to_output: Callable[[EnvTransition], TOutput] = field(
@@ -419,7 +421,7 @@ class RobotProcessor(ModelHubMixin, Generic[TOutput]):
to_transition: Callable[[dict[str, Any]], EnvTransition] | None = None, to_transition: Callable[[dict[str, Any]], EnvTransition] | None = None,
to_output: Callable[[EnvTransition], TOutput] | None = None, to_output: Callable[[EnvTransition], TOutput] | None = None,
**kwargs, **kwargs,
) -> RobotProcessor[TOutput]: ) -> DataProcessorPipeline[TOutput]:
"""Load a serialized processor from source (local path or Hugging Face Hub identifier). """Load a serialized processor from source (local path or Hugging Face Hub identifier).
Args: Args:
@@ -440,7 +442,7 @@ class RobotProcessor(ModelHubMixin, Generic[TOutput]):
Use identity function (lambda x: x) for EnvTransition output. Use identity function (lambda x: x) for EnvTransition output.
Returns: Returns:
A RobotProcessor[TOutput] instance loaded from the saved configuration. A DataProcessorPipeline[TOutput] instance loaded from the saved configuration.
Raises: Raises:
ImportError: If a processor step class cannot be loaded or imported. ImportError: If a processor step class cannot be loaded or imported.
@@ -450,12 +452,12 @@ class RobotProcessor(ModelHubMixin, Generic[TOutput]):
Examples: Examples:
Basic loading: Basic loading:
```python ```python
processor = RobotProcessor.from_pretrained("path/to/processor") processor = DataProcessorPipeline.from_pretrained("path/to/processor")
``` ```
Loading specific config file: Loading specific config file:
```python ```python
processor = RobotProcessor.from_pretrained( processor = DataProcessorPipeline.from_pretrained(
"username/multi-processor-repo", config_filename="preprocessor.json" "username/multi-processor-repo", config_filename="preprocessor.json"
) )
``` ```
@@ -465,14 +467,14 @@ class RobotProcessor(ModelHubMixin, Generic[TOutput]):
import gym import gym
env = gym.make("CartPole-v1") env = gym.make("CartPole-v1")
processor = RobotProcessor.from_pretrained( processor = DataProcessorPipeline.from_pretrained(
"username/cartpole-processor", overrides={"ActionRepeatStep": {"env": env}} "username/cartpole-processor", overrides={"ActionRepeatStep": {"env": env}}
) )
``` ```
Multiple overrides: Multiple overrides:
```python ```python
processor = RobotProcessor.from_pretrained( processor = DataProcessorPipeline.from_pretrained(
"path/to/processor", "path/to/processor",
overrides={ overrides={
"CustomStep": {"param1": "new_value"}, "CustomStep": {"param1": "new_value"},
@@ -656,7 +658,7 @@ class RobotProcessor(ModelHubMixin, Generic[TOutput]):
return cls( return cls(
steps=steps, steps=steps,
name=loaded_config.get("name", "RobotProcessor"), name=loaded_config.get("name", "DataProcessorPipeline"),
to_transition=to_transition or batch_to_transition, to_transition=to_transition or batch_to_transition,
# Cast is necessary here: Same type-checker limitation as above. # Cast is necessary here: Same type-checker limitation as above.
# When to_output is None, we use the default which returns dict[str, Any]. # When to_output is None, we use the default which returns dict[str, Any].
@@ -668,13 +670,13 @@ class RobotProcessor(ModelHubMixin, Generic[TOutput]):
"""Return the number of steps in the processor.""" """Return the number of steps in the processor."""
return len(self.steps) return len(self.steps)
def __getitem__(self, idx: int | slice) -> ProcessorStep | RobotProcessor[TOutput]: def __getitem__(self, idx: int | slice) -> ProcessorStep | DataProcessorPipeline[TOutput]:
"""Indexing helper exposing underlying steps. """Indexing helper exposing underlying steps.
* ``int`` returns the idx-th ProcessorStep. * ``int`` returns the idx-th ProcessorStep.
* ``slice`` returns a new RobotProcessor with the sliced steps. * ``slice`` returns a new DataProcessorPipeline with the sliced steps.
""" """
if isinstance(idx, slice): if isinstance(idx, slice):
return RobotProcessor( return DataProcessorPipeline(
steps=self.steps[idx], steps=self.steps[idx],
name=self.name, name=self.name,
to_transition=self.to_transition, to_transition=self.to_transition,
@@ -745,7 +747,7 @@ class RobotProcessor(ModelHubMixin, Generic[TOutput]):
parts = [f"name='{self.name}'", steps_repr] parts = [f"name='{self.name}'", steps_repr]
return f"RobotProcessor({', '.join(parts)})" return f"DataProcessorPipeline({', '.join(parts)})"
def __post_init__(self): def __post_init__(self):
for i, step in enumerate(self.steps): for i, step in enumerate(self.steps):
+9 -9
View File
@@ -76,7 +76,7 @@ from lerobot.datasets.utils import hw_to_dataset_features
from lerobot.datasets.video_utils import VideoEncodingManager from lerobot.datasets.video_utils import VideoEncodingManager
from lerobot.policies.factory import make_policy, make_pre_post_processors from lerobot.policies.factory import make_policy, make_pre_post_processors
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.processor import IdentityProcessor, RobotProcessor, TransitionKey from lerobot.processor import DataProcessorPipeline, IdentityProcessor, TransitionKey
from lerobot.processor.converters import ( from lerobot.processor.converters import (
to_output_robot_action, to_output_robot_action,
to_transition_robot_observation, to_transition_robot_observation,
@@ -235,22 +235,22 @@ def record_loop(
dataset: LeRobotDataset | None = None, dataset: LeRobotDataset | None = None,
teleop: Teleoperator | list[Teleoperator] | None = None, teleop: Teleoperator | list[Teleoperator] | None = None,
policy: PreTrainedPolicy | None = None, policy: PreTrainedPolicy | None = None,
preprocessor: RobotProcessor | None = None, preprocessor: DataProcessorPipeline | None = None,
postprocessor: RobotProcessor | None = None, postprocessor: DataProcessorPipeline | None = None,
control_time_s: int | None = None, control_time_s: int | None = None,
teleop_action_processor: RobotProcessor | None = None, # runs after teleop teleop_action_processor: DataProcessorPipeline | None = None, # runs after teleop
robot_action_processor: RobotProcessor | None = None, # runs before robot robot_action_processor: DataProcessorPipeline | None = None, # runs before robot
robot_observation_processor: RobotProcessor | None = None, # runs after robot robot_observation_processor: DataProcessorPipeline | None = None, # runs after robot
single_task: str | None = None, single_task: str | None = None,
display_data: bool = False, display_data: bool = False,
): ):
teleop_action_processor = teleop_action_processor or RobotProcessor( teleop_action_processor = teleop_action_processor or DataProcessorPipeline(
steps=[IdentityProcessor()], to_transition=to_transition_teleop_action, to_output=lambda tr: tr steps=[IdentityProcessor()], to_transition=to_transition_teleop_action, to_output=lambda tr: tr
) )
robot_action_processor = robot_action_processor or RobotProcessor( robot_action_processor = robot_action_processor or DataProcessorPipeline(
steps=[IdentityProcessor()], to_transition=lambda tr: tr, to_output=to_output_robot_action steps=[IdentityProcessor()], to_transition=lambda tr: tr, to_output=to_output_robot_action
) )
robot_observation_processor = robot_observation_processor or RobotProcessor( robot_observation_processor = robot_observation_processor or DataProcessorPipeline(
steps=[IdentityProcessor()], to_transition=to_transition_robot_observation, to_output=lambda tr: tr steps=[IdentityProcessor()], to_transition=to_transition_robot_observation, to_output=lambda tr: tr
) )
+3 -3
View File
@@ -47,7 +47,7 @@ from pprint import pformat
from lerobot.configs import parser from lerobot.configs import parser
from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.processor import IdentityProcessor, RobotProcessor from lerobot.processor import DataProcessorPipeline, IdentityProcessor
from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action
from lerobot.robots import ( # noqa: F401 from lerobot.robots import ( # noqa: F401
Robot, Robot,
@@ -85,7 +85,7 @@ class ReplayConfig:
# Use vocal synthesis to read events. # Use vocal synthesis to read events.
play_sounds: bool = True play_sounds: bool = True
# Optional processor for actions before sending to robot # Optional processor for actions before sending to robot
robot_action_processor: RobotProcessor | None = None robot_action_processor: DataProcessorPipeline | None = None
@parser.wrap() @parser.wrap()
@@ -94,7 +94,7 @@ def replay(cfg: ReplayConfig):
logging.info(pformat(asdict(cfg))) logging.info(pformat(asdict(cfg)))
# Initialize robot action processor with default if not provided # Initialize robot action processor with default if not provided
robot_action_processor = cfg.robot_action_processor or RobotProcessor( robot_action_processor = cfg.robot_action_processor or DataProcessorPipeline(
steps=[IdentityProcessor()], steps=[IdentityProcessor()],
to_transition=to_transition_teleop_action, to_transition=to_transition_teleop_action,
to_output=to_output_robot_action, # type: ignore[arg-type] to_output=to_output_robot_action, # type: ignore[arg-type]
+12 -8
View File
@@ -31,6 +31,7 @@ from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import ( from lerobot.processor import (
AddTeleopActionAsComplimentaryData, AddTeleopActionAsComplimentaryData,
AddTeleopEventsAsInfo, AddTeleopEventsAsInfo,
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
EnvTransition, EnvTransition,
GripperPenaltyProcessor, GripperPenaltyProcessor,
@@ -42,7 +43,6 @@ from lerobot.processor import (
MotorCurrentProcessor, MotorCurrentProcessor,
Numpy2TorchActionProcessor, Numpy2TorchActionProcessor,
RewardClassifierProcessor, RewardClassifierProcessor,
RobotProcessor,
TimeLimitProcessor, TimeLimitProcessor,
ToBatchProcessor, ToBatchProcessor,
Torch2NumpyActionProcessor, Torch2NumpyActionProcessor,
@@ -374,7 +374,9 @@ def make_processors(
DeviceProcessor(device=device), DeviceProcessor(device=device),
] ]
return RobotProcessor(steps=env_pipeline_steps), RobotProcessor(steps=action_pipeline_steps) return DataProcessorPipeline(steps=env_pipeline_steps), DataProcessorPipeline(
steps=action_pipeline_steps
)
# Full processor pipeline for real robot environment # Full processor pipeline for real robot environment
# Get robot and motor information for kinematics # Get robot and motor information for kinematics
@@ -486,15 +488,15 @@ def make_processors(
] ]
action_pipeline_steps.extend(inverse_kinematics_steps) action_pipeline_steps.extend(inverse_kinematics_steps)
return RobotProcessor(steps=env_pipeline_steps), RobotProcessor(steps=action_pipeline_steps) return DataProcessorPipeline(steps=env_pipeline_steps), DataProcessorPipeline(steps=action_pipeline_steps)
def step_env_and_process_transition( def step_env_and_process_transition(
env: gym.Env, env: gym.Env,
transition: EnvTransition, transition: EnvTransition,
action: torch.Tensor, action: torch.Tensor,
env_processor: RobotProcessor, env_processor: DataProcessorPipeline,
action_processor: RobotProcessor, action_processor: DataProcessorPipeline,
): ):
""" """
Execute one step with processor pipeline. Execute one step with processor pipeline.
@@ -543,8 +545,8 @@ def step_env_and_process_transition(
def control_loop( def control_loop(
env: gym.Env, env: gym.Env,
env_processor: RobotProcessor, env_processor: DataProcessorPipeline,
action_processor: RobotProcessor, action_processor: DataProcessorPipeline,
teleop_device: Teleoperator, teleop_device: Teleoperator,
cfg: GymManipulatorConfig, cfg: GymManipulatorConfig,
) -> None: ) -> None:
@@ -698,7 +700,9 @@ def control_loop(
dataset.push_to_hub() dataset.push_to_hub()
def replay_trajectory(env: gym.Env, action_processor: RobotProcessor, cfg: GymManipulatorConfig) -> None: def replay_trajectory(
env: gym.Env, action_processor: DataProcessorPipeline, cfg: GymManipulatorConfig
) -> None:
"""Replay recorded trajectory on robot environment.""" """Replay recorded trajectory on robot environment."""
assert cfg.dataset.replay_episode is not None, "Replay episode must be provided for replay" assert cfg.dataset.replay_episode is not None, "Replay episode must be provided for replay"
+10 -10
View File
@@ -61,7 +61,7 @@ import rerun as rr
from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401
from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401
from lerobot.configs import parser from lerobot.configs import parser
from lerobot.processor import IdentityProcessor, RobotProcessor from lerobot.processor import DataProcessorPipeline, IdentityProcessor
from lerobot.processor.converters import ( from lerobot.processor.converters import (
to_output_robot_action, to_output_robot_action,
to_transition_robot_observation, to_transition_robot_observation,
@@ -104,9 +104,9 @@ class TeleoperateConfig:
# Display all cameras on screen # Display all cameras on screen
display_data: bool = False display_data: bool = False
# Optional processors for data transformation # Optional processors for data transformation
teleop_action_processor: RobotProcessor | None = None # runs after teleop teleop_action_processor: DataProcessorPipeline | None = None # runs after teleop
robot_action_processor: RobotProcessor | None = None # runs before robot robot_action_processor: DataProcessorPipeline | None = None # runs before robot
robot_observation_processor: RobotProcessor | None = None # runs after robot robot_observation_processor: DataProcessorPipeline | None = None # runs after robot
def teleop_loop( def teleop_loop(
@@ -115,20 +115,20 @@ def teleop_loop(
fps: int, fps: int,
display_data: bool = False, display_data: bool = False,
duration: float | None = None, duration: float | None = None,
teleop_action_processor: RobotProcessor | None = None, teleop_action_processor: DataProcessorPipeline | None = None,
robot_action_processor: RobotProcessor | None = None, robot_action_processor: DataProcessorPipeline | None = None,
robot_observation_processor: RobotProcessor | None = None, robot_observation_processor: DataProcessorPipeline | None = None,
): ):
# Initialize processors with defaults if not provided # Initialize processors with defaults if not provided
teleop_action_processor = teleop_action_processor or RobotProcessor( teleop_action_processor = teleop_action_processor or DataProcessorPipeline(
steps=[IdentityProcessor()], to_transition=to_transition_teleop_action, to_output=lambda tr: tr steps=[IdentityProcessor()], to_transition=to_transition_teleop_action, to_output=lambda tr: tr
) )
robot_action_processor = robot_action_processor or RobotProcessor( robot_action_processor = robot_action_processor or DataProcessorPipeline(
steps=[IdentityProcessor()], steps=[IdentityProcessor()],
to_transition=lambda tr: tr, to_transition=lambda tr: tr,
to_output=to_output_robot_action, # type: ignore[arg-type] to_output=to_output_robot_action, # type: ignore[arg-type]
) )
robot_observation_processor = robot_observation_processor or RobotProcessor( robot_observation_processor = robot_observation_processor or DataProcessorPipeline(
steps=[IdentityProcessor()], to_transition=to_transition_robot_observation, to_output=lambda tr: tr steps=[IdentityProcessor()], to_transition=to_transition_robot_observation, to_output=lambda tr: tr
) )
+3 -3
View File
@@ -31,7 +31,7 @@ from termcolor import colored
from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.utils import DEFAULT_FEATURES from lerobot.datasets.utils import DEFAULT_FEATURES
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.processor import RobotProcessor, TransitionKey from lerobot.processor import DataProcessorPipeline, TransitionKey
from lerobot.robots import Robot from lerobot.robots import Robot
@@ -102,8 +102,8 @@ def predict_action(
observation: dict[str, np.ndarray], observation: dict[str, np.ndarray],
policy: PreTrainedPolicy, policy: PreTrainedPolicy,
device: torch.device, device: torch.device,
preprocessor: RobotProcessor, preprocessor: DataProcessorPipeline,
postprocessor: RobotProcessor, postprocessor: DataProcessorPipeline,
use_amp: bool, use_amp: bool,
task: str | None = None, task: str | None = None,
robot_type: str | None = None, robot_type: str | None = None,
+3 -3
View File
@@ -32,7 +32,7 @@ from lerobot.datasets.utils import load_json, write_json
from lerobot.optim.optimizers import load_optimizer_state, save_optimizer_state from lerobot.optim.optimizers import load_optimizer_state, save_optimizer_state
from lerobot.optim.schedulers import load_scheduler_state, save_scheduler_state from lerobot.optim.schedulers import load_scheduler_state, save_scheduler_state
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.processor import RobotProcessor from lerobot.processor import DataProcessorPipeline
from lerobot.utils.random_utils import load_rng_state, save_rng_state from lerobot.utils.random_utils import load_rng_state, save_rng_state
@@ -75,8 +75,8 @@ def save_checkpoint(
policy: PreTrainedPolicy, policy: PreTrainedPolicy,
optimizer: Optimizer, optimizer: Optimizer,
scheduler: LRScheduler | None = None, scheduler: LRScheduler | None = None,
preprocessor: RobotProcessor | None = None, preprocessor: DataProcessorPipeline | None = None,
postprocessor: RobotProcessor | None = None, postprocessor: DataProcessorPipeline | None = None,
) -> None: ) -> None:
"""This function creates the following directory structure: """This function creates the following directory structure:
+2 -2
View File
@@ -25,10 +25,10 @@ from lerobot.constants import ACTION, OBS_STATE
from lerobot.policies.act.configuration_act import ACTConfig from lerobot.policies.act.configuration_act import ACTConfig
from lerobot.policies.act.processor_act import make_act_pre_post_processors from lerobot.policies.act.processor_act import make_act_pre_post_processors
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
NormalizerProcessor, NormalizerProcessor,
RenameProcessor, RenameProcessor,
RobotProcessor,
ToBatchProcessor, ToBatchProcessor,
TransitionKey, TransitionKey,
UnnormalizerProcessor, UnnormalizerProcessor,
@@ -250,7 +250,7 @@ def test_act_processor_save_and_load():
preprocessor.save_pretrained(tmpdir) preprocessor.save_pretrained(tmpdir)
# Load preprocessor # Load preprocessor
loaded_preprocessor = RobotProcessor.from_pretrained( loaded_preprocessor = DataProcessorPipeline.from_pretrained(
tmpdir, to_transition=lambda x: x, to_output=lambda x: x tmpdir, to_transition=lambda x: x, to_output=lambda x: x
) )
+3 -3
View File
@@ -1,6 +1,6 @@
import torch import torch
from lerobot.processor import RobotProcessor, TransitionKey from lerobot.processor import DataProcessorPipeline, TransitionKey
from lerobot.processor.converters import batch_to_transition, transition_to_batch from lerobot.processor.converters import batch_to_transition, transition_to_batch
@@ -20,7 +20,7 @@ def _dummy_batch():
def test_observation_grouping_roundtrip(): def test_observation_grouping_roundtrip():
"""Test that observation.* keys are properly grouped and ungrouped.""" """Test that observation.* keys are properly grouped and ungrouped."""
proc = RobotProcessor([]) proc = DataProcessorPipeline([])
batch_in = _dummy_batch() batch_in = _dummy_batch()
batch_out = proc(batch_in) batch_out = proc(batch_in)
@@ -261,7 +261,7 @@ def test_custom_converter():
batch = transition_to_batch(tr) batch = transition_to_batch(tr)
return batch return batch
processor = RobotProcessor(steps=[], to_transition=to_tr, to_output=to_batch) processor = DataProcessorPipeline(steps=[], to_transition=to_tr, to_output=to_batch)
batch = { batch = {
"observation.state": torch.randn(1, 4), "observation.state": torch.randn(1, 4),
+6 -6
View File
@@ -22,7 +22,7 @@ import pytest
import torch import torch
from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
from lerobot.processor import ProcessorStepRegistry, RobotProcessor, ToBatchProcessor, TransitionKey from lerobot.processor import DataProcessorPipeline, ProcessorStepRegistry, ToBatchProcessor, TransitionKey
def create_transition( def create_transition(
@@ -243,7 +243,7 @@ def test_mixed_observation():
def test_integration_with_robot_processor(): def test_integration_with_robot_processor():
"""Test ToBatchProcessor integration with RobotProcessor.""" """Test ToBatchProcessor integration with RobotProcessor."""
to_batch_processor = ToBatchProcessor() to_batch_processor = ToBatchProcessor()
pipeline = RobotProcessor([to_batch_processor], to_transition=lambda x: x, to_output=lambda x: x) pipeline = DataProcessorPipeline([to_batch_processor], to_transition=lambda x: x, to_output=lambda x: x)
# Create unbatched observation # Create unbatched observation
observation = { observation = {
@@ -283,7 +283,7 @@ def test_serialization_methods():
def test_save_and_load_pretrained(): def test_save_and_load_pretrained():
"""Test saving and loading ToBatchProcessor with RobotProcessor.""" """Test saving and loading ToBatchProcessor with RobotProcessor."""
processor = ToBatchProcessor() processor = ToBatchProcessor()
pipeline = RobotProcessor( pipeline = DataProcessorPipeline(
[processor], name="BatchPipeline", to_transition=lambda x: x, to_output=lambda x: x [processor], name="BatchPipeline", to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -296,7 +296,7 @@ def test_save_and_load_pretrained():
assert config_path.exists() assert config_path.exists()
# Load pipeline # Load pipeline
loaded_pipeline = RobotProcessor.from_pretrained( loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir, to_transition=lambda x: x, to_output=lambda x: x tmp_dir, to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -325,11 +325,11 @@ def test_registry_functionality():
def test_registry_based_save_load(): def test_registry_based_save_load():
"""Test saving and loading using registry name.""" """Test saving and loading using registry name."""
processor = ToBatchProcessor() processor = ToBatchProcessor()
pipeline = RobotProcessor([processor], to_transition=lambda x: x, to_output=lambda x: x) pipeline = DataProcessorPipeline([processor], to_transition=lambda x: x, to_output=lambda x: x)
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
loaded_pipeline = RobotProcessor.from_pretrained( loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir, to_transition=lambda x: x, to_output=lambda x: x tmp_dir, to_transition=lambda x: x, to_output=lambda x: x
) )
+4 -4
View File
@@ -25,10 +25,10 @@ from lerobot.constants import OBS_IMAGE, OBS_STATE
from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
from lerobot.policies.sac.reward_model.processor_classifier import make_classifier_processor from lerobot.policies.sac.reward_model.processor_classifier import make_classifier_processor
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
IdentityProcessor, IdentityProcessor,
NormalizerProcessor, NormalizerProcessor,
RobotProcessor,
TransitionKey, TransitionKey,
) )
@@ -254,7 +254,7 @@ def test_classifier_processor_save_and_load():
factory_preprocessor, factory_postprocessor = make_classifier_processor(config, stats) factory_preprocessor, factory_postprocessor = make_classifier_processor(config, stats)
# Create new processors with EnvTransition input/output # Create new processors with EnvTransition input/output
preprocessor = RobotProcessor( preprocessor = DataProcessorPipeline(
factory_preprocessor.steps, to_transition=lambda x: x, to_output=lambda x: x factory_preprocessor.steps, to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -263,7 +263,7 @@ def test_classifier_processor_save_and_load():
preprocessor.save_pretrained(tmpdir) preprocessor.save_pretrained(tmpdir)
# Load preprocessor # Load preprocessor
loaded_preprocessor = RobotProcessor.from_pretrained( loaded_preprocessor = DataProcessorPipeline.from_pretrained(
tmpdir, to_transition=lambda x: x, to_output=lambda x: x tmpdir, to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -300,7 +300,7 @@ def test_classifier_processor_mixed_precision():
modified_steps.append(step) modified_steps.append(step)
# Create new processors with EnvTransition input/output # Create new processors with EnvTransition input/output
preprocessor = RobotProcessor(modified_steps, to_transition=lambda x: x, to_output=lambda x: x) preprocessor = DataProcessorPipeline(modified_steps, to_transition=lambda x: x, to_output=lambda x: x)
# Create test data # Create test data
observation = { observation = {
+6 -6
View File
@@ -19,7 +19,7 @@ import pytest
import torch import torch
from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.configs.types import FeatureType, PolicyFeature
from lerobot.processor import DeviceProcessor, RobotProcessor, TransitionKey from lerobot.processor import DataProcessorPipeline, DeviceProcessor, TransitionKey
def create_transition( def create_transition(
@@ -310,7 +310,7 @@ def test_integration_with_robot_processor():
device_processor = DeviceProcessor(device="cpu") device_processor = DeviceProcessor(device="cpu")
batch_processor = ToBatchProcessor() batch_processor = ToBatchProcessor()
processor = RobotProcessor( processor = DataProcessorPipeline(
steps=[batch_processor, device_processor], steps=[batch_processor, device_processor],
name="test_pipeline", name="test_pipeline",
to_transition=lambda x: x, to_transition=lambda x: x,
@@ -336,14 +336,14 @@ def test_save_and_load_pretrained():
"""Test saving and loading processor with DeviceProcessor.""" """Test saving and loading processor with DeviceProcessor."""
device = "cuda:0" if torch.cuda.is_available() else "cpu" device = "cuda:0" if torch.cuda.is_available() else "cpu"
processor = DeviceProcessor(device=device, float_dtype="float16") processor = DeviceProcessor(device=device, float_dtype="float16")
robot_processor = RobotProcessor(steps=[processor], name="device_test_processor") robot_processor = DataProcessorPipeline(steps=[processor], name="device_test_processor")
with tempfile.TemporaryDirectory() as tmpdir: with tempfile.TemporaryDirectory() as tmpdir:
# Save # Save
robot_processor.save_pretrained(tmpdir) robot_processor.save_pretrained(tmpdir)
# Load # Load
loaded_processor = RobotProcessor.from_pretrained(tmpdir) loaded_processor = DataProcessorPipeline.from_pretrained(tmpdir)
assert len(loaded_processor.steps) == 1 assert len(loaded_processor.steps) == 1
loaded_device_processor = loaded_processor.steps[0] loaded_device_processor = loaded_processor.steps[0]
@@ -982,7 +982,7 @@ def test_policy_processor_integration():
norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD, FeatureType.ACTION: NormalizationMode.MEAN_STD} norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD, FeatureType.ACTION: NormalizationMode.MEAN_STD}
# Create input processor (preprocessor) that moves to GPU # Create input processor (preprocessor) that moves to GPU
input_processor = RobotProcessor( input_processor = DataProcessorPipeline(
steps=[ steps=[
NormalizerProcessor(features=features, norm_map=norm_map, stats=stats), NormalizerProcessor(features=features, norm_map=norm_map, stats=stats),
ToBatchProcessor(), ToBatchProcessor(),
@@ -994,7 +994,7 @@ def test_policy_processor_integration():
) )
# Create output processor (postprocessor) that moves to CPU # Create output processor (postprocessor) that moves to CPU
output_processor = RobotProcessor( output_processor = DataProcessorPipeline(
steps=[ steps=[
DeviceProcessor(device="cpu"), DeviceProcessor(device="cpu"),
UnnormalizerProcessor(features={ACTION: features[ACTION]}, norm_map=norm_map, stats=stats), UnnormalizerProcessor(features={ACTION: features[ACTION]}, norm_map=norm_map, stats=stats),
+4 -4
View File
@@ -25,10 +25,10 @@ from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE
from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
from lerobot.policies.diffusion.processor_diffusion import make_diffusion_pre_post_processors from lerobot.policies.diffusion.processor_diffusion import make_diffusion_pre_post_processors
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
NormalizerProcessor, NormalizerProcessor,
RenameProcessor, RenameProcessor,
RobotProcessor,
ToBatchProcessor, ToBatchProcessor,
TransitionKey, TransitionKey,
UnnormalizerProcessor, UnnormalizerProcessor,
@@ -257,7 +257,7 @@ def test_diffusion_processor_save_and_load():
factory_preprocessor, factory_postprocessor = make_diffusion_pre_post_processors(config, stats) factory_preprocessor, factory_postprocessor = make_diffusion_pre_post_processors(config, stats)
# Create new processors with EnvTransition input/output # Create new processors with EnvTransition input/output
preprocessor = RobotProcessor( preprocessor = DataProcessorPipeline(
factory_preprocessor.steps, to_transition=lambda x: x, to_output=lambda x: x factory_preprocessor.steps, to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -266,7 +266,7 @@ def test_diffusion_processor_save_and_load():
preprocessor.save_pretrained(tmpdir) preprocessor.save_pretrained(tmpdir)
# Load preprocessor # Load preprocessor
loaded_preprocessor = RobotProcessor.from_pretrained( loaded_preprocessor = DataProcessorPipeline.from_pretrained(
tmpdir, to_transition=lambda x: x, to_output=lambda x: x tmpdir, to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -303,7 +303,7 @@ def test_diffusion_processor_mixed_precision():
modified_steps.append(step) modified_steps.append(step)
# Create new processors with EnvTransition input/output # Create new processors with EnvTransition input/output
preprocessor = RobotProcessor(modified_steps, to_transition=lambda x: x, to_output=lambda x: x) preprocessor = DataProcessorPipeline(modified_steps, to_transition=lambda x: x, to_output=lambda x: x)
# Create test data # Create test data
observation = { observation = {
+15 -11
View File
@@ -21,9 +21,9 @@ import torch
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
IdentityProcessor, IdentityProcessor,
NormalizerProcessor, NormalizerProcessor,
RobotProcessor,
TransitionKey, TransitionKey,
UnnormalizerProcessor, UnnormalizerProcessor,
hotswap_stats, hotswap_stats,
@@ -508,7 +508,9 @@ def test_get_config(full_stats):
def test_integration_with_robot_processor(normalizer_processor): def test_integration_with_robot_processor(normalizer_processor):
"""Test integration with RobotProcessor pipeline""" """Test integration with RobotProcessor pipeline"""
robot_processor = RobotProcessor([normalizer_processor], to_transition=lambda x: x, to_output=lambda x: x) robot_processor = DataProcessorPipeline(
[normalizer_processor], to_transition=lambda x: x, to_output=lambda x: x
)
observation = { observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]), "observation.image": torch.tensor([0.7, 0.5, 0.3]),
@@ -1009,7 +1011,7 @@ def test_hotswap_stats_basic_functionality():
identity = IdentityProcessor() identity = IdentityProcessor()
# Create robot processor # Create robot processor
robot_processor = RobotProcessor(steps=[normalizer, unnormalizer, identity]) robot_processor = DataProcessorPipeline(steps=[normalizer, unnormalizer, identity])
# Hotswap stats # Hotswap stats
new_processor = hotswap_stats(robot_processor, new_stats) new_processor = hotswap_stats(robot_processor, new_stats)
@@ -1046,7 +1048,7 @@ def test_hotswap_stats_deep_copy():
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats)
original_processor = RobotProcessor(steps=[normalizer]) original_processor = DataProcessorPipeline(steps=[normalizer])
# Store reference to original stats # Store reference to original stats
original_stats_reference = original_processor.steps[0].stats original_stats_reference = original_processor.steps[0].stats
@@ -1089,7 +1091,7 @@ def test_hotswap_stats_only_affects_normalizer_steps():
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
identity = IdentityProcessor() identity = IdentityProcessor()
robot_processor = RobotProcessor(steps=[normalizer, identity, unnormalizer]) robot_processor = DataProcessorPipeline(steps=[normalizer, identity, unnormalizer])
# Hotswap stats # Hotswap stats
new_processor = hotswap_stats(robot_processor, new_stats) new_processor = hotswap_stats(robot_processor, new_stats)
@@ -1116,7 +1118,7 @@ def test_hotswap_stats_empty_stats():
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats)
robot_processor = RobotProcessor(steps=[normalizer]) robot_processor = DataProcessorPipeline(steps=[normalizer])
# Hotswap with empty stats # Hotswap with empty stats
new_processor = hotswap_stats(robot_processor, empty_stats) new_processor = hotswap_stats(robot_processor, empty_stats)
@@ -1133,7 +1135,7 @@ def test_hotswap_stats_no_normalizer_steps():
} }
# Create processor with only identity steps # Create processor with only identity steps
robot_processor = RobotProcessor(steps=[IdentityProcessor(), IdentityProcessor()]) robot_processor = DataProcessorPipeline(steps=[IdentityProcessor(), IdentityProcessor()])
# Hotswap stats - should work without error # Hotswap stats - should work without error
new_processor = hotswap_stats(robot_processor, stats) new_processor = hotswap_stats(robot_processor, stats)
@@ -1172,7 +1174,7 @@ def test_hotswap_stats_preserves_other_attributes():
normalize_observation_keys=normalize_observation_keys, normalize_observation_keys=normalize_observation_keys,
eps=eps, eps=eps,
) )
robot_processor = RobotProcessor(steps=[normalizer]) robot_processor = DataProcessorPipeline(steps=[normalizer])
# Hotswap stats # Hotswap stats
new_processor = hotswap_stats(robot_processor, new_stats) new_processor = hotswap_stats(robot_processor, new_stats)
@@ -1215,7 +1217,7 @@ def test_hotswap_stats_multiple_normalizer_types():
unnormalizer1 = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) unnormalizer1 = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats)
unnormalizer2 = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) unnormalizer2 = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats)
robot_processor = RobotProcessor(steps=[normalizer1, unnormalizer1, normalizer2, unnormalizer2]) robot_processor = DataProcessorPipeline(steps=[normalizer1, unnormalizer1, normalizer2, unnormalizer2])
# Hotswap stats # Hotswap stats
new_processor = hotswap_stats(robot_processor, new_stats) new_processor = hotswap_stats(robot_processor, new_stats)
@@ -1263,7 +1265,7 @@ def test_hotswap_stats_with_different_data_types():
} }
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats)
robot_processor = RobotProcessor(steps=[normalizer]) robot_processor = DataProcessorPipeline(steps=[normalizer])
# Hotswap stats # Hotswap stats
new_processor = hotswap_stats(robot_processor, new_stats) new_processor = hotswap_stats(robot_processor, new_stats)
@@ -1319,7 +1321,9 @@ def test_hotswap_stats_functional_test():
# Create original processor # Create original processor
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats)
original_processor = RobotProcessor(steps=[normalizer], to_transition=lambda x: x, to_output=lambda x: x) original_processor = DataProcessorPipeline(
steps=[normalizer], to_transition=lambda x: x, to_output=lambda x: x
)
# Process with original stats # Process with original stats
original_result = original_processor(transition) original_result = original_processor(transition)
+127 -119
View File
@@ -27,7 +27,7 @@ import torch.nn as nn
from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.configs.types import FeatureType, PolicyFeature
from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features
from lerobot.processor import EnvTransition, ProcessorStepRegistry, RobotProcessor, TransitionKey from lerobot.processor import DataProcessorPipeline, EnvTransition, ProcessorStepRegistry, TransitionKey
from tests.conftest import assert_contract_is_typed from tests.conftest import assert_contract_is_typed
@@ -175,7 +175,7 @@ class MockStepWithTensorState:
def test_empty_pipeline(): def test_empty_pipeline():
"""Test pipeline with no steps.""" """Test pipeline with no steps."""
pipeline = RobotProcessor([], to_transition=lambda x: x, to_output=lambda x: x) pipeline = DataProcessorPipeline([], to_transition=lambda x: x, to_output=lambda x: x)
transition = create_transition() transition = create_transition()
result = pipeline(transition) result = pipeline(transition)
@@ -187,7 +187,7 @@ def test_empty_pipeline():
def test_single_step_pipeline(): def test_single_step_pipeline():
"""Test pipeline with a single step.""" """Test pipeline with a single step."""
step = MockStep("test_step") step = MockStep("test_step")
pipeline = RobotProcessor([step], to_transition=lambda x: x, to_output=lambda x: x) pipeline = DataProcessorPipeline([step], to_transition=lambda x: x, to_output=lambda x: x)
transition = create_transition() transition = create_transition()
result = pipeline(transition) result = pipeline(transition)
@@ -204,7 +204,7 @@ def test_multiple_steps_pipeline():
"""Test pipeline with multiple steps.""" """Test pipeline with multiple steps."""
step1 = MockStep("step1") step1 = MockStep("step1")
step2 = MockStep("step2") step2 = MockStep("step2")
pipeline = RobotProcessor([step1, step2], to_transition=lambda x: x, to_output=lambda x: x) pipeline = DataProcessorPipeline([step1, step2], to_transition=lambda x: x, to_output=lambda x: x)
transition = create_transition() transition = create_transition()
result = pipeline(transition) result = pipeline(transition)
@@ -216,7 +216,7 @@ def test_multiple_steps_pipeline():
def test_invalid_transition_format(): def test_invalid_transition_format():
"""Test pipeline with invalid transition format.""" """Test pipeline with invalid transition format."""
pipeline = RobotProcessor([MockStep()]) pipeline = DataProcessorPipeline([MockStep()])
# Test with wrong type (tuple instead of dict) # Test with wrong type (tuple instead of dict)
with pytest.raises(ValueError, match="EnvTransition must be a dictionary"): with pytest.raises(ValueError, match="EnvTransition must be a dictionary"):
@@ -231,7 +231,7 @@ def test_step_through():
"""Test step_through method with dict input.""" """Test step_through method with dict input."""
step1 = MockStep("step1") step1 = MockStep("step1")
step2 = MockStep("step2") step2 = MockStep("step2")
pipeline = RobotProcessor([step1, step2]) pipeline = DataProcessorPipeline([step1, step2])
transition = create_transition() transition = create_transition()
@@ -252,7 +252,7 @@ def test_step_through_with_dict():
"""Test step_through method with dict input.""" """Test step_through method with dict input."""
step1 = MockStep("step1") step1 = MockStep("step1")
step2 = MockStep("step2") step2 = MockStep("step2")
pipeline = RobotProcessor([step1, step2]) pipeline = DataProcessorPipeline([step1, step2])
batch = { batch = {
"observation.image": None, "observation.image": None,
@@ -291,7 +291,7 @@ def test_step_through_with_dict():
def test_step_through_no_hooks(): def test_step_through_no_hooks():
"""Test that step_through doesn't execute hooks.""" """Test that step_through doesn't execute hooks."""
step = MockStep("test_step") step = MockStep("test_step")
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
hook_calls = [] hook_calls = []
@@ -326,7 +326,7 @@ def test_indexing():
"""Test pipeline indexing.""" """Test pipeline indexing."""
step1 = MockStep("step1") step1 = MockStep("step1")
step2 = MockStep("step2") step2 = MockStep("step2")
pipeline = RobotProcessor([step1, step2]) pipeline = DataProcessorPipeline([step1, step2])
# Test integer indexing # Test integer indexing
assert pipeline[0] is step1 assert pipeline[0] is step1
@@ -334,7 +334,7 @@ def test_indexing():
# Test slice indexing # Test slice indexing
sub_pipeline = pipeline[0:1] sub_pipeline = pipeline[0:1]
assert isinstance(sub_pipeline, RobotProcessor) assert isinstance(sub_pipeline, DataProcessorPipeline)
assert len(sub_pipeline) == 1 assert len(sub_pipeline) == 1
assert sub_pipeline[0] is step1 assert sub_pipeline[0] is step1
@@ -342,7 +342,7 @@ def test_indexing():
def test_hooks(): def test_hooks():
"""Test before/after step hooks.""" """Test before/after step hooks."""
step = MockStep("test_step") step = MockStep("test_step")
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
before_calls = [] before_calls = []
after_calls = [] after_calls = []
@@ -366,7 +366,7 @@ def test_hooks():
def test_unregister_hooks(): def test_unregister_hooks():
"""Test unregistering hooks from the pipeline.""" """Test unregistering hooks from the pipeline."""
step = MockStep("test_step") step = MockStep("test_step")
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
# Test before_step_hook # Test before_step_hook
before_calls = [] before_calls = []
@@ -405,7 +405,7 @@ def test_unregister_hooks():
def test_unregister_nonexistent_hook(): def test_unregister_nonexistent_hook():
"""Test error handling when unregistering hooks that don't exist.""" """Test error handling when unregistering hooks that don't exist."""
pipeline = RobotProcessor([MockStep()]) pipeline = DataProcessorPipeline([MockStep()])
def some_hook(idx: int, transition: EnvTransition): def some_hook(idx: int, transition: EnvTransition):
pass pass
@@ -423,7 +423,7 @@ def test_unregister_nonexistent_hook():
def test_multiple_hooks_and_selective_unregister(): def test_multiple_hooks_and_selective_unregister():
"""Test registering multiple hooks and selectively unregistering them.""" """Test registering multiple hooks and selectively unregistering them."""
pipeline = RobotProcessor([MockStep("step1"), MockStep("step2")]) pipeline = DataProcessorPipeline([MockStep("step1"), MockStep("step2")])
calls_1 = [] calls_1 = []
calls_2 = [] calls_2 = []
@@ -469,7 +469,7 @@ def test_multiple_hooks_and_selective_unregister():
def test_hook_execution_order_documentation(): def test_hook_execution_order_documentation():
"""Test and document that hooks are executed sequentially in registration order.""" """Test and document that hooks are executed sequentially in registration order."""
pipeline = RobotProcessor([MockStep("step")]) pipeline = DataProcessorPipeline([MockStep("step")])
execution_order = [] execution_order = []
@@ -521,7 +521,7 @@ def test_save_and_load_pretrained():
step1.counter = 5 step1.counter = 5
step2.counter = 10 step2.counter = 10
pipeline = RobotProcessor([step1, step2], name="TestPipeline") pipeline = DataProcessorPipeline([step1, step2], name="TestPipeline")
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
# Save pipeline # Save pipeline
@@ -543,7 +543,7 @@ def test_save_and_load_pretrained():
assert config["steps"][1]["config"]["counter"] == 10 assert config["steps"][1]["config"]["counter"] == 10
# Load pipeline # Load pipeline
loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir) loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir)
assert loaded_pipeline.name == "TestPipeline" assert loaded_pipeline.name == "TestPipeline"
assert len(loaded_pipeline) == 2 assert len(loaded_pipeline) == 2
@@ -556,7 +556,7 @@ def test_save_and_load_pretrained():
def test_step_without_optional_methods(): def test_step_without_optional_methods():
"""Test pipeline with steps that don't implement optional methods.""" """Test pipeline with steps that don't implement optional methods."""
step = MockStepWithoutOptionalMethods(multiplier=3.0) step = MockStepWithoutOptionalMethods(multiplier=3.0)
pipeline = RobotProcessor( pipeline = DataProcessorPipeline(
[step], to_transition=lambda x: x, to_output=lambda x: x [step], to_transition=lambda x: x, to_output=lambda x: x
) # Identity for EnvTransition input/output ) # Identity for EnvTransition input/output
@@ -571,14 +571,14 @@ def test_step_without_optional_methods():
# Save/load should work even without optional methods # Save/load should work even without optional methods
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir) loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir)
assert len(loaded_pipeline) == 1 assert len(loaded_pipeline) == 1
def test_mixed_json_and_tensor_state(): def test_mixed_json_and_tensor_state():
"""Test step with both JSON attributes and tensor state.""" """Test step with both JSON attributes and tensor state."""
step = MockStepWithTensorState(name="stats", learning_rate=0.05, window_size=5) step = MockStepWithTensorState(name="stats", learning_rate=0.05, window_size=5)
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
# Process some transitions with rewards # Process some transitions with rewards
for i in range(10): for i in range(10):
@@ -594,13 +594,13 @@ def test_mixed_json_and_tensor_state():
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
# Check that both config and state files were created # Check that both config and state files were created
config_path = Path(tmp_dir) / "robotprocessor.json" # Default name is "RobotProcessor" config_path = Path(tmp_dir) / "dataprocessorpipeline.json" # Default name is "RobotProcessor"
state_path = Path(tmp_dir) / "robotprocessor_step_0.safetensors" state_path = Path(tmp_dir) / "dataprocessorpipeline_step_0.safetensors"
assert config_path.exists() assert config_path.exists()
assert state_path.exists() assert state_path.exists()
# Load and verify # Load and verify
loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir) loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir)
loaded_step = loaded_pipeline.steps[0] loaded_step = loaded_pipeline.steps[0]
# Check JSON attributes were restored # Check JSON attributes were restored
@@ -861,7 +861,7 @@ def test_from_pretrained_with_overrides():
env_step = MockStepWithNonSerializableParam(name="env_step", multiplier=2.0) env_step = MockStepWithNonSerializableParam(name="env_step", multiplier=2.0)
registered_step = RegisteredMockStep(value=100, device="cpu") registered_step = RegisteredMockStep(value=100, device="cpu")
pipeline = RobotProcessor([env_step, registered_step], name="TestOverrides") pipeline = DataProcessorPipeline([env_step, registered_step], name="TestOverrides")
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
# Save the pipeline # Save the pipeline
@@ -879,7 +879,7 @@ def test_from_pretrained_with_overrides():
"registered_mock_step": {"device": "cuda", "value": 200}, "registered_mock_step": {"device": "cuda", "value": 200},
} }
loaded_pipeline = RobotProcessor.from_pretrained( loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir, overrides=overrides, to_transition=lambda x: x, to_output=lambda x: x tmp_dir, overrides=overrides, to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -907,7 +907,7 @@ def test_from_pretrained_with_partial_overrides():
step1 = MockStepWithNonSerializableParam(name="step1", multiplier=1.0) step1 = MockStepWithNonSerializableParam(name="step1", multiplier=1.0)
step2 = MockStepWithNonSerializableParam(name="step2", multiplier=2.0) step2 = MockStepWithNonSerializableParam(name="step2", multiplier=2.0)
pipeline = RobotProcessor([step1, step2]) pipeline = DataProcessorPipeline([step1, step2])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
@@ -917,7 +917,7 @@ def test_from_pretrained_with_partial_overrides():
# The current implementation applies overrides to ALL steps with the same class name # The current implementation applies overrides to ALL steps with the same class name
# Both steps will get the override # Both steps will get the override
loaded_pipeline = RobotProcessor.from_pretrained( loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir, overrides=overrides, to_transition=lambda x: x, to_output=lambda x: x tmp_dir, overrides=overrides, to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -933,7 +933,7 @@ def test_from_pretrained_with_partial_overrides():
def test_from_pretrained_invalid_override_key(): def test_from_pretrained_invalid_override_key():
"""Test that invalid override keys raise KeyError.""" """Test that invalid override keys raise KeyError."""
step = MockStepWithNonSerializableParam() step = MockStepWithNonSerializableParam()
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
@@ -942,13 +942,13 @@ def test_from_pretrained_invalid_override_key():
overrides = {"NonExistentStep": {"param": "value"}} overrides = {"NonExistentStep": {"param": "value"}}
with pytest.raises(KeyError, match="Override keys.*do not match any step"): with pytest.raises(KeyError, match="Override keys.*do not match any step"):
RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
def test_from_pretrained_multiple_invalid_override_keys(): def test_from_pretrained_multiple_invalid_override_keys():
"""Test that multiple invalid override keys are reported.""" """Test that multiple invalid override keys are reported."""
step = MockStepWithNonSerializableParam() step = MockStepWithNonSerializableParam()
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
@@ -957,7 +957,7 @@ def test_from_pretrained_multiple_invalid_override_keys():
overrides = {"NonExistentStep1": {"param": "value1"}, "NonExistentStep2": {"param": "value2"}} overrides = {"NonExistentStep1": {"param": "value1"}, "NonExistentStep2": {"param": "value2"}}
with pytest.raises(KeyError) as exc_info: with pytest.raises(KeyError) as exc_info:
RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
error_msg = str(exc_info.value) error_msg = str(exc_info.value)
assert "NonExistentStep1" in error_msg assert "NonExistentStep1" in error_msg
@@ -968,7 +968,7 @@ def test_from_pretrained_multiple_invalid_override_keys():
def test_from_pretrained_registered_step_override(): def test_from_pretrained_registered_step_override():
"""Test overriding registered steps using registry names.""" """Test overriding registered steps using registry names."""
registered_step = RegisteredMockStep(value=50, device="cpu") registered_step = RegisteredMockStep(value=50, device="cpu")
pipeline = RobotProcessor([registered_step]) pipeline = DataProcessorPipeline([registered_step])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
@@ -976,7 +976,7 @@ def test_from_pretrained_registered_step_override():
# Override using registry name # Override using registry name
overrides = {"registered_mock_step": {"value": 999, "device": "cuda"}} overrides = {"registered_mock_step": {"value": 999, "device": "cuda"}}
loaded_pipeline = RobotProcessor.from_pretrained( loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir, overrides=overrides, to_transition=lambda x: x, to_output=lambda x: x tmp_dir, overrides=overrides, to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -994,7 +994,7 @@ def test_from_pretrained_mixed_registered_and_unregistered():
unregistered_step = MockStepWithNonSerializableParam(name="unregistered", multiplier=1.0) unregistered_step = MockStepWithNonSerializableParam(name="unregistered", multiplier=1.0)
registered_step = RegisteredMockStep(value=10, device="cpu") registered_step = RegisteredMockStep(value=10, device="cpu")
pipeline = RobotProcessor([unregistered_step, registered_step]) pipeline = DataProcessorPipeline([unregistered_step, registered_step])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
@@ -1006,7 +1006,7 @@ def test_from_pretrained_mixed_registered_and_unregistered():
"registered_mock_step": {"value": 777}, "registered_mock_step": {"value": 777},
} }
loaded_pipeline = RobotProcessor.from_pretrained( loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir, overrides=overrides, to_transition=lambda x: x, to_output=lambda x: x tmp_dir, overrides=overrides, to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -1023,13 +1023,13 @@ def test_from_pretrained_mixed_registered_and_unregistered():
def test_from_pretrained_no_overrides(): def test_from_pretrained_no_overrides():
"""Test that from_pretrained works without overrides (backward compatibility).""" """Test that from_pretrained works without overrides (backward compatibility)."""
step = MockStepWithNonSerializableParam(name="no_override", multiplier=3.0) step = MockStepWithNonSerializableParam(name="no_override", multiplier=3.0)
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
# Load without overrides # Load without overrides
loaded_pipeline = RobotProcessor.from_pretrained( loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir, to_transition=lambda x: x, to_output=lambda x: x tmp_dir, to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -1045,13 +1045,13 @@ def test_from_pretrained_no_overrides():
def test_from_pretrained_empty_overrides(): def test_from_pretrained_empty_overrides():
"""Test that from_pretrained works with empty overrides dict.""" """Test that from_pretrained works with empty overrides dict."""
step = MockStepWithNonSerializableParam(multiplier=2.0) step = MockStepWithNonSerializableParam(multiplier=2.0)
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
# Load with empty overrides # Load with empty overrides
loaded_pipeline = RobotProcessor.from_pretrained( loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir, overrides={}, to_transition=lambda x: x, to_output=lambda x: x tmp_dir, overrides={}, to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -1067,7 +1067,7 @@ def test_from_pretrained_empty_overrides():
def test_from_pretrained_override_instantiation_error(): def test_from_pretrained_override_instantiation_error():
"""Test that instantiation errors with overrides are properly reported.""" """Test that instantiation errors with overrides are properly reported."""
step = MockStepWithNonSerializableParam(multiplier=1.0) step = MockStepWithNonSerializableParam(multiplier=1.0)
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
@@ -1080,13 +1080,13 @@ def test_from_pretrained_override_instantiation_error():
} }
with pytest.raises(ValueError, match="Failed to instantiate processor step"): with pytest.raises(ValueError, match="Failed to instantiate processor step"):
RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
def test_from_pretrained_with_state_and_overrides(): def test_from_pretrained_with_state_and_overrides():
"""Test that overrides work correctly with steps that have tensor state.""" """Test that overrides work correctly with steps that have tensor state."""
step = MockStepWithTensorState(name="tensor_step", learning_rate=0.01, window_size=5) step = MockStepWithTensorState(name="tensor_step", learning_rate=0.01, window_size=5)
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
# Process some data to create state # Process some data to create state
for i in range(10): for i in range(10):
@@ -1104,7 +1104,7 @@ def test_from_pretrained_with_state_and_overrides():
} }
} }
loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
loaded_step = loaded_pipeline.steps[0] loaded_step = loaded_pipeline.steps[0]
# Check that config overrides were applied # Check that config overrides were applied
@@ -1123,7 +1123,7 @@ def test_from_pretrained_override_error_messages():
"""Test that error messages for override failures are helpful.""" """Test that error messages for override failures are helpful."""
step1 = MockStepWithNonSerializableParam(name="step1") step1 = MockStepWithNonSerializableParam(name="step1")
step2 = RegisteredMockStep() step2 = RegisteredMockStep()
pipeline = RobotProcessor([step1, step2]) pipeline = DataProcessorPipeline([step1, step2])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
@@ -1132,7 +1132,7 @@ def test_from_pretrained_override_error_messages():
overrides = {"WrongStepName": {"param": "value"}} overrides = {"WrongStepName": {"param": "value"}}
with pytest.raises(KeyError) as exc_info: with pytest.raises(KeyError) as exc_info:
RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
error_msg = str(exc_info.value) error_msg = str(exc_info.value)
assert "WrongStepName" in error_msg assert "WrongStepName" in error_msg
@@ -1143,20 +1143,20 @@ def test_from_pretrained_override_error_messages():
def test_repr_empty_processor(): def test_repr_empty_processor():
"""Test __repr__ with empty processor.""" """Test __repr__ with empty processor."""
pipeline = RobotProcessor() pipeline = DataProcessorPipeline()
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='RobotProcessor', steps=0: [])" expected = "DataProcessorPipeline(name='DataProcessorPipeline', steps=0: [])"
assert repr_str == expected assert repr_str == expected
def test_repr_single_step(): def test_repr_single_step():
"""Test __repr__ with single step.""" """Test __repr__ with single step."""
step = MockStep("test_step") step = MockStep("test_step")
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='RobotProcessor', steps=1: [MockStep])" expected = "DataProcessorPipeline(name='DataProcessorPipeline', steps=1: [MockStep])"
assert repr_str == expected assert repr_str == expected
@@ -1164,18 +1164,18 @@ def test_repr_multiple_steps_under_limit():
"""Test __repr__ with 2-3 steps (all shown).""" """Test __repr__ with 2-3 steps (all shown)."""
step1 = MockStep("step1") step1 = MockStep("step1")
step2 = MockStepWithoutOptionalMethods() step2 = MockStepWithoutOptionalMethods()
pipeline = RobotProcessor([step1, step2]) pipeline = DataProcessorPipeline([step1, step2])
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='RobotProcessor', steps=2: [MockStep, MockStepWithoutOptionalMethods])" expected = "DataProcessorPipeline(name='DataProcessorPipeline', steps=2: [MockStep, MockStepWithoutOptionalMethods])"
assert repr_str == expected assert repr_str == expected
# Test with 3 steps (boundary case) # Test with 3 steps (boundary case)
step3 = MockStepWithTensorState() step3 = MockStepWithTensorState()
pipeline = RobotProcessor([step1, step2, step3]) pipeline = DataProcessorPipeline([step1, step2, step3])
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='RobotProcessor', steps=3: [MockStep, MockStepWithoutOptionalMethods, MockStepWithTensorState])" expected = "DataProcessorPipeline(name='DataProcessorPipeline', steps=3: [MockStep, MockStepWithoutOptionalMethods, MockStepWithTensorState])"
assert repr_str == expected assert repr_str == expected
@@ -1187,30 +1187,30 @@ def test_repr_many_steps_truncated():
step4 = MockModuleStep() step4 = MockModuleStep()
step5 = MockNonModuleStepWithState() step5 = MockNonModuleStepWithState()
pipeline = RobotProcessor([step1, step2, step3, step4, step5]) pipeline = DataProcessorPipeline([step1, step2, step3, step4, step5])
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='RobotProcessor', steps=5: [MockStep, MockStepWithoutOptionalMethods, ..., MockNonModuleStepWithState])" expected = "DataProcessorPipeline(name='DataProcessorPipeline', steps=5: [MockStep, MockStepWithoutOptionalMethods, ..., MockNonModuleStepWithState])"
assert repr_str == expected assert repr_str == expected
def test_repr_with_custom_name(): def test_repr_with_custom_name():
"""Test __repr__ with custom processor name.""" """Test __repr__ with custom processor name."""
step = MockStep("test_step") step = MockStep("test_step")
pipeline = RobotProcessor([step], name="CustomProcessor") pipeline = DataProcessorPipeline([step], name="CustomProcessor")
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='CustomProcessor', steps=1: [MockStep])" expected = "DataProcessorPipeline(name='CustomProcessor', steps=1: [MockStep])"
assert repr_str == expected assert repr_str == expected
def test_repr_with_seed(): def test_repr_with_seed():
"""Test __repr__ with seed parameter.""" """Test __repr__ with seed parameter."""
step = MockStep("test_step") step = MockStep("test_step")
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='RobotProcessor', steps=1: [MockStep])" expected = "DataProcessorPipeline(name='DataProcessorPipeline', steps=1: [MockStep])"
assert repr_str == expected assert repr_str == expected
@@ -1218,20 +1218,22 @@ def test_repr_with_custom_name_and_seed():
"""Test __repr__ with both custom name and seed.""" """Test __repr__ with both custom name and seed."""
step1 = MockStep("step1") step1 = MockStep("step1")
step2 = MockStepWithoutOptionalMethods() step2 = MockStepWithoutOptionalMethods()
pipeline = RobotProcessor([step1, step2], name="MyProcessor") pipeline = DataProcessorPipeline([step1, step2], name="MyProcessor")
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='MyProcessor', steps=2: [MockStep, MockStepWithoutOptionalMethods])" expected = (
"DataProcessorPipeline(name='MyProcessor', steps=2: [MockStep, MockStepWithoutOptionalMethods])"
)
assert repr_str == expected assert repr_str == expected
def test_repr_without_seed(): def test_repr_without_seed():
"""Test __repr__ when seed is explicitly None (should not show seed).""" """Test __repr__ when seed is explicitly None (should not show seed)."""
step = MockStep("test_step") step = MockStep("test_step")
pipeline = RobotProcessor([step], name="TestProcessor") pipeline = DataProcessorPipeline([step], name="TestProcessor")
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='TestProcessor', steps=1: [MockStep])" expected = "DataProcessorPipeline(name='TestProcessor', steps=1: [MockStep])"
assert repr_str == expected assert repr_str == expected
@@ -1242,10 +1244,10 @@ def test_repr_various_step_types():
step3 = MockModuleStep() step3 = MockModuleStep()
step4 = MockNonModuleStepWithState() step4 = MockNonModuleStepWithState()
pipeline = RobotProcessor([step1, step2, step3, step4], name="MixedSteps") pipeline = DataProcessorPipeline([step1, step2, step3, step4], name="MixedSteps")
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='MixedSteps', steps=4: [MockStep, MockStepWithTensorState, ..., MockNonModuleStepWithState])" expected = "DataProcessorPipeline(name='MixedSteps', steps=4: [MockStep, MockStepWithTensorState, ..., MockNonModuleStepWithState])"
assert repr_str == expected assert repr_str == expected
@@ -1256,10 +1258,10 @@ def test_repr_edge_case_long_names():
step3 = MockStepWithTensorState() step3 = MockStepWithTensorState()
step4 = MockNonModuleStepWithState() step4 = MockNonModuleStepWithState()
pipeline = RobotProcessor([step1, step2, step3, step4], name="LongNames") pipeline = DataProcessorPipeline([step1, step2, step3, step4], name="LongNames")
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='LongNames', steps=4: [MockStepWithNonSerializableParam, MockStepWithoutOptionalMethods, ..., MockNonModuleStepWithState])" expected = "DataProcessorPipeline(name='LongNames', steps=4: [MockStepWithNonSerializableParam, MockStepWithoutOptionalMethods, ..., MockNonModuleStepWithState])"
assert repr_str == expected assert repr_str == expected
@@ -1267,7 +1269,7 @@ def test_repr_edge_case_long_names():
def test_save_with_custom_config_filename(): def test_save_with_custom_config_filename():
"""Test saving processor with custom config filename.""" """Test saving processor with custom config filename."""
step = MockStep("test") step = MockStep("test")
pipeline = RobotProcessor([step], name="TestProcessor") pipeline = DataProcessorPipeline([step], name="TestProcessor")
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
# Save with custom filename # Save with custom filename
@@ -1283,16 +1285,18 @@ def test_save_with_custom_config_filename():
assert config["name"] == "TestProcessor" assert config["name"] == "TestProcessor"
# Load with specific filename # Load with specific filename
loaded = RobotProcessor.from_pretrained(tmp_dir, config_filename="my_custom_config.json") loaded = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="my_custom_config.json")
assert loaded.name == "TestProcessor" assert loaded.name == "TestProcessor"
def test_multiple_processors_same_directory(): def test_multiple_processors_same_directory():
"""Test saving multiple processors to the same directory with different config files.""" """Test saving multiple processors to the same directory with different config files."""
# Create different processors # Create different processors
preprocessor = RobotProcessor([MockStep("pre1"), MockStep("pre2")], name="preprocessor") preprocessor = DataProcessorPipeline([MockStep("pre1"), MockStep("pre2")], name="preprocessor")
postprocessor = RobotProcessor([MockStepWithoutOptionalMethods(multiplier=0.5)], name="postprocessor") postprocessor = DataProcessorPipeline(
[MockStepWithoutOptionalMethods(multiplier=0.5)], name="postprocessor"
)
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
# Save both to same directory # Save both to same directory
@@ -1304,8 +1308,8 @@ def test_multiple_processors_same_directory():
assert (Path(tmp_dir) / "postprocessor.json").exists() assert (Path(tmp_dir) / "postprocessor.json").exists()
# Load them back # Load them back
loaded_pre = RobotProcessor.from_pretrained(tmp_dir, config_filename="preprocessor.json") loaded_pre = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="preprocessor.json")
loaded_post = RobotProcessor.from_pretrained(tmp_dir, config_filename="postprocessor.json") loaded_post = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="postprocessor.json")
assert loaded_pre.name == "preprocessor" assert loaded_pre.name == "preprocessor"
assert loaded_post.name == "postprocessor" assert loaded_post.name == "postprocessor"
@@ -1316,20 +1320,20 @@ def test_multiple_processors_same_directory():
def test_auto_detect_single_config(): def test_auto_detect_single_config():
"""Test automatic config detection when there's only one JSON file.""" """Test automatic config detection when there's only one JSON file."""
step = MockStepWithTensorState() step = MockStepWithTensorState()
pipeline = RobotProcessor([step], name="SingleConfig") pipeline = DataProcessorPipeline([step], name="SingleConfig")
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
# Load without specifying config_filename # Load without specifying config_filename
loaded = RobotProcessor.from_pretrained(tmp_dir) loaded = DataProcessorPipeline.from_pretrained(tmp_dir)
assert loaded.name == "SingleConfig" assert loaded.name == "SingleConfig"
def test_error_multiple_configs_no_filename(): def test_error_multiple_configs_no_filename():
"""Test error when multiple configs exist and no filename specified.""" """Test error when multiple configs exist and no filename specified."""
proc1 = RobotProcessor([MockStep()], name="processor1") proc1 = DataProcessorPipeline([MockStep()], name="processor1")
proc2 = RobotProcessor([MockStep()], name="processor2") proc2 = DataProcessorPipeline([MockStep()], name="processor2")
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
proc1.save_pretrained(tmp_dir) proc1.save_pretrained(tmp_dir)
@@ -1337,7 +1341,7 @@ def test_error_multiple_configs_no_filename():
# Should raise error # Should raise error
with pytest.raises(ValueError, match="Multiple .json files found"): with pytest.raises(ValueError, match="Multiple .json files found"):
RobotProcessor.from_pretrained(tmp_dir) DataProcessorPipeline.from_pretrained(tmp_dir)
def test_state_file_naming_with_indices(): def test_state_file_naming_with_indices():
@@ -1347,7 +1351,7 @@ def test_state_file_naming_with_indices():
step2 = MockStepWithTensorState(name="norm2", window_size=10) step2 = MockStepWithTensorState(name="norm2", window_size=10)
step3 = MockModuleStep(input_dim=5) step3 = MockModuleStep(input_dim=5)
pipeline = RobotProcessor([step1, step2, step3]) pipeline = DataProcessorPipeline([step1, step2, step3])
# Process some data to create state # Process some data to create state
for i in range(5): for i in range(5):
@@ -1363,9 +1367,9 @@ def test_state_file_naming_with_indices():
# Files should be named with pipeline name prefix and indices # Files should be named with pipeline name prefix and indices
expected_names = [ expected_names = [
"robotprocessor_step_0.safetensors", "dataprocessorpipeline_step_0.safetensors",
"robotprocessor_step_1.safetensors", "dataprocessorpipeline_step_1.safetensors",
"robotprocessor_step_2.safetensors", "dataprocessorpipeline_step_2.safetensors",
] ]
actual_names = [f.name for f in state_files] actual_names = [f.name for f in state_files]
assert actual_names == expected_names assert actual_names == expected_names
@@ -1404,7 +1408,7 @@ def test_state_file_naming_with_registry():
# Create pipeline with registered steps # Create pipeline with registered steps
step1 = TestStatefulStep(1) step1 = TestStatefulStep(1)
step2 = TestStatefulStep(2) step2 = TestStatefulStep(2)
pipeline = RobotProcessor([step1, step2]) pipeline = DataProcessorPipeline([step1, step2])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
@@ -1415,8 +1419,8 @@ def test_state_file_naming_with_registry():
# Should include pipeline name, index and registry name # Should include pipeline name, index and registry name
expected_names = [ expected_names = [
"robotprocessor_step_0_test_stateful_step.safetensors", "dataprocessorpipeline_step_0_test_stateful_step.safetensors",
"robotprocessor_step_1_test_stateful_step.safetensors", "dataprocessorpipeline_step_1_test_stateful_step.safetensors",
] ]
actual_names = [f.name for f in state_files] actual_names = [f.name for f in state_files]
assert actual_names == expected_names assert actual_names == expected_names
@@ -1459,13 +1463,13 @@ def test_override_with_nested_config():
try: try:
step = ComplexConfigStep() step = ComplexConfigStep()
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
# Load with nested override # Load with nested override
loaded = RobotProcessor.from_pretrained( loaded = DataProcessorPipeline.from_pretrained(
tmp_dir, tmp_dir,
overrides={"complex_config_step": {"nested_config": {"level1": {"level2": "overridden"}}}}, overrides={"complex_config_step": {"nested_config": {"level1": {"level2": "overridden"}}}},
to_transition=lambda x: x, to_transition=lambda x: x,
@@ -1483,13 +1487,13 @@ def test_override_with_nested_config():
def test_override_preserves_defaults(): def test_override_preserves_defaults():
"""Test that overrides only affect specified parameters.""" """Test that overrides only affect specified parameters."""
step = MockStepWithNonSerializableParam(name="test", multiplier=2.0) step = MockStepWithNonSerializableParam(name="test", multiplier=2.0)
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
# Override only one parameter # Override only one parameter
loaded = RobotProcessor.from_pretrained( loaded = DataProcessorPipeline.from_pretrained(
tmp_dir, tmp_dir,
overrides={ overrides={
"MockStepWithNonSerializableParam": { "MockStepWithNonSerializableParam": {
@@ -1507,7 +1511,7 @@ def test_override_preserves_defaults():
def test_override_type_validation(): def test_override_type_validation():
"""Test that type errors in overrides are caught properly.""" """Test that type errors in overrides are caught properly."""
step = MockStepWithTensorState(learning_rate=0.01) step = MockStepWithTensorState(learning_rate=0.01)
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
@@ -1520,7 +1524,7 @@ def test_override_type_validation():
} }
with pytest.raises(ValueError, match="Failed to instantiate"): with pytest.raises(ValueError, match="Failed to instantiate"):
RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
def test_override_with_callables(): def test_override_with_callables():
@@ -1553,7 +1557,7 @@ def test_override_with_callables():
try: try:
step = CallableStep() step = CallableStep()
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
@@ -1567,7 +1571,7 @@ def test_override_with_callables():
return x return x
# Load with callable override # Load with callable override
loaded = RobotProcessor.from_pretrained( loaded = DataProcessorPipeline.from_pretrained(
tmp_dir, tmp_dir,
overrides={"callable_step": {"transform_fn": double_values}}, overrides={"callable_step": {"transform_fn": double_values}},
to_transition=lambda x: x, to_transition=lambda x: x,
@@ -1586,13 +1590,13 @@ def test_override_multiple_same_class_warning():
"""Test behavior when multiple steps of same class exist.""" """Test behavior when multiple steps of same class exist."""
step1 = MockStepWithNonSerializableParam(name="step1", multiplier=1.0) step1 = MockStepWithNonSerializableParam(name="step1", multiplier=1.0)
step2 = MockStepWithNonSerializableParam(name="step2", multiplier=2.0) step2 = MockStepWithNonSerializableParam(name="step2", multiplier=2.0)
pipeline = RobotProcessor([step1, step2]) pipeline = DataProcessorPipeline([step1, step2])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
# Override affects all instances of the class # Override affects all instances of the class
loaded = RobotProcessor.from_pretrained( loaded = DataProcessorPipeline.from_pretrained(
tmp_dir, overrides={"MockStepWithNonSerializableParam": {"multiplier": 10.0}} tmp_dir, overrides={"MockStepWithNonSerializableParam": {"multiplier": 10.0}}
) )
@@ -1608,7 +1612,7 @@ def test_override_multiple_same_class_warning():
def test_config_filename_special_characters(): def test_config_filename_special_characters():
"""Test config filenames with special characters are sanitized.""" """Test config filenames with special characters are sanitized."""
# Processor name with special characters # Processor name with special characters
pipeline = RobotProcessor([MockStep()], name="My/Processor\\With:Special*Chars") pipeline = DataProcessorPipeline([MockStep()], name="My/Processor\\With:Special*Chars")
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
@@ -1626,10 +1630,10 @@ def test_state_file_naming_with_multiple_processors():
"""Test that state files are properly prefixed with pipeline names to avoid conflicts.""" """Test that state files are properly prefixed with pipeline names to avoid conflicts."""
# Create two processors with state # Create two processors with state
step1 = MockStepWithTensorState(name="norm", window_size=5) step1 = MockStepWithTensorState(name="norm", window_size=5)
preprocessor = RobotProcessor([step1], name="PreProcessor") preprocessor = DataProcessorPipeline([step1], name="PreProcessor")
step2 = MockStepWithTensorState(name="norm", window_size=10) step2 = MockStepWithTensorState(name="norm", window_size=10)
postprocessor = RobotProcessor([step2], name="PostProcessor") postprocessor = DataProcessorPipeline([step2], name="PostProcessor")
# Process some data to create state # Process some data to create state
for i in range(3): for i in range(3):
@@ -1649,8 +1653,8 @@ def test_state_file_naming_with_multiple_processors():
assert (Path(tmp_dir) / "postprocessor_step_0.safetensors").exists() assert (Path(tmp_dir) / "postprocessor_step_0.safetensors").exists()
# Load both back and verify they work correctly # Load both back and verify they work correctly
loaded_pre = RobotProcessor.from_pretrained(tmp_dir, config_filename="preprocessor.json") loaded_pre = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="preprocessor.json")
loaded_post = RobotProcessor.from_pretrained(tmp_dir, config_filename="postprocessor.json") loaded_post = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="postprocessor.json")
assert loaded_pre.name == "PreProcessor" assert loaded_pre.name == "PreProcessor"
assert loaded_post.name == "PostProcessor" assert loaded_post.name == "PostProcessor"
@@ -1688,14 +1692,14 @@ def test_override_with_device_strings():
try: try:
step = DeviceAwareStep(device="cpu") step = DeviceAwareStep(device="cpu")
pipeline = RobotProcessor([step]) pipeline = DataProcessorPipeline([step])
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
# Override device # Override device
if torch.cuda.is_available(): if torch.cuda.is_available():
loaded = RobotProcessor.from_pretrained( loaded = DataProcessorPipeline.from_pretrained(
tmp_dir, overrides={"device_aware_step": {"device": "cuda:0"}} tmp_dir, overrides={"device_aware_step": {"device": "cuda:0"}}
) )
@@ -1714,16 +1718,16 @@ def test_from_pretrained_nonexistent_path():
# Test with an invalid repo ID (too many slashes) - caught by HF validation # Test with an invalid repo ID (too many slashes) - caught by HF validation
with pytest.raises(HFValidationError): with pytest.raises(HFValidationError):
RobotProcessor.from_pretrained("/path/that/does/not/exist") DataProcessorPipeline.from_pretrained("/path/that/does/not/exist")
# Test with a non-existent but valid Hub repo format # Test with a non-existent but valid Hub repo format
with pytest.raises((FileNotFoundError, HfHubHTTPError)): with pytest.raises((FileNotFoundError, HfHubHTTPError)):
RobotProcessor.from_pretrained("nonexistent-user/nonexistent-repo") DataProcessorPipeline.from_pretrained("nonexistent-user/nonexistent-repo")
# Test with a local directory that exists but has no config files # Test with a local directory that exists but has no config files
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
with pytest.raises(FileNotFoundError, match="No .json configuration files found"): with pytest.raises(FileNotFoundError, match="No .json configuration files found"):
RobotProcessor.from_pretrained(tmp_dir) DataProcessorPipeline.from_pretrained(tmp_dir)
def test_save_load_with_custom_converter_functions(): def test_save_load_with_custom_converter_functions():
@@ -1752,13 +1756,15 @@ def test_save_load_with_custom_converter_functions():
} }
# Create processor with custom converters # Create processor with custom converters
pipeline = RobotProcessor([MockStep()], to_transition=custom_to_transition, to_output=custom_to_output) pipeline = DataProcessorPipeline(
[MockStep()], to_transition=custom_to_transition, to_output=custom_to_output
)
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
# Load - should use default converters # Load - should use default converters
loaded = RobotProcessor.from_pretrained(tmp_dir) loaded = DataProcessorPipeline.from_pretrained(tmp_dir)
# Verify it uses default converters by checking with standard batch format # Verify it uses default converters by checking with standard batch format
batch = { batch = {
@@ -1792,7 +1798,7 @@ class NonCallableStep:
def test_construction_rejects_step_without_call(): def test_construction_rejects_step_without_call():
with pytest.raises(TypeError, match=r"must define __call__"): with pytest.raises(TypeError, match=r"must define __call__"):
RobotProcessor([NonCallableStep()]) DataProcessorPipeline([NonCallableStep()])
@dataclass @dataclass
@@ -1851,7 +1857,7 @@ class FeatureContractRemoveStep:
def test_features_orders_and_merges(policy_feature_factory): def test_features_orders_and_merges(policy_feature_factory):
p = RobotProcessor( p = DataProcessorPipeline(
[ [
FeatureContractAddStep("a", policy_feature_factory(FeatureType.STATE, (1,))), FeatureContractAddStep("a", policy_feature_factory(FeatureType.STATE, (1,))),
FeatureContractMutateStep("a", lambda v: PolicyFeature(type=v.type, shape=(3,))), FeatureContractMutateStep("a", lambda v: PolicyFeature(type=v.type, shape=(3,))),
@@ -1870,7 +1876,7 @@ def test_features_respects_initial_without_mutation(policy_feature_factory):
"seed": policy_feature_factory(FeatureType.STATE, (7,)), "seed": policy_feature_factory(FeatureType.STATE, (7,)),
"nested": policy_feature_factory(FeatureType.ENV, (0,)), "nested": policy_feature_factory(FeatureType.ENV, (0,)),
} }
p = RobotProcessor( p = DataProcessorPipeline(
[ [
FeatureContractMutateStep("seed", lambda v: PolicyFeature(type=v.type, shape=(v.shape[0] + 1,))), FeatureContractMutateStep("seed", lambda v: PolicyFeature(type=v.type, shape=(v.shape[0] + 1,))),
FeatureContractMutateStep( FeatureContractMutateStep(
@@ -1903,12 +1909,12 @@ def test_features_execution_order_tracking():
features["order"] = PolicyFeature(type=pf.type, shape=pf.shape + (code,)) features["order"] = PolicyFeature(type=pf.type, shape=pf.shape + (code,))
return features return features
out = RobotProcessor([Track("A"), Track("B"), Track("C")]).transform_features({}) out = DataProcessorPipeline([Track("A"), Track("B"), Track("C")]).transform_features({})
assert out["order"].shape == (1, 2, 3) assert out["order"].shape == (1, 2, 3)
def test_features_remove_key(policy_feature_factory): def test_features_remove_key(policy_feature_factory):
p = RobotProcessor( p = DataProcessorPipeline(
[ [
FeatureContractAddStep("a", policy_feature_factory(FeatureType.STATE, (1,))), FeatureContractAddStep("a", policy_feature_factory(FeatureType.STATE, (1,))),
FeatureContractRemoveStep("a"), FeatureContractRemoveStep("a"),
@@ -1923,7 +1929,7 @@ def test_features_remove_from_initial(policy_feature_factory):
"keep": policy_feature_factory(FeatureType.STATE, (1,)), "keep": policy_feature_factory(FeatureType.STATE, (1,)),
"drop": policy_feature_factory(FeatureType.STATE, (1,)), "drop": policy_feature_factory(FeatureType.STATE, (1,)),
} }
p = RobotProcessor([FeatureContractRemoveStep("drop")]) p = DataProcessorPipeline([FeatureContractRemoveStep("drop")])
out = p.transform_features(initial_features=initial) out = p.transform_features(initial_features=initial)
assert "drop" not in out and out["keep"] == initial["keep"] assert "drop" not in out and out["keep"] == initial["keep"]
@@ -1965,7 +1971,7 @@ class AddObservationStateFeatures:
def test_aggregate_joint_action_only(): def test_aggregate_joint_action_only():
rp = RobotProcessor([AddActionEEAndJointFeatures()]) rp = DataProcessorPipeline([AddActionEEAndJointFeatures()])
initial = {"front": (480, 640, 3)} initial = {"front": (480, 640, 3)}
out = aggregate_pipeline_dataset_features( out = aggregate_pipeline_dataset_features(
@@ -1983,7 +1989,7 @@ def test_aggregate_joint_action_only():
def test_aggregate_ee_action_and_observation_with_videos(): def test_aggregate_ee_action_and_observation_with_videos():
rp = RobotProcessor([AddActionEEAndJointFeatures(), AddObservationStateFeatures()]) rp = DataProcessorPipeline([AddActionEEAndJointFeatures(), AddObservationStateFeatures()])
initial = {"front": (480, 640, 3), "side": (720, 1280, 3)} initial = {"front": (480, 640, 3), "side": (720, 1280, 3)}
out = aggregate_pipeline_dataset_features( out = aggregate_pipeline_dataset_features(
@@ -2013,7 +2019,7 @@ def test_aggregate_ee_action_and_observation_with_videos():
def test_aggregate_both_action_types(): def test_aggregate_both_action_types():
rp = RobotProcessor([AddActionEEAndJointFeatures()]) rp = DataProcessorPipeline([AddActionEEAndJointFeatures()])
out = aggregate_pipeline_dataset_features( out = aggregate_pipeline_dataset_features(
pipeline=rp, pipeline=rp,
initial_features={}, initial_features={},
@@ -2028,7 +2034,7 @@ def test_aggregate_both_action_types():
def test_aggregate_images_when_use_videos_false(): def test_aggregate_images_when_use_videos_false():
rp = RobotProcessor([AddObservationStateFeatures(add_front_image=True)]) rp = DataProcessorPipeline([AddObservationStateFeatures(add_front_image=True)])
initial = {"back": (480, 640, 3)} initial = {"back": (480, 640, 3)}
out = aggregate_pipeline_dataset_features( out = aggregate_pipeline_dataset_features(
@@ -2045,7 +2051,7 @@ def test_aggregate_images_when_use_videos_false():
def test_aggregate_images_when_use_videos_true(): def test_aggregate_images_when_use_videos_true():
rp = RobotProcessor([AddObservationStateFeatures(add_front_image=True)]) rp = DataProcessorPipeline([AddObservationStateFeatures(add_front_image=True)])
initial = {"back": (480, 640, 3)} initial = {"back": (480, 640, 3)}
out = aggregate_pipeline_dataset_features( out = aggregate_pipeline_dataset_features(
@@ -2067,7 +2073,9 @@ def test_aggregate_images_when_use_videos_true():
def test_initial_camera_not_overridden_by_step_image(): def test_initial_camera_not_overridden_by_step_image():
# Step explicitly sets a different front image shape; initial has another shape. # Step explicitly sets a different front image shape; initial has another shape.
# aggregate_pipeline_dataset_features should keep the step's value (setdefault behavior on initial cams). # aggregate_pipeline_dataset_features should keep the step's value (setdefault behavior on initial cams).
rp = RobotProcessor([AddObservationStateFeatures(add_front_image=True, front_image_shape=(240, 320, 3))]) rp = DataProcessorPipeline(
[AddObservationStateFeatures(add_front_image=True, front_image_shape=(240, 320, 3))]
)
initial = {"front": (480, 640, 3)} # should NOT override the step-provided (240, 320, 3) initial = {"front": (480, 640, 3)} # should NOT override the step-provided (240, 320, 3)
out = aggregate_pipeline_dataset_features( out = aggregate_pipeline_dataset_features(
+11 -9
View File
@@ -21,9 +21,9 @@ import torch
from lerobot.configs.types import FeatureType from lerobot.configs.types import FeatureType
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
ProcessorStepRegistry, ProcessorStepRegistry,
RenameProcessor, RenameProcessor,
RobotProcessor,
TransitionKey, TransitionKey,
) )
from lerobot.processor.rename_processor import rename_stats from lerobot.processor.rename_processor import rename_stats
@@ -193,7 +193,7 @@ def test_integration_with_robot_processor():
} }
rename_processor = RenameProcessor(rename_map=rename_map) rename_processor = RenameProcessor(rename_map=rename_map)
pipeline = RobotProcessor([rename_processor], to_transition=lambda x: x, to_output=lambda x: x) pipeline = DataProcessorPipeline([rename_processor], to_transition=lambda x: x, to_output=lambda x: x)
observation = { observation = {
"agent_pos": np.array([1.0, 2.0, 3.0]), "agent_pos": np.array([1.0, 2.0, 3.0]),
@@ -226,7 +226,7 @@ def test_save_and_load_pretrained():
"old_image": "observation.image", "old_image": "observation.image",
} }
processor = RenameProcessor(rename_map=rename_map) processor = RenameProcessor(rename_map=rename_map)
pipeline = RobotProcessor([processor], name="TestRenameProcessor") pipeline = DataProcessorPipeline([processor], name="TestRenameProcessor")
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
# Save pipeline # Save pipeline
@@ -241,7 +241,7 @@ def test_save_and_load_pretrained():
assert len(state_files) == 0 assert len(state_files) == 0
# Load pipeline # Load pipeline
loaded_pipeline = RobotProcessor.from_pretrained( loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir, to_transition=lambda x: x, to_output=lambda x: x tmp_dir, to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -284,7 +284,7 @@ def test_registry_functionality():
def test_registry_based_save_load(): def test_registry_based_save_load():
"""Test save/load using registry name instead of module path.""" """Test save/load using registry name instead of module path."""
processor = RenameProcessor(rename_map={"key1": "renamed_key1"}) processor = RenameProcessor(rename_map={"key1": "renamed_key1"})
pipeline = RobotProcessor([processor], to_transition=lambda x: x, to_output=lambda x: x) pipeline = DataProcessorPipeline([processor], to_transition=lambda x: x, to_output=lambda x: x)
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
# Save and load # Save and load
@@ -293,7 +293,7 @@ def test_registry_based_save_load():
# Verify config uses registry name # Verify config uses registry name
import json import json
with open(Path(tmp_dir) / "robotprocessor.json") as f: # Default name is "RobotProcessor" with open(Path(tmp_dir) / "dataprocessorpipeline.json") as f: # Default name is "RobotProcessor"
config = json.load(f) config = json.load(f)
assert "registry_name" in config["steps"][0] assert "registry_name" in config["steps"][0]
@@ -301,7 +301,7 @@ def test_registry_based_save_load():
assert "class" not in config["steps"][0] # Should use registry, not module path assert "class" not in config["steps"][0] # Should use registry, not module path
# Load should work # Load should work
loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir) loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir)
loaded_processor = loaded_pipeline.steps[0] loaded_processor = loaded_pipeline.steps[0]
assert isinstance(loaded_processor, RenameProcessor) assert isinstance(loaded_processor, RenameProcessor)
assert loaded_processor.rename_map == {"key1": "renamed_key1"} assert loaded_processor.rename_map == {"key1": "renamed_key1"}
@@ -325,7 +325,9 @@ def test_chained_rename_processors():
} }
) )
pipeline = RobotProcessor([processor1, processor2], to_transition=lambda x: x, to_output=lambda x: x) pipeline = DataProcessorPipeline(
[processor1, processor2], to_transition=lambda x: x, to_output=lambda x: x
)
observation = { observation = {
"pos": np.array([1.0, 2.0]), "pos": np.array([1.0, 2.0]),
@@ -459,7 +461,7 @@ def test_features_chained_processors(policy_feature_factory):
processor2 = RenameProcessor( processor2 = RenameProcessor(
rename_map={"agent_position": "observation.state", "camera_image": "observation.image"} rename_map={"agent_position": "observation.state", "camera_image": "observation.image"}
) )
pipeline = RobotProcessor([processor1, processor2]) pipeline = DataProcessorPipeline([processor1, processor2])
spec = { spec = {
"pos": policy_feature_factory(FeatureType.STATE, (7,)), "pos": policy_feature_factory(FeatureType.STATE, (7,)),
+4 -4
View File
@@ -25,10 +25,10 @@ from lerobot.constants import ACTION, OBS_STATE
from lerobot.policies.sac.configuration_sac import SACConfig from lerobot.policies.sac.configuration_sac import SACConfig
from lerobot.policies.sac.processor_sac import make_sac_pre_post_processors from lerobot.policies.sac.processor_sac import make_sac_pre_post_processors
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
NormalizerProcessor, NormalizerProcessor,
RenameProcessor, RenameProcessor,
RobotProcessor,
ToBatchProcessor, ToBatchProcessor,
TransitionKey, TransitionKey,
UnnormalizerProcessor, UnnormalizerProcessor,
@@ -234,13 +234,13 @@ def test_sac_processor_without_stats():
factory_preprocessor, factory_postprocessor = make_sac_pre_post_processors(config, dataset_stats=None) factory_preprocessor, factory_postprocessor = make_sac_pre_post_processors(config, dataset_stats=None)
# Create new processors with EnvTransition input/output # Create new processors with EnvTransition input/output
preprocessor = RobotProcessor( preprocessor = DataProcessorPipeline(
factory_preprocessor.steps, factory_preprocessor.steps,
name=factory_preprocessor.name, name=factory_preprocessor.name,
to_transition=lambda x: x, to_transition=lambda x: x,
to_output=lambda x: x, to_output=lambda x: x,
) )
postprocessor = RobotProcessor( postprocessor = DataProcessorPipeline(
factory_postprocessor.steps, factory_postprocessor.steps,
name=factory_postprocessor.name, name=factory_postprocessor.name,
to_transition=lambda x: x, to_transition=lambda x: x,
@@ -277,7 +277,7 @@ def test_sac_processor_save_and_load():
preprocessor.save_pretrained(tmpdir) preprocessor.save_pretrained(tmpdir)
# Load preprocessor # Load preprocessor
loaded_preprocessor = RobotProcessor.from_pretrained( loaded_preprocessor = DataProcessorPipeline.from_pretrained(
tmpdir, to_transition=lambda x: x, to_output=lambda x: x tmpdir, to_transition=lambda x: x, to_output=lambda x: x
) )
+4 -4
View File
@@ -25,10 +25,10 @@ from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE
from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
from lerobot.policies.tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors from lerobot.policies.tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
NormalizerProcessor, NormalizerProcessor,
RenameProcessor, RenameProcessor,
RobotProcessor,
ToBatchProcessor, ToBatchProcessor,
TransitionKey, TransitionKey,
UnnormalizerProcessor, UnnormalizerProcessor,
@@ -251,13 +251,13 @@ def test_tdmpc_processor_without_stats():
factory_preprocessor, factory_postprocessor = make_tdmpc_pre_post_processors(config, dataset_stats=None) factory_preprocessor, factory_postprocessor = make_tdmpc_pre_post_processors(config, dataset_stats=None)
# Create new processors with EnvTransition input/output # Create new processors with EnvTransition input/output
preprocessor = RobotProcessor( preprocessor = DataProcessorPipeline(
factory_preprocessor.steps, factory_preprocessor.steps,
name=factory_preprocessor.name, name=factory_preprocessor.name,
to_transition=lambda x: x, to_transition=lambda x: x,
to_output=lambda x: x, to_output=lambda x: x,
) )
postprocessor = RobotProcessor( postprocessor = DataProcessorPipeline(
factory_postprocessor.steps, factory_postprocessor.steps,
name=factory_postprocessor.name, name=factory_postprocessor.name,
to_transition=lambda x: x, to_transition=lambda x: x,
@@ -297,7 +297,7 @@ def test_tdmpc_processor_save_and_load():
preprocessor.save_pretrained(tmpdir) preprocessor.save_pretrained(tmpdir)
# Load preprocessor # Load preprocessor
loaded_preprocessor = RobotProcessor.from_pretrained( loaded_preprocessor = DataProcessorPipeline.from_pretrained(
tmpdir, to_transition=lambda x: x, to_output=lambda x: x tmpdir, to_transition=lambda x: x, to_output=lambda x: x
) )
+13 -7
View File
@@ -10,7 +10,7 @@ import torch
from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.configs.types import FeatureType, PolicyFeature
from lerobot.constants import OBS_LANGUAGE from lerobot.constants import OBS_LANGUAGE
from lerobot.processor import RobotProcessor, TokenizerProcessor, TransitionKey from lerobot.processor import DataProcessorPipeline, TokenizerProcessor, TransitionKey
from tests.utils import require_package from tests.utils import require_package
@@ -388,7 +388,9 @@ def test_integration_with_robot_processor(mock_auto_tokenizer):
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
tokenizer_processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=6) tokenizer_processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=6)
robot_processor = RobotProcessor([tokenizer_processor], to_transition=lambda x: x, to_output=lambda x: x) robot_processor = DataProcessorPipeline(
[tokenizer_processor], to_transition=lambda x: x, to_output=lambda x: x
)
transition = create_transition( transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])}, observation={"state": torch.tensor([1.0, 2.0])},
@@ -426,14 +428,16 @@ def test_save_and_load_pretrained_with_tokenizer_name(mock_auto_tokenizer):
tokenizer_name="test-tokenizer", max_length=32, task_key="instruction" tokenizer_name="test-tokenizer", max_length=32, task_key="instruction"
) )
robot_processor = RobotProcessor([original_processor], to_transition=lambda x: x, to_output=lambda x: x) robot_processor = DataProcessorPipeline(
[original_processor], to_transition=lambda x: x, to_output=lambda x: x
)
with tempfile.TemporaryDirectory() as temp_dir: with tempfile.TemporaryDirectory() as temp_dir:
# Save processor # Save processor
robot_processor.save_pretrained(temp_dir) robot_processor.save_pretrained(temp_dir)
# Load processor - tokenizer will be recreated from saved config # Load processor - tokenizer will be recreated from saved config
loaded_processor = RobotProcessor.from_pretrained( loaded_processor = DataProcessorPipeline.from_pretrained(
temp_dir, to_transition=lambda x: x, to_output=lambda x: x temp_dir, to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -457,14 +461,16 @@ def test_save_and_load_pretrained_with_tokenizer_object():
original_processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=32, task_key="instruction") original_processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=32, task_key="instruction")
robot_processor = RobotProcessor([original_processor], to_transition=lambda x: x, to_output=lambda x: x) robot_processor = DataProcessorPipeline(
[original_processor], to_transition=lambda x: x, to_output=lambda x: x
)
with tempfile.TemporaryDirectory() as temp_dir: with tempfile.TemporaryDirectory() as temp_dir:
# Save processor # Save processor
robot_processor.save_pretrained(temp_dir) robot_processor.save_pretrained(temp_dir)
# Load processor with tokenizer override (since tokenizer object wasn't saved) # Load processor with tokenizer override (since tokenizer object wasn't saved)
loaded_processor = RobotProcessor.from_pretrained( loaded_processor = DataProcessorPipeline.from_pretrained(
temp_dir, temp_dir,
overrides={"tokenizer_processor": {"tokenizer": mock_tokenizer}}, overrides={"tokenizer_processor": {"tokenizer": mock_tokenizer}},
to_transition=lambda x: x, to_transition=lambda x: x,
@@ -956,7 +962,7 @@ def test_integration_with_device_processor(mock_auto_tokenizer):
# Create pipeline with TokenizerProcessor then DeviceProcessor # Create pipeline with TokenizerProcessor then DeviceProcessor
tokenizer_processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=6) tokenizer_processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=6)
device_processor = DeviceProcessor(device="cuda:0") device_processor = DeviceProcessor(device="cuda:0")
robot_processor = RobotProcessor( robot_processor = DataProcessorPipeline(
[tokenizer_processor, device_processor], to_transition=lambda x: x, to_output=lambda x: x [tokenizer_processor, device_processor], to_transition=lambda x: x, to_output=lambda x: x
) )
+4 -4
View File
@@ -25,10 +25,10 @@ from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE
from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
from lerobot.policies.vqbet.processor_vqbet import make_vqbet_pre_post_processors from lerobot.policies.vqbet.processor_vqbet import make_vqbet_pre_post_processors
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline,
DeviceProcessor, DeviceProcessor,
NormalizerProcessor, NormalizerProcessor,
RenameProcessor, RenameProcessor,
RobotProcessor,
ToBatchProcessor, ToBatchProcessor,
TransitionKey, TransitionKey,
UnnormalizerProcessor, UnnormalizerProcessor,
@@ -244,13 +244,13 @@ def test_vqbet_processor_without_stats():
factory_preprocessor, factory_postprocessor = make_vqbet_pre_post_processors(config, dataset_stats=None) factory_preprocessor, factory_postprocessor = make_vqbet_pre_post_processors(config, dataset_stats=None)
# Create new processors with EnvTransition input/output # Create new processors with EnvTransition input/output
preprocessor = RobotProcessor( preprocessor = DataProcessorPipeline(
factory_preprocessor.steps, factory_preprocessor.steps,
name=factory_preprocessor.name, name=factory_preprocessor.name,
to_transition=lambda x: x, to_transition=lambda x: x,
to_output=lambda x: x, to_output=lambda x: x,
) )
postprocessor = RobotProcessor( postprocessor = DataProcessorPipeline(
factory_postprocessor.steps, factory_postprocessor.steps,
name=factory_postprocessor.name, name=factory_postprocessor.name,
to_transition=lambda x: x, to_transition=lambda x: x,
@@ -290,7 +290,7 @@ def test_vqbet_processor_save_and_load():
preprocessor.save_pretrained(tmpdir) preprocessor.save_pretrained(tmpdir)
# Load preprocessor # Load preprocessor
loaded_preprocessor = RobotProcessor.from_pretrained( loaded_preprocessor = DataProcessorPipeline.from_pretrained(
tmpdir, to_transition=lambda x: x, to_output=lambda x: x tmpdir, to_transition=lambda x: x, to_output=lambda x: x
) )