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

chore(processor): add Step suffix to all processors (#1854)

Browse Source 
* refactor(processor): rename MapDeltaActionToRobotAction and MapTensorToDeltaActionDict for consistency

* refactor(processor): rename DeviceProcessor to DeviceProcessorStep for consistency across modules

* refactor(processor): rename Torch2NumpyActionProcessor to Torch2NumpyActionProcessorStep for consistency

* refactor(processor): rename Numpy2TorchActionProcessor to Numpy2TorchActionProcessorStep for consistency

* refactor(processor): rename AddTeleopActionAsComplimentaryData to AddTeleopActionAsComplimentaryDataStep for consistency

* refactor(processor): rename ImageCropResizeProcessor and AddTeleopEventsAsInfo for consistency

* refactor(processor): rename TimeLimitProcessor to TimeLimitProcessorStep for consistency

* refactor(processor): rename GripperPenaltyProcessor to GripperPenaltyProcessorStep for consistency

* refactor(processor): rename InterventionActionProcessor to InterventionActionProcessorStep for consistency

* refactor(processor): rename RewardClassifierProcessor to RewardClassifierProcessorStep for consistency

* refactor(processor): rename JointVelocityProcessor to JointVelocityProcessorStep for consistency

* refactor(processor): rename MotorCurrentProcessor to MotorCurrentProcessorStep for consistency

* refactor(processor): rename NormalizerProcessor and UnnormalizerProcessor to NormalizerProcessorStep and UnnormalizerProcessorStep for consistency

* refactor(processor): rename VanillaObservationProcessor to VanillaObservationProcessorStep for consistency

* refactor(processor): rename RenameProcessor to RenameProcessorStep for consistency

* refactor(processor): rename TokenizerProcessor to TokenizerProcessorStep for consistency

* refactor(processor): rename ToBatchProcessor to AddBatchDimensionProcessorStep for consistency

* refactor(processor): update config file name in test for RenameProcessorStep consistency
This commit is contained in:
Adil Zouitine
2025-09-03 18:12:11 +02:00
committed by GitHub
parent 029c4a9a76
commit ce793cde64
37 changed files with 653 additions and 628 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/hilserl.mdx
+13 -13
View File
@@ -143,27 +143,27 @@ HIL-SERL uses a modular processor pipeline architecture that processes robot obs
The environment processor (`env_processor`) handles incoming observations and environment state: The environment processor (`env_processor`) handles incoming observations and environment state:
1. **VanillaObservationProcessor**: Converts raw robot observations into standardized format 1. **VanillaObservationProcessorStep**: Converts raw robot observations into standardized format
2. **JointVelocityProcessor** (optional): Adds joint velocity information to observations 2. **JointVelocityProcessorStep** (optional): Adds joint velocity information to observations
3. **MotorCurrentProcessor** (optional): Adds motor current readings to observations 3. **MotorCurrentProcessorStep** (optional): Adds motor current readings to observations
4. **ForwardKinematicsJointsToEE** (optional): Computes end-effector pose from joint positions 4. **ForwardKinematicsJointsToEE** (optional): Computes end-effector pose from joint positions
5. **ImageCropResizeProcessor** (optional): Crops and resizes camera images 5. **ImageCropResizeProcessorStep** (optional): Crops and resizes camera images
6. **TimeLimitProcessor** (optional): Enforces episode time limits 6. **TimeLimitProcessorStep** (optional): Enforces episode time limits
7. **GripperPenaltyProcessor** (optional): Applies penalties for inappropriate gripper usage 7. **GripperPenaltyProcessorStep** (optional): Applies penalties for inappropriate gripper usage
8. **RewardClassifierProcessor** (optional): Automated reward detection using vision models 8. **RewardClassifierProcessorStep** (optional): Automated reward detection using vision models
9. **ToBatchProcessor**: Converts data to batch format for neural network processing 9. **AddBatchDimensionProcessorStep**: Converts data to batch format for neural network processing
10. **DeviceProcessor**: Moves data to the specified compute device (CPU/GPU) 10. **DeviceProcessorStep**: Moves data to the specified compute device (CPU/GPU)
#### Action Processor Pipeline #### Action Processor Pipeline
The action processor (`action_processor`) handles outgoing actions and human interventions: The action processor (`action_processor`) handles outgoing actions and human interventions:
1. **AddTeleopActionAsComplimentaryData**: Captures teleoperator actions for logging 1. **AddTeleopActionAsComplimentaryDataStep**: Captures teleoperator actions for logging
2. **AddTeleopEventsAsInfo**: Records intervention events and episode control signals 2. **AddTeleopEventsAsInfoStep**: Records intervention events and episode control signals
3. **AddRobotObservationAsComplimentaryData**: Stores raw robot state for processing 3. **AddRobotObservationAsComplimentaryData**: Stores raw robot state for processing
4. **InterventionActionProcessor**: Handles human interventions and episode termination 4. **InterventionActionProcessorStep**: Handles human interventions and episode termination
5. **Inverse Kinematics Pipeline** (when enabled): 5. **Inverse Kinematics Pipeline** (when enabled):
- **MapDeltaActionToRobotAction**: Converts delta actions to robot action format - **MapDeltaActionToRobotActionStep**: Converts delta actions to robot action format
- **EEReferenceAndDelta**: Computes end-effector reference and delta movements - **EEReferenceAndDelta**: Computes end-effector reference and delta movements
- **EEBoundsAndSafety**: Enforces workspace safety bounds - **EEBoundsAndSafety**: Enforces workspace safety bounds
- **InverseKinematicsEEToJoints**: Converts end-effector actions to joint targets - **InverseKinematicsEEToJoints**: Converts end-effector actions to joint targets
src/lerobot/policies/act/processor_act.py
+11 -11
View File
@@ -18,13 +18,13 @@ 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 (
AddBatchDimensionProcessorStep,
DataProcessorPipeline, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
ProcessorKwargs, ProcessorKwargs,
RenameProcessor, RenameProcessorStep,
ToBatchProcessor, UnnormalizerProcessorStep,
UnnormalizerProcessor,
) )
@@ -40,18 +40,18 @@ def make_act_pre_post_processors(
postprocessor_kwargs = {} postprocessor_kwargs = {}
input_steps = [ input_steps = [
RenameProcessor(rename_map={}), RenameProcessorStep(rename_map={}),
NormalizerProcessor( NormalizerProcessorStep(
features={**config.input_features, **config.output_features}, features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping, norm_map=config.normalization_mapping,
stats=dataset_stats, stats=dataset_stats,
), ),
ToBatchProcessor(), AddBatchDimensionProcessorStep(),
DeviceProcessor(device=config.device), DeviceProcessorStep(device=config.device),
] ]
output_steps = [ output_steps = [
DeviceProcessor(device="cpu"), DeviceProcessorStep(device="cpu"),
UnnormalizerProcessor( UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
), ),
] ]
src/lerobot/policies/diffusion/processor_diffusion.py
+11 -11
View File
@@ -19,13 +19,13 @@ 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 (
AddBatchDimensionProcessorStep,
DataProcessorPipeline, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
ProcessorKwargs, ProcessorKwargs,
RenameProcessor, RenameProcessorStep,
ToBatchProcessor, UnnormalizerProcessorStep,
UnnormalizerProcessor,
) )
@@ -41,18 +41,18 @@ def make_diffusion_pre_post_processors(
postprocessor_kwargs = {} postprocessor_kwargs = {}
input_steps = [ input_steps = [
RenameProcessor(rename_map={}), RenameProcessorStep(rename_map={}),
NormalizerProcessor( NormalizerProcessorStep(
features={**config.input_features, **config.output_features}, features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping, norm_map=config.normalization_mapping,
stats=dataset_stats, stats=dataset_stats,
), ),
ToBatchProcessor(), AddBatchDimensionProcessorStep(),
DeviceProcessor(device=config.device), DeviceProcessorStep(device=config.device),
] ]
output_steps = [ output_steps = [
DeviceProcessor(device="cpu"), DeviceProcessorStep(device="cpu"),
UnnormalizerProcessor( UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
), ),
] ]
src/lerobot/policies/pi0/processor_pi0.py
+13 -13
View File
@@ -21,17 +21,17 @@ from lerobot.configs.types import PolicyFeature
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 (
AddBatchDimensionProcessorStep,
ComplementaryDataProcessorStep, ComplementaryDataProcessorStep,
DataProcessorPipeline, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
ProcessorKwargs, ProcessorKwargs,
ProcessorStep, ProcessorStep,
ProcessorStepRegistry, ProcessorStepRegistry,
RenameProcessor, RenameProcessorStep,
ToBatchProcessor, TokenizerProcessorStep,
TokenizerProcessor, UnnormalizerProcessorStep,
UnnormalizerProcessor,
) )
@@ -80,26 +80,26 @@ def make_pi0_pre_post_processors(
# Add remaining processors # Add remaining processors
input_steps: list[ProcessorStep] = [ input_steps: list[ProcessorStep] = [
RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one RenameProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
NormalizerProcessor( NormalizerProcessorStep(
features={**config.input_features, **config.output_features}, features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping, norm_map=config.normalization_mapping,
stats=dataset_stats, stats=dataset_stats,
), ),
ToBatchProcessor(), AddBatchDimensionProcessorStep(),
Pi0NewLineProcessor(), # Add newlines before tokenization for PaliGemma Pi0NewLineProcessor(), # Add newlines before tokenization for PaliGemma
TokenizerProcessor( TokenizerProcessorStep(
tokenizer_name="google/paligemma-3b-pt-224", tokenizer_name="google/paligemma-3b-pt-224",
max_length=config.tokenizer_max_length, max_length=config.tokenizer_max_length,
padding_side="right", padding_side="right",
padding="max_length", padding="max_length",
), ),
DeviceProcessor(device=config.device), DeviceProcessorStep(device=config.device),
] ]
output_steps: list[ProcessorStep] = [ output_steps: list[ProcessorStep] = [
DeviceProcessor(device="cpu"), DeviceProcessorStep(device="cpu"),
UnnormalizerProcessor( UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
), ),
] ]
src/lerobot/policies/pi0fast/processor_pi0fast.py
+11 -11
View File
@@ -19,13 +19,13 @@ 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 (
AddBatchDimensionProcessorStep,
DataProcessorPipeline, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
ProcessorKwargs, ProcessorKwargs,
RenameProcessor, RenameProcessorStep,
ToBatchProcessor, UnnormalizerProcessorStep,
UnnormalizerProcessor,
) )
@@ -41,18 +41,18 @@ def make_pi0fast_pre_post_processors(
postprocessor_kwargs = {} postprocessor_kwargs = {}
input_steps = [ input_steps = [
RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one RenameProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
NormalizerProcessor( NormalizerProcessorStep(
features={**config.input_features, **config.output_features}, features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping, norm_map=config.normalization_mapping,
stats=dataset_stats, stats=dataset_stats,
), ),
ToBatchProcessor(), AddBatchDimensionProcessorStep(),
DeviceProcessor(device=config.device), DeviceProcessorStep(device=config.device),
] ]
output_steps = [ output_steps = [
DeviceProcessor(device="cpu"), DeviceProcessorStep(device="cpu"),
UnnormalizerProcessor( UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
), ),
] ]
src/lerobot/policies/sac/processor_sac.py
+11 -11
View File
@@ -20,13 +20,13 @@ 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 (
AddBatchDimensionProcessorStep,
DataProcessorPipeline, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
ProcessorKwargs, ProcessorKwargs,
RenameProcessor, RenameProcessorStep,
ToBatchProcessor, UnnormalizerProcessorStep,
UnnormalizerProcessor,
) )
@@ -42,18 +42,18 @@ def make_sac_pre_post_processors(
postprocessor_kwargs = {} postprocessor_kwargs = {}
input_steps = [ input_steps = [
RenameProcessor(rename_map={}), RenameProcessorStep(rename_map={}),
NormalizerProcessor( NormalizerProcessorStep(
features={**config.input_features, **config.output_features}, features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping, norm_map=config.normalization_mapping,
stats=dataset_stats, stats=dataset_stats,
), ),
ToBatchProcessor(), AddBatchDimensionProcessorStep(),
DeviceProcessor(device=config.device), DeviceProcessorStep(device=config.device),
] ]
output_steps = [ output_steps = [
DeviceProcessor(device="cpu"), DeviceProcessorStep(device="cpu"),
UnnormalizerProcessor( UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
), ),
] ]
src/lerobot/policies/sac/reward_model/processor_classifier.py
+6 -6
View File
@@ -18,9 +18,9 @@ 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, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
IdentityProcessorStep, IdentityProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
ProcessorKwargs, ProcessorKwargs,
) )
@@ -37,15 +37,15 @@ def make_classifier_processor(
postprocessor_kwargs = {} postprocessor_kwargs = {}
input_steps = [ input_steps = [
NormalizerProcessor( NormalizerProcessorStep(
features=config.input_features, norm_map=config.normalization_mapping, stats=dataset_stats features=config.input_features, norm_map=config.normalization_mapping, stats=dataset_stats
), ),
NormalizerProcessor( NormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
), ),
DeviceProcessor(device=config.device), DeviceProcessorStep(device=config.device),
] ]
output_steps = [DeviceProcessor(device="cpu"), IdentityProcessorStep()] output_steps = [DeviceProcessorStep(device="cpu"), IdentityProcessorStep()]
return ( return (
DataProcessorPipeline( DataProcessorPipeline(
src/lerobot/policies/smolvla/processor_smolvla.py
+13 -13
View File
@@ -20,16 +20,16 @@ from lerobot.configs.types import PolicyFeature
from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME
from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
from lerobot.processor import ( from lerobot.processor import (
AddBatchDimensionProcessorStep,
ComplementaryDataProcessorStep, ComplementaryDataProcessorStep,
DataProcessorPipeline, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
ProcessorKwargs, ProcessorKwargs,
ProcessorStepRegistry, ProcessorStepRegistry,
RenameProcessor, RenameProcessorStep,
ToBatchProcessor, TokenizerProcessorStep,
TokenizerProcessor, UnnormalizerProcessorStep,
UnnormalizerProcessor,
) )
@@ -45,25 +45,25 @@ def make_smolvla_pre_post_processors(
postprocessor_kwargs = {} postprocessor_kwargs = {}
input_steps = [ input_steps = [
RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one RenameProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
NormalizerProcessor( NormalizerProcessorStep(
features={**config.input_features, **config.output_features}, features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping, norm_map=config.normalization_mapping,
stats=dataset_stats, stats=dataset_stats,
), ),
ToBatchProcessor(), AddBatchDimensionProcessorStep(),
SmolVLANewLineProcessor(), SmolVLANewLineProcessor(),
TokenizerProcessor( TokenizerProcessorStep(
tokenizer_name=config.vlm_model_name, tokenizer_name=config.vlm_model_name,
padding=config.pad_language_to, padding=config.pad_language_to,
padding_side="right", padding_side="right",
max_length=config.tokenizer_max_length, max_length=config.tokenizer_max_length,
), ),
DeviceProcessor(device=config.device), DeviceProcessorStep(device=config.device),
] ]
output_steps = [ output_steps = [
DeviceProcessor(device="cpu"), DeviceProcessorStep(device="cpu"),
UnnormalizerProcessor( UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
), ),
] ]
src/lerobot/policies/tdmpc/processor_tdmpc.py
+11 -11
View File
@@ -19,13 +19,13 @@ 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 (
AddBatchDimensionProcessorStep,
DataProcessorPipeline, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
ProcessorKwargs, ProcessorKwargs,
RenameProcessor, RenameProcessorStep,
ToBatchProcessor, UnnormalizerProcessorStep,
UnnormalizerProcessor,
) )
@@ -41,18 +41,18 @@ def make_tdmpc_pre_post_processors(
postprocessor_kwargs = {} postprocessor_kwargs = {}
input_steps = [ input_steps = [
RenameProcessor(rename_map={}), RenameProcessorStep(rename_map={}),
NormalizerProcessor( NormalizerProcessorStep(
features={**config.input_features, **config.output_features}, features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping, norm_map=config.normalization_mapping,
stats=dataset_stats, stats=dataset_stats,
), ),
ToBatchProcessor(), AddBatchDimensionProcessorStep(),
DeviceProcessor(device=config.device), DeviceProcessorStep(device=config.device),
] ]
output_steps = [ output_steps = [
DeviceProcessor(device="cpu"), DeviceProcessorStep(device="cpu"),
UnnormalizerProcessor( UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
), ),
] ]
src/lerobot/policies/vqbet/processor_vqbet.py
+11 -11
View File
@@ -20,13 +20,13 @@ 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 (
AddBatchDimensionProcessorStep,
DataProcessorPipeline, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
ProcessorKwargs, ProcessorKwargs,
RenameProcessor, RenameProcessorStep,
ToBatchProcessor, UnnormalizerProcessorStep,
UnnormalizerProcessor,
) )
@@ -42,18 +42,18 @@ def make_vqbet_pre_post_processors(
postprocessor_kwargs = {} postprocessor_kwargs = {}
input_steps = [ input_steps = [
RenameProcessor(rename_map={}), # Let the possibility to the user to rename the keys RenameProcessorStep(rename_map={}), # Let the possibility to the user to rename the keys
NormalizerProcessor( NormalizerProcessorStep(
features={**config.input_features, **config.output_features}, features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping, norm_map=config.normalization_mapping,
stats=dataset_stats, stats=dataset_stats,
), ),
ToBatchProcessor(), AddBatchDimensionProcessorStep(),
DeviceProcessor(device=config.device), DeviceProcessorStep(device=config.device),
] ]
output_steps = [ output_steps = [
DeviceProcessor(device="cpu"), DeviceProcessorStep(device="cpu"),
UnnormalizerProcessor( UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
), ),
] ]
src/lerobot/processor/__init__.py
+36 -36
View File
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
from .batch_processor import ToBatchProcessor from .batch_processor import AddBatchDimensionProcessorStep
from .converters import ( from .converters import (
batch_to_transition, batch_to_transition,
create_transition, create_transition,
@@ -23,21 +23,21 @@ from .converters import (
transition_to_dataset_frame, transition_to_dataset_frame,
) )
from .core import EnvTransition, TransitionKey from .core import EnvTransition, TransitionKey
from .delta_action_processor import MapDeltaActionToRobotAction, MapTensorToDeltaActionDict from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep
from .device_processor import DeviceProcessor from .device_processor import DeviceProcessorStep
from .gym_action_processor import Numpy2TorchActionProcessor, Torch2NumpyActionProcessor from .gym_action_processor import Numpy2TorchActionProcessorStep, Torch2NumpyActionProcessorStep
from .hil_processor import ( from .hil_processor import (
AddTeleopActionAsComplimentaryData, AddTeleopActionAsComplimentaryDataStep,
AddTeleopEventsAsInfo, AddTeleopEventsAsInfoStep,
GripperPenaltyProcessor, GripperPenaltyProcessorStep,
ImageCropResizeProcessor, ImageCropResizeProcessorStep,
InterventionActionProcessor, InterventionActionProcessorStep,
RewardClassifierProcessor, RewardClassifierProcessorStep,
TimeLimitProcessor, TimeLimitProcessorStep,
) )
from .joint_observations_processor import JointVelocityProcessor, MotorCurrentProcessor from .joint_observations_processor import JointVelocityProcessorStep, MotorCurrentProcessorStep
from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor, hotswap_stats from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep, hotswap_stats
from .observation_processor import VanillaObservationProcessor from .observation_processor import VanillaObservationProcessorStep
from .pipeline import ( from .pipeline import (
ActionProcessorStep, ActionProcessorStep,
ComplementaryDataProcessorStep, ComplementaryDataProcessorStep,
@@ -52,48 +52,48 @@ from .pipeline import (
RewardProcessorStep, RewardProcessorStep,
TruncatedProcessorStep, TruncatedProcessorStep,
) )
from .rename_processor import RenameProcessor from .rename_processor import RenameProcessorStep
from .tokenizer_processor import TokenizerProcessor from .tokenizer_processor import TokenizerProcessorStep
__all__ = [ __all__ = [
"ActionProcessorStep", "ActionProcessorStep",
"AddTeleopActionAsComplimentaryData", "AddTeleopActionAsComplimentaryDataStep",
"AddTeleopEventsAsInfo", "AddTeleopEventsAsInfoStep",
"ComplementaryDataProcessorStep", "ComplementaryDataProcessorStep",
"batch_to_transition", "batch_to_transition",
"create_transition", "create_transition",
"DeviceProcessor", "DeviceProcessorStep",
"DoneProcessorStep", "DoneProcessorStep",
"EnvTransition", "EnvTransition",
"GripperPenaltyProcessor", "GripperPenaltyProcessorStep",
"hotswap_stats", "hotswap_stats",
"IdentityProcessorStep", "IdentityProcessorStep",
"ImageCropResizeProcessor", "ImageCropResizeProcessorStep",
"InfoProcessorStep", "InfoProcessorStep",
"InterventionActionProcessor", "InterventionActionProcessorStep",
"JointVelocityProcessor", "JointVelocityProcessorStep",
"MapDeltaActionToRobotAction", "MapDeltaActionToRobotActionStep",
"MapTensorToDeltaActionDict", "MapTensorToDeltaActionDictStep",
"merge_transitions", "merge_transitions",
"MotorCurrentProcessor", "MotorCurrentProcessorStep",
"NormalizerProcessor", "NormalizerProcessorStep",
"Numpy2TorchActionProcessor", "Numpy2TorchActionProcessorStep",
"ObservationProcessorStep", "ObservationProcessorStep",
"ProcessorKwargs", "ProcessorKwargs",
"ProcessorStep", "ProcessorStep",
"ProcessorStepRegistry", "ProcessorStepRegistry",
"RenameProcessor", "RenameProcessorStep",
"RewardClassifierProcessor", "RewardClassifierProcessorStep",
"RewardProcessorStep", "RewardProcessorStep",
"DataProcessorPipeline", "DataProcessorPipeline",
"TimeLimitProcessor", "TimeLimitProcessorStep",
"ToBatchProcessor", "AddBatchDimensionProcessorStep",
"TokenizerProcessor", "TokenizerProcessorStep",
"Torch2NumpyActionProcessor", "Torch2NumpyActionProcessorStep",
"transition_to_batch", "transition_to_batch",
"transition_to_dataset_frame", "transition_to_dataset_frame",
"TransitionKey", "TransitionKey",
"TruncatedProcessorStep", "TruncatedProcessorStep",
"UnnormalizerProcessor", "UnnormalizerProcessorStep",
"VanillaObservationProcessor", "VanillaObservationProcessorStep",
] ]
src/lerobot/processor/batch_processor.py
+11 -9
View File
@@ -30,7 +30,7 @@ from .pipeline import (
@dataclass @dataclass
@ProcessorStepRegistry.register(name="to_batch_processor_action") @ProcessorStepRegistry.register(name="to_batch_processor_action")
class ToBatchProcessorAction(ActionProcessorStep): class AddBatchDimensionActionStep(ActionProcessorStep):
"""Process action component in-place, adding batch dimension if needed.""" """Process action component in-place, adding batch dimension if needed."""
def action(self, action): def action(self, action):
@@ -45,7 +45,7 @@ class ToBatchProcessorAction(ActionProcessorStep):
@dataclass @dataclass
@ProcessorStepRegistry.register(name="to_batch_processor_observation") @ProcessorStepRegistry.register(name="to_batch_processor_observation")
class ToBatchProcessorObservation(ObservationProcessorStep): class AddBatchDimensionObservationStep(ObservationProcessorStep):
"""Process observation component in-place, adding batch dimensions where needed.""" """Process observation component in-place, adding batch dimensions where needed."""
def observation(self, observation): def observation(self, observation):
@@ -74,7 +74,7 @@ class ToBatchProcessorObservation(ObservationProcessorStep):
@dataclass @dataclass
@ProcessorStepRegistry.register(name="to_batch_processor_complementary_data") @ProcessorStepRegistry.register(name="to_batch_processor_complementary_data")
class ToBatchProcessorComplementaryData(ComplementaryDataProcessorStep): class AddBatchDimensionComplementaryDataStep(ComplementaryDataProcessorStep):
"""Process complementary data in-place, handling task field batching.""" """Process complementary data in-place, handling task field batching."""
def complementary_data(self, complementary_data): def complementary_data(self, complementary_data):
@@ -103,7 +103,7 @@ class ToBatchProcessorComplementaryData(ComplementaryDataProcessorStep):
@dataclass @dataclass
@ProcessorStepRegistry.register(name="to_batch_processor") @ProcessorStepRegistry.register(name="to_batch_processor")
class ToBatchProcessor(ProcessorStep): class AddBatchDimensionProcessorStep(ProcessorStep):
"""Processor that adds batch dimensions to observations and actions when needed. """Processor that adds batch dimensions to observations and actions when needed.
This processor ensures that observations and actions have proper batch dimensions for model processing: This processor ensures that observations and actions have proper batch dimensions for model processing:
@@ -138,12 +138,14 @@ class ToBatchProcessor(ProcessorStep):
``` ```
""" """
to_batch_action_processor: ToBatchProcessorAction = field(default_factory=ToBatchProcessorAction) to_batch_action_processor: AddBatchDimensionActionStep = field(
to_batch_observation_processor: ToBatchProcessorObservation = field( default_factory=AddBatchDimensionActionStep
default_factory=ToBatchProcessorObservation
) )
to_batch_complementary_data_processor: ToBatchProcessorComplementaryData = field( to_batch_observation_processor: AddBatchDimensionObservationStep = field(
default_factory=ToBatchProcessorComplementaryData default_factory=AddBatchDimensionObservationStep
)
to_batch_complementary_data_processor: AddBatchDimensionComplementaryDataStep = field(
default_factory=AddBatchDimensionComplementaryDataStep
) )
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
src/lerobot/processor/delta_action_processor.py
+2 -2
View File
@@ -26,7 +26,7 @@ from .pipeline import ActionProcessorStep, ProcessorStepRegistry
@ProcessorStepRegistry.register("map_tensor_to_delta_action_dict") @ProcessorStepRegistry.register("map_tensor_to_delta_action_dict")
@dataclass @dataclass
class MapTensorToDeltaActionDict(ActionProcessorStep): class MapTensorToDeltaActionDictStep(ActionProcessorStep):
""" """
Map a tensor to a delta action dictionary. Map a tensor to a delta action dictionary.
""" """
@@ -58,7 +58,7 @@ class MapTensorToDeltaActionDict(ActionProcessorStep):
@ProcessorStepRegistry.register("map_delta_action_to_robot_action") @ProcessorStepRegistry.register("map_delta_action_to_robot_action")
@dataclass @dataclass
class MapDeltaActionToRobotAction(ActionProcessorStep): class MapDeltaActionToRobotActionStep(ActionProcessorStep):
""" """
Map delta actions from teleoperators (gamepad, keyboard) to robot target actions Map delta actions from teleoperators (gamepad, keyboard) to robot target actions
for use with inverse kinematics processors. for use with inverse kinematics processors.
src/lerobot/processor/device_processor.py
+1 -1
View File
@@ -27,7 +27,7 @@ from .pipeline import ProcessorStep, ProcessorStepRegistry
@ProcessorStepRegistry.register("device_processor") @ProcessorStepRegistry.register("device_processor")
@dataclass @dataclass
class DeviceProcessor(ProcessorStep): class DeviceProcessorStep(ProcessorStep):
"""Processes transitions by moving tensors to the specified device and optionally converting float dtypes. """Processes transitions by moving tensors to the specified device and optionally converting float dtypes.
This processor ensures that all tensors in the transition are moved to the This processor ensures that all tensors in the transition are moved to the
src/lerobot/processor/gym_action_processor.py
+2 -2
View File
@@ -24,7 +24,7 @@ from .pipeline import ActionProcessorStep, ProcessorStepRegistry
@ProcessorStepRegistry.register("torch2numpy_action_processor") @ProcessorStepRegistry.register("torch2numpy_action_processor")
@dataclass @dataclass
class Torch2NumpyActionProcessor(ActionProcessorStep): class Torch2NumpyActionProcessorStep(ActionProcessorStep):
"""Convert PyTorch tensor actions to NumPy arrays.""" """Convert PyTorch tensor actions to NumPy arrays."""
squeeze_batch_dim: bool = True squeeze_batch_dim: bool = True
@@ -56,7 +56,7 @@ class Torch2NumpyActionProcessor(ActionProcessorStep):
@ProcessorStepRegistry.register("numpy2torch_action_processor") @ProcessorStepRegistry.register("numpy2torch_action_processor")
@dataclass @dataclass
class Numpy2TorchActionProcessor(ActionProcessorStep): class Numpy2TorchActionProcessorStep(ActionProcessorStep):
"""Convert NumPy array action to PyTorch tensor.""" """Convert NumPy array action to PyTorch tensor."""
def action(self, action: np.ndarray) -> torch.Tensor: def action(self, action: np.ndarray) -> torch.Tensor:
src/lerobot/processor/hil_processor.py
+7 -7
View File
@@ -29,7 +29,7 @@ TELEOP_ACTION_KEY = "teleop_action"
@ProcessorStepRegistry.register("add_teleop_action_as_complementary_data") @ProcessorStepRegistry.register("add_teleop_action_as_complementary_data")
@dataclass @dataclass
class AddTeleopActionAsComplimentaryData(ComplementaryDataProcessorStep): class AddTeleopActionAsComplimentaryDataStep(ComplementaryDataProcessorStep):
"""Add teleoperator action to transition complementary data.""" """Add teleoperator action to transition complementary data."""
teleop_device: Teleoperator teleop_device: Teleoperator
@@ -45,7 +45,7 @@ class AddTeleopActionAsComplimentaryData(ComplementaryDataProcessorStep):
@ProcessorStepRegistry.register("add_teleop_action_as_info") @ProcessorStepRegistry.register("add_teleop_action_as_info")
@dataclass @dataclass
class AddTeleopEventsAsInfo(InfoProcessorStep): class AddTeleopEventsAsInfoStep(InfoProcessorStep):
"""Add teleoperator control events to transition info.""" """Add teleoperator control events to transition info."""
teleop_device: Teleoperator teleop_device: Teleoperator
@@ -62,7 +62,7 @@ class AddTeleopEventsAsInfo(InfoProcessorStep):
@ProcessorStepRegistry.register("image_crop_resize_processor") @ProcessorStepRegistry.register("image_crop_resize_processor")
@dataclass @dataclass
class ImageCropResizeProcessor(ObservationProcessorStep): class ImageCropResizeProcessorStep(ObservationProcessorStep):
"""Crop and resize image observations.""" """Crop and resize image observations."""
crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
@@ -112,7 +112,7 @@ class ImageCropResizeProcessor(ObservationProcessorStep):
@dataclass @dataclass
@ProcessorStepRegistry.register("time_limit_processor") @ProcessorStepRegistry.register("time_limit_processor")
class TimeLimitProcessor(TruncatedProcessorStep): class TimeLimitProcessorStep(TruncatedProcessorStep):
"""Track episode steps and enforce time limits.""" """Track episode steps and enforce time limits."""
max_episode_steps: int max_episode_steps: int
@@ -139,7 +139,7 @@ class TimeLimitProcessor(TruncatedProcessorStep):
@dataclass @dataclass
@ProcessorStepRegistry.register("gripper_penalty_processor") @ProcessorStepRegistry.register("gripper_penalty_processor")
class GripperPenaltyProcessor(ComplementaryDataProcessorStep): class GripperPenaltyProcessorStep(ComplementaryDataProcessorStep):
"""Apply penalty for inappropriate gripper usage.""" """Apply penalty for inappropriate gripper usage."""
penalty: float = -0.01 penalty: float = -0.01
@@ -188,7 +188,7 @@ class GripperPenaltyProcessor(ComplementaryDataProcessorStep):
@dataclass @dataclass
@ProcessorStepRegistry.register("intervention_action_processor") @ProcessorStepRegistry.register("intervention_action_processor")
class InterventionActionProcessor(ProcessorStep): class InterventionActionProcessorStep(ProcessorStep):
"""Handle human intervention actions and episode termination.""" """Handle human intervention actions and episode termination."""
use_gripper: bool = False use_gripper: bool = False
@@ -261,7 +261,7 @@ class InterventionActionProcessor(ProcessorStep):
@dataclass @dataclass
@ProcessorStepRegistry.register("reward_classifier_processor") @ProcessorStepRegistry.register("reward_classifier_processor")
class RewardClassifierProcessor(ProcessorStep): class RewardClassifierProcessorStep(ProcessorStep):
"""Apply reward classification to image observations.""" """Apply reward classification to image observations."""
pretrained_path: str | None = None pretrained_path: str | None = None
src/lerobot/processor/joint_observations_processor.py
+2 -2
View File
@@ -14,7 +14,7 @@ from lerobot.robots import Robot
@dataclass @dataclass
@ProcessorStepRegistry.register("joint_velocity_processor") @ProcessorStepRegistry.register("joint_velocity_processor")
class JointVelocityProcessor(ObservationProcessorStep): class JointVelocityProcessorStep(ObservationProcessorStep):
"""Add joint velocity information to observations.""" """Add joint velocity information to observations."""
dt: float = 0.1 dt: float = 0.1
@@ -67,7 +67,7 @@ class JointVelocityProcessor(ObservationProcessorStep):
@dataclass @dataclass
@ProcessorStepRegistry.register("current_processor") @ProcessorStepRegistry.register("current_processor")
class MotorCurrentProcessor(ObservationProcessorStep): class MotorCurrentProcessorStep(ObservationProcessorStep):
"""Add motor current information to observations.""" """Add motor current information to observations."""
robot: Robot | None = None robot: Robot | None = None
src/lerobot/processor/migrate_policy_normalization.py
+12 -12
View File
@@ -47,11 +47,11 @@ from safetensors.torch import load_file as load_safetensors
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from .batch_processor import ToBatchProcessor from .batch_processor import AddBatchDimensionProcessorStep
from .device_processor import DeviceProcessor from .device_processor import DeviceProcessorStep
from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
from .pipeline import DataProcessorPipeline from .pipeline import DataProcessorPipeline
from .rename_processor import RenameProcessor from .rename_processor import RenameProcessorStep
# Policy type to class mapping # Policy type to class mapping
POLICY_CLASSES = { POLICY_CLASSES = {
@@ -404,8 +404,8 @@ def main():
# Now create preprocessor and postprocessor with cleaned_config available # Now create preprocessor and postprocessor with cleaned_config available
print("Creating preprocessor and postprocessor...") print("Creating preprocessor and postprocessor...")
# The pattern from existing processor factories: # The pattern from existing processor factories:
# - Preprocessor has two NormalizerProcessors: one for input_features, one for output_features # - Preprocessor has two NormalizerProcessorSteps: one for input_features, one for output_features
# - Postprocessor has one UnnormalizerProcessor for output_features only # - Postprocessor has one UnnormalizerProcessorStep for output_features only
# Get features from cleaned_config (now they're PolicyFeature objects) # Get features from cleaned_config (now they're PolicyFeature objects)
input_features = cleaned_config.get("input_features", {}) input_features = cleaned_config.get("input_features", {})
@@ -413,21 +413,21 @@ def main():
# Create preprocessor with two normalizers (following the pattern from processor factories) # Create preprocessor with two normalizers (following the pattern from processor factories)
preprocessor_steps = [ preprocessor_steps = [
RenameProcessor(rename_map={}), RenameProcessorStep(rename_map={}),
NormalizerProcessor( NormalizerProcessorStep(
features={**input_features, **output_features}, features={**input_features, **output_features},
norm_map=norm_map, norm_map=norm_map,
stats=stats, stats=stats,
), ),
ToBatchProcessor(), AddBatchDimensionProcessorStep(),
DeviceProcessor(device=policy_config.device), DeviceProcessorStep(device=policy_config.device),
] ]
preprocessor = DataProcessorPipeline(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"), DeviceProcessorStep(device="cpu"),
UnnormalizerProcessor(features=output_features, norm_map=norm_map, stats=stats), UnnormalizerProcessorStep(features=output_features, norm_map=norm_map, stats=stats),
] ]
postprocessor = DataProcessorPipeline(steps=postprocessor_steps, name="robot_postprocessor") postprocessor = DataProcessorPipeline(steps=postprocessor_steps, name="robot_postprocessor")
src/lerobot/processor/normalize_processor.py
+5 -5
View File
@@ -161,7 +161,7 @@ class _NormalizationMixin:
@dataclass @dataclass
@ProcessorStepRegistry.register(name="normalizer_processor") @ProcessorStepRegistry.register(name="normalizer_processor")
class NormalizerProcessor(_NormalizationMixin, ProcessorStep): class NormalizerProcessorStep(_NormalizationMixin, ProcessorStep):
""" """
A processor that applies normalization to observations and actions in a transition. A processor that applies normalization to observations and actions in a transition.
@@ -180,7 +180,7 @@ class NormalizerProcessor(_NormalizationMixin, ProcessorStep):
normalize_observation_keys: set[str] | None = None, normalize_observation_keys: set[str] | None = None,
eps: float = 1e-8, eps: float = 1e-8,
device: torch.device | str | None = None, device: torch.device | str | None = None,
) -> NormalizerProcessor: ) -> NormalizerProcessorStep:
return cls( return cls(
features=features, features=features,
norm_map=norm_map, norm_map=norm_map,
@@ -213,7 +213,7 @@ class NormalizerProcessor(_NormalizationMixin, ProcessorStep):
@dataclass @dataclass
@ProcessorStepRegistry.register(name="unnormalizer_processor") @ProcessorStepRegistry.register(name="unnormalizer_processor")
class UnnormalizerProcessor(_NormalizationMixin, ProcessorStep): class UnnormalizerProcessorStep(_NormalizationMixin, ProcessorStep):
""" """
A processor that applies unnormalization (the inverse of normalization) to A processor that applies unnormalization (the inverse of normalization) to
observations and actions in a transition. observations and actions in a transition.
@@ -230,7 +230,7 @@ class UnnormalizerProcessor(_NormalizationMixin, ProcessorStep):
norm_map: dict[FeatureType, NormalizationMode], norm_map: dict[FeatureType, NormalizationMode],
*, *,
device: torch.device | str | None = None, device: torch.device | str | None = None,
) -> UnnormalizerProcessor: ) -> UnnormalizerProcessorStep:
return cls(features=features, norm_map=norm_map, stats=dataset.meta.stats, device=device) return cls(features=features, norm_map=norm_map, stats=dataset.meta.stats, device=device)
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
@@ -259,7 +259,7 @@ def hotswap_stats(
Replaces normalization statistics in a RobotProcessor pipeline. Replaces normalization statistics in a RobotProcessor pipeline.
This function creates a deep copy of the provided `RobotProcessor` and updates the This function creates a deep copy of the provided `RobotProcessor` and updates the
statistics of any `NormalizerProcessor` or `UnnormalizerProcessor` steps within it. statistics of any `NormalizerProcessorStep` or `UnnormalizerProcessorStep` steps within it.
It's useful for adapting a trained policy to a new environment or dataset with It's useful for adapting a trained policy to a new environment or dataset with
different data distributions. different data distributions.
""" """
src/lerobot/processor/observation_processor.py
+1 -1
View File
@@ -28,7 +28,7 @@ from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
@dataclass @dataclass
@ProcessorStepRegistry.register(name="observation_processor") @ProcessorStepRegistry.register(name="observation_processor")
class VanillaObservationProcessor(ObservationProcessorStep): class VanillaObservationProcessorStep(ObservationProcessorStep):
""" """
Processes environment observations into the LeRobot format by handling both images and states. Processes environment observations into the LeRobot format by handling both images and states.
src/lerobot/processor/rename_processor.py
+1 -1
View File
@@ -24,7 +24,7 @@ from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
@dataclass @dataclass
@ProcessorStepRegistry.register(name="rename_processor") @ProcessorStepRegistry.register(name="rename_processor")
class RenameProcessor(ObservationProcessorStep): class RenameProcessorStep(ObservationProcessorStep):
"""Rename processor that renames keys in the observation.""" """Rename processor that renames keys in the observation."""
rename_map: dict[str, str] = field(default_factory=dict) rename_map: dict[str, str] = field(default_factory=dict)
src/lerobot/processor/tokenizer_processor.py
+4 -4
View File
@@ -24,7 +24,7 @@ else:
@dataclass @dataclass
@ProcessorStepRegistry.register(name="tokenizer_processor") @ProcessorStepRegistry.register(name="tokenizer_processor")
class TokenizerProcessor(ObservationProcessorStep): class TokenizerProcessorStep(ObservationProcessorStep):
"""Tokenizes text tasks in complementary data using a huggingface tokenizer. """Tokenizes text tasks in complementary data using a huggingface tokenizer.
This processor handles tokenization of task strings found in the complementary_data This processor handles tokenization of task strings found in the complementary_data
@@ -48,7 +48,7 @@ class TokenizerProcessor(ObservationProcessorStep):
Examples: Examples:
Using tokenizer name (auto-loaded): Using tokenizer name (auto-loaded):
```python ```python
processor = TokenizerProcessor(tokenizer_name="bert-base-uncased", max_length=128) processor = TokenizerProcessorStep(tokenizer_name="bert-base-uncased", max_length=128)
``` ```
Using custom tokenizer object: Using custom tokenizer object:
@@ -56,7 +56,7 @@ class TokenizerProcessor(ObservationProcessorStep):
from transformers import AutoTokenizer from transformers import AutoTokenizer
custom_tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") custom_tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
processor = TokenizerProcessor(tokenizer=custom_tokenizer, max_length=128) processor = TokenizerProcessorStep(tokenizer=custom_tokenizer, max_length=128)
``` ```
""" """
@@ -76,7 +76,7 @@ class TokenizerProcessor(ObservationProcessorStep):
if not _transformers_available: if not _transformers_available:
raise ImportError( raise ImportError(
"The 'transformers' library is not installed. " "The 'transformers' library is not installed. "
"Please install it with `pip install 'lerobot[transformers-dep]'` to use TokenizerProcessor." "Please install it with `pip install 'lerobot[transformers-dep]'` to use TokenizerProcessorStep."
) )
if self.tokenizer is not None: if self.tokenizer is not None:
src/lerobot/scripts/rl/gym_manipulator.py
+36 -36
View File
@@ -29,25 +29,25 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.envs.configs import HILSerlRobotEnvConfig from lerobot.envs.configs import HILSerlRobotEnvConfig
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import ( from lerobot.processor import (
AddTeleopActionAsComplimentaryData, AddBatchDimensionProcessorStep,
AddTeleopEventsAsInfo, AddTeleopActionAsComplimentaryDataStep,
AddTeleopEventsAsInfoStep,
DataProcessorPipeline, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
EnvTransition, EnvTransition,
GripperPenaltyProcessor, GripperPenaltyProcessorStep,
ImageCropResizeProcessor, ImageCropResizeProcessorStep,
InterventionActionProcessor, InterventionActionProcessorStep,
JointVelocityProcessor, JointVelocityProcessorStep,
MapDeltaActionToRobotAction, MapDeltaActionToRobotActionStep,
MapTensorToDeltaActionDict, MapTensorToDeltaActionDictStep,
MotorCurrentProcessor, MotorCurrentProcessorStep,
Numpy2TorchActionProcessor, Numpy2TorchActionProcessorStep,
RewardClassifierProcessor, RewardClassifierProcessorStep,
TimeLimitProcessor, TimeLimitProcessorStep,
ToBatchProcessor, Torch2NumpyActionProcessorStep,
Torch2NumpyActionProcessor,
TransitionKey, TransitionKey,
VanillaObservationProcessor, VanillaObservationProcessorStep,
create_transition, create_transition,
) )
from lerobot.robots import ( # noqa: F401 from lerobot.robots import ( # noqa: F401
@@ -362,16 +362,16 @@ def make_processors(
if cfg.name == "gym_hil": if cfg.name == "gym_hil":
action_pipeline_steps = [ action_pipeline_steps = [
InterventionActionProcessor(terminate_on_success=terminate_on_success), InterventionActionProcessorStep(terminate_on_success=terminate_on_success),
Torch2NumpyActionProcessor(), Torch2NumpyActionProcessorStep(),
] ]
# Minimal processor pipeline for GymHIL simulation # Minimal processor pipeline for GymHIL simulation
env_pipeline_steps = [ env_pipeline_steps = [
Numpy2TorchActionProcessor(), Numpy2TorchActionProcessorStep(),
VanillaObservationProcessor(), VanillaObservationProcessorStep(),
ToBatchProcessor(), AddBatchDimensionProcessorStep(),
DeviceProcessor(device=device), DeviceProcessorStep(device=device),
] ]
return DataProcessorPipeline(steps=env_pipeline_steps), DataProcessorPipeline( return DataProcessorPipeline(steps=env_pipeline_steps), DataProcessorPipeline(
@@ -391,13 +391,13 @@ def make_processors(
joint_names=motor_names, joint_names=motor_names,
) )
env_pipeline_steps = [VanillaObservationProcessor()] env_pipeline_steps = [VanillaObservationProcessorStep()]
if cfg.processor.observation is not None: if cfg.processor.observation is not None:
if cfg.processor.observation.add_joint_velocity_to_observation: if cfg.processor.observation.add_joint_velocity_to_observation:
env_pipeline_steps.append(JointVelocityProcessor(dt=1.0 / cfg.fps)) env_pipeline_steps.append(JointVelocityProcessorStep(dt=1.0 / cfg.fps))
if cfg.processor.observation.add_current_to_observation: if cfg.processor.observation.add_current_to_observation:
env_pipeline_steps.append(MotorCurrentProcessor(robot=env.robot)) env_pipeline_steps.append(MotorCurrentProcessorStep(robot=env.robot))
if kinematics_solver is not None: if kinematics_solver is not None:
env_pipeline_steps.append( env_pipeline_steps.append(
@@ -409,7 +409,7 @@ def make_processors(
if cfg.processor.image_preprocessing is not None: if cfg.processor.image_preprocessing is not None:
env_pipeline_steps.append( env_pipeline_steps.append(
ImageCropResizeProcessor( ImageCropResizeProcessorStep(
crop_params_dict=cfg.processor.image_preprocessing.crop_params_dict, crop_params_dict=cfg.processor.image_preprocessing.crop_params_dict,
resize_size=cfg.processor.image_preprocessing.resize_size, resize_size=cfg.processor.image_preprocessing.resize_size,
) )
@@ -418,13 +418,13 @@ def make_processors(
# Add time limit processor if reset config exists # Add time limit processor if reset config exists
if cfg.processor.reset is not None: if cfg.processor.reset is not None:
env_pipeline_steps.append( env_pipeline_steps.append(
TimeLimitProcessor(max_episode_steps=int(cfg.processor.reset.control_time_s * cfg.fps)) TimeLimitProcessorStep(max_episode_steps=int(cfg.processor.reset.control_time_s * cfg.fps))
) )
# Add gripper penalty processor if gripper config exists and enabled # Add gripper penalty processor if gripper config exists and enabled
if cfg.processor.gripper is not None and cfg.processor.gripper.use_gripper: if cfg.processor.gripper is not None and cfg.processor.gripper.use_gripper:
env_pipeline_steps.append( env_pipeline_steps.append(
GripperPenaltyProcessor( GripperPenaltyProcessorStep(
penalty=cfg.processor.gripper.gripper_penalty, penalty=cfg.processor.gripper.gripper_penalty,
max_gripper_pos=cfg.processor.max_gripper_pos, max_gripper_pos=cfg.processor.max_gripper_pos,
) )
@@ -435,7 +435,7 @@ def make_processors(
and cfg.processor.reward_classifier.pretrained_path is not None and cfg.processor.reward_classifier.pretrained_path is not None
): ):
env_pipeline_steps.append( env_pipeline_steps.append(
RewardClassifierProcessor( RewardClassifierProcessorStep(
pretrained_path=cfg.processor.reward_classifier.pretrained_path, pretrained_path=cfg.processor.reward_classifier.pretrained_path,
device=device, device=device,
success_threshold=cfg.processor.reward_classifier.success_threshold, success_threshold=cfg.processor.reward_classifier.success_threshold,
@@ -444,14 +444,14 @@ def make_processors(
) )
) )
env_pipeline_steps.append(ToBatchProcessor()) env_pipeline_steps.append(AddBatchDimensionProcessorStep())
env_pipeline_steps.append(DeviceProcessor(device=device)) env_pipeline_steps.append(DeviceProcessorStep(device=device))
action_pipeline_steps = [ action_pipeline_steps = [
AddTeleopActionAsComplimentaryData(teleop_device=teleop_device), AddTeleopActionAsComplimentaryDataStep(teleop_device=teleop_device),
AddTeleopEventsAsInfo(teleop_device=teleop_device), AddTeleopEventsAsInfoStep(teleop_device=teleop_device),
AddRobotObservationAsComplimentaryData(robot=env.robot), AddRobotObservationAsComplimentaryData(robot=env.robot),
InterventionActionProcessor( InterventionActionProcessorStep(
use_gripper=cfg.processor.gripper.use_gripper if cfg.processor.gripper is not None else False, use_gripper=cfg.processor.gripper.use_gripper if cfg.processor.gripper is not None else False,
terminate_on_success=terminate_on_success, terminate_on_success=terminate_on_success,
), ),
@@ -461,10 +461,10 @@ def make_processors(
if cfg.processor.inverse_kinematics is not None and kinematics_solver is not None: if cfg.processor.inverse_kinematics is not None and kinematics_solver is not None:
# Add EE bounds and safety processor # Add EE bounds and safety processor
inverse_kinematics_steps = [ inverse_kinematics_steps = [
MapTensorToDeltaActionDict( MapTensorToDeltaActionDictStep(
use_gripper=cfg.processor.gripper.use_gripper if cfg.processor.gripper is not None else False use_gripper=cfg.processor.gripper.use_gripper if cfg.processor.gripper is not None else False
), ),
MapDeltaActionToRobotAction(), MapDeltaActionToRobotActionStep(),
EEReferenceAndDelta( EEReferenceAndDelta(
kinematics=kinematics_solver, kinematics=kinematics_solver,
end_effector_step_sizes=cfg.processor.inverse_kinematics.end_effector_step_sizes, end_effector_step_sizes=cfg.processor.inverse_kinematics.end_effector_step_sizes,
tests/processor/test_act_processor.py
+14 -14
View File
@@ -25,13 +25,13 @@ 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 (
AddBatchDimensionProcessorStep,
DataProcessorPipeline, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
RenameProcessor, RenameProcessorStep,
ToBatchProcessor,
TransitionKey, TransitionKey,
UnnormalizerProcessor, UnnormalizerProcessorStep,
) )
@@ -86,15 +86,15 @@ def test_make_act_processor_basic():
# Check steps in preprocessor # Check steps in preprocessor
assert len(preprocessor.steps) == 4 assert len(preprocessor.steps) == 4
assert isinstance(preprocessor.steps[0], RenameProcessor) assert isinstance(preprocessor.steps[0], RenameProcessorStep)
assert isinstance(preprocessor.steps[1], NormalizerProcessor) assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
assert isinstance(preprocessor.steps[2], ToBatchProcessor) assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep)
assert isinstance(preprocessor.steps[3], DeviceProcessor) assert isinstance(preprocessor.steps[3], DeviceProcessorStep)
# Check steps in postprocessor # Check steps in postprocessor
assert len(postprocessor.steps) == 2 assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor) assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_act_processor_normalization(): def test_act_processor_normalization():
@@ -303,11 +303,11 @@ def test_act_processor_mixed_precision():
postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x}, postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
) )
# Replace DeviceProcessor with one that uses float16 # Replace DeviceProcessorStep with one that uses float16
modified_steps = [] modified_steps = []
for step in preprocessor.steps: for step in preprocessor.steps:
if isinstance(step, DeviceProcessor): if isinstance(step, DeviceProcessorStep):
modified_steps.append(DeviceProcessor(device=config.device, float_dtype="float16")) modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16"))
else: else:
modified_steps.append(step) modified_steps.append(step)
preprocessor.steps = modified_steps preprocessor.steps = modified_steps
tests/processor/test_batch_processor.py
+61 -56
View File
@@ -22,7 +22,12 @@ 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 DataProcessorPipeline, ProcessorStepRegistry, ToBatchProcessor, TransitionKey from lerobot.processor import (
AddBatchDimensionProcessorStep,
DataProcessorPipeline,
ProcessorStepRegistry,
TransitionKey,
)
def create_transition( def create_transition(
@@ -42,7 +47,7 @@ def create_transition(
def test_state_1d_to_2d(): def test_state_1d_to_2d():
"""Test that 1D state tensors get unsqueezed to 2D.""" """Test that 1D state tensors get unsqueezed to 2D."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Test observation.state # Test observation.state
state_1d = torch.randn(7) state_1d = torch.randn(7)
@@ -58,7 +63,7 @@ def test_state_1d_to_2d():
def test_env_state_1d_to_2d(): def test_env_state_1d_to_2d():
"""Test that 1D environment state tensors get unsqueezed to 2D.""" """Test that 1D environment state tensors get unsqueezed to 2D."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Test observation.environment_state # Test observation.environment_state
env_state_1d = torch.randn(10) env_state_1d = torch.randn(10)
@@ -74,7 +79,7 @@ def test_env_state_1d_to_2d():
def test_image_3d_to_4d(): def test_image_3d_to_4d():
"""Test that 3D image tensors get unsqueezed to 4D.""" """Test that 3D image tensors get unsqueezed to 4D."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Test observation.image # Test observation.image
image_3d = torch.randn(224, 224, 3) image_3d = torch.randn(224, 224, 3)
@@ -90,7 +95,7 @@ def test_image_3d_to_4d():
def test_multiple_images_3d_to_4d(): def test_multiple_images_3d_to_4d():
"""Test that 3D image tensors in observation.images.* get unsqueezed to 4D.""" """Test that 3D image tensors in observation.images.* get unsqueezed to 4D."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Test observation.images.camera1 and observation.images.camera2 # Test observation.images.camera1 and observation.images.camera2
image1_3d = torch.randn(64, 64, 3) image1_3d = torch.randn(64, 64, 3)
@@ -115,7 +120,7 @@ def test_multiple_images_3d_to_4d():
def test_already_batched_tensors_unchanged(): def test_already_batched_tensors_unchanged():
"""Test that already batched tensors remain unchanged.""" """Test that already batched tensors remain unchanged."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Create already batched tensors # Create already batched tensors
state_2d = torch.randn(1, 7) state_2d = torch.randn(1, 7)
@@ -141,7 +146,7 @@ def test_already_batched_tensors_unchanged():
def test_higher_dimensional_tensors_unchanged(): def test_higher_dimensional_tensors_unchanged():
"""Test that tensors with more dimensions than expected remain unchanged.""" """Test that tensors with more dimensions than expected remain unchanged."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Create tensors with more dimensions # Create tensors with more dimensions
state_3d = torch.randn(2, 7, 5) # More than 1D state_3d = torch.randn(2, 7, 5) # More than 1D
@@ -164,7 +169,7 @@ def test_higher_dimensional_tensors_unchanged():
def test_non_tensor_values_unchanged(): def test_non_tensor_values_unchanged():
"""Test that non-tensor values in observations remain unchanged.""" """Test that non-tensor values in observations remain unchanged."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
observation = { observation = {
OBS_STATE: [1, 2, 3], # List, not tensor OBS_STATE: [1, 2, 3], # List, not tensor
@@ -187,7 +192,7 @@ def test_non_tensor_values_unchanged():
def test_none_observation(): def test_none_observation():
"""Test processor handles None observation gracefully.""" """Test processor handles None observation gracefully."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
transition = create_transition(observation=None) transition = create_transition(observation=None)
result = processor(transition) result = processor(transition)
@@ -197,7 +202,7 @@ def test_none_observation():
def test_empty_observation(): def test_empty_observation():
"""Test processor handles empty observation dict.""" """Test processor handles empty observation dict."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
observation = {} observation = {}
transition = create_transition(observation=observation) transition = create_transition(observation=observation)
@@ -209,7 +214,7 @@ def test_empty_observation():
def test_mixed_observation(): def test_mixed_observation():
"""Test processor with mixed observation containing various types and dimensions.""" """Test processor with mixed observation containing various types and dimensions."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
state_1d = torch.randn(5) state_1d = torch.randn(5)
env_state_2d = torch.randn(1, 8) # Already batched env_state_2d = torch.randn(1, 8) # Already batched
@@ -241,8 +246,8 @@ def test_mixed_observation():
def test_integration_with_robot_processor(): def test_integration_with_robot_processor():
"""Test ToBatchProcessor integration with RobotProcessor.""" """Test AddBatchDimensionProcessorStep integration with RobotProcessor."""
to_batch_processor = ToBatchProcessor() to_batch_processor = AddBatchDimensionProcessorStep()
pipeline = DataProcessorPipeline([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
@@ -261,7 +266,7 @@ def test_integration_with_robot_processor():
def test_serialization_methods(): def test_serialization_methods():
"""Test get_config, state_dict, load_state_dict, and reset methods.""" """Test get_config, state_dict, load_state_dict, and reset methods."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Test get_config # Test get_config
config = processor.get_config() config = processor.get_config()
@@ -281,8 +286,8 @@ 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 AddBatchDimensionProcessorStep with RobotProcessor."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
pipeline = DataProcessorPipeline( 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
) )
@@ -302,7 +307,7 @@ def test_save_and_load_pretrained():
assert loaded_pipeline.name == "BatchPipeline" assert loaded_pipeline.name == "BatchPipeline"
assert len(loaded_pipeline) == 1 assert len(loaded_pipeline) == 1
assert isinstance(loaded_pipeline.steps[0], ToBatchProcessor) assert isinstance(loaded_pipeline.steps[0], AddBatchDimensionProcessorStep)
# Test functionality of loaded processor # Test functionality of loaded processor
observation = {OBS_STATE: torch.randn(5)} observation = {OBS_STATE: torch.randn(5)}
@@ -313,10 +318,10 @@ def test_save_and_load_pretrained():
def test_registry_functionality(): def test_registry_functionality():
"""Test that ToBatchProcessor is properly registered.""" """Test that AddBatchDimensionProcessorStep is properly registered."""
# Check that the processor is registered # Check that the processor is registered
registered_class = ProcessorStepRegistry.get("to_batch_processor") registered_class = ProcessorStepRegistry.get("to_batch_processor")
assert registered_class is ToBatchProcessor assert registered_class is AddBatchDimensionProcessorStep
# Check that it's in the list of registered processors # Check that it's in the list of registered processors
assert "to_batch_processor" in ProcessorStepRegistry.list() assert "to_batch_processor" in ProcessorStepRegistry.list()
@@ -324,7 +329,7 @@ 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 = AddBatchDimensionProcessorStep()
pipeline = DataProcessorPipeline([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:
@@ -350,7 +355,7 @@ def test_registry_based_save_load():
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_device_compatibility(): def test_device_compatibility():
"""Test processor works with tensors on different devices.""" """Test processor works with tensors on different devices."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Create tensors on GPU # Create tensors on GPU
state_1d = torch.randn(7, device="cuda") state_1d = torch.randn(7, device="cuda")
@@ -374,7 +379,7 @@ def test_device_compatibility():
def test_processor_preserves_other_transition_keys(): def test_processor_preserves_other_transition_keys():
"""Test that processor only modifies observation and preserves other transition keys.""" """Test that processor only modifies observation and preserves other transition keys."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
action = torch.randn(5) action = torch.randn(5)
reward = 1.5 reward = 1.5
@@ -411,7 +416,7 @@ def test_processor_preserves_other_transition_keys():
def test_edge_case_zero_dimensional_tensors(): def test_edge_case_zero_dimensional_tensors():
"""Test processor handles 0D tensors (scalars) correctly.""" """Test processor handles 0D tensors (scalars) correctly."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# 0D tensors should not be modified # 0D tensors should not be modified
scalar_tensor = torch.tensor(42.0) scalar_tensor = torch.tensor(42.0)
@@ -433,7 +438,7 @@ def test_edge_case_zero_dimensional_tensors():
# Action-specific tests # Action-specific tests
def test_action_1d_to_2d(): def test_action_1d_to_2d():
"""Test that 1D action tensors get batch dimension added.""" """Test that 1D action tensors get batch dimension added."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Create 1D action tensor # Create 1D action tensor
action_1d = torch.randn(4) action_1d = torch.randn(4)
@@ -448,7 +453,7 @@ def test_action_1d_to_2d():
def test_action_already_batched(): def test_action_already_batched():
"""Test that already batched action tensors remain unchanged.""" """Test that already batched action tensors remain unchanged."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Test various batch sizes # Test various batch sizes
action_batched_1 = torch.randn(1, 4) action_batched_1 = torch.randn(1, 4)
@@ -467,7 +472,7 @@ def test_action_already_batched():
def test_action_higher_dimensional(): def test_action_higher_dimensional():
"""Test that higher dimensional action tensors remain unchanged.""" """Test that higher dimensional action tensors remain unchanged."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# 3D action tensor (e.g., sequence of actions) # 3D action tensor (e.g., sequence of actions)
action_3d = torch.randn(2, 4, 3) action_3d = torch.randn(2, 4, 3)
@@ -484,7 +489,7 @@ def test_action_higher_dimensional():
def test_action_scalar_tensor(): def test_action_scalar_tensor():
"""Test that scalar (0D) action tensors remain unchanged.""" """Test that scalar (0D) action tensors remain unchanged."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
action_scalar = torch.tensor(1.5) action_scalar = torch.tensor(1.5)
transition = create_transition(action=action_scalar) transition = create_transition(action=action_scalar)
@@ -497,7 +502,7 @@ def test_action_scalar_tensor():
def test_action_non_tensor(): def test_action_non_tensor():
"""Test that non-tensor actions remain unchanged.""" """Test that non-tensor actions remain unchanged."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# List action # List action
action_list = [0.1, 0.2, 0.3, 0.4] action_list = [0.1, 0.2, 0.3, 0.4]
@@ -526,7 +531,7 @@ def test_action_non_tensor():
def test_action_none(): def test_action_none():
"""Test that None action is handled correctly.""" """Test that None action is handled correctly."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
transition = create_transition(action=None) transition = create_transition(action=None)
result = processor(transition) result = processor(transition)
@@ -535,7 +540,7 @@ def test_action_none():
def test_action_with_observation(): def test_action_with_observation():
"""Test action processing together with observation processing.""" """Test action processing together with observation processing."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Both need batching # Both need batching
observation = { observation = {
@@ -555,7 +560,7 @@ def test_action_with_observation():
def test_action_different_sizes(): def test_action_different_sizes():
"""Test action processing with various action dimensions.""" """Test action processing with various action dimensions."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Different action sizes (robot with different DOF) # Different action sizes (robot with different DOF)
action_sizes = [1, 2, 4, 7, 10, 20] action_sizes = [1, 2, 4, 7, 10, 20]
@@ -572,7 +577,7 @@ def test_action_different_sizes():
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_action_device_compatibility(): def test_action_device_compatibility():
"""Test action processing on different devices.""" """Test action processing on different devices."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# CUDA action # CUDA action
action_cuda = torch.randn(4, device="cuda") action_cuda = torch.randn(4, device="cuda")
@@ -593,7 +598,7 @@ def test_action_device_compatibility():
def test_action_dtype_preservation(): def test_action_dtype_preservation():
"""Test that action dtype is preserved during processing.""" """Test that action dtype is preserved during processing."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Different dtypes # Different dtypes
dtypes = [torch.float32, torch.float64, torch.int32, torch.int64] dtypes = [torch.float32, torch.float64, torch.int32, torch.int64]
@@ -609,7 +614,7 @@ def test_action_dtype_preservation():
def test_empty_action_tensor(): def test_empty_action_tensor():
"""Test handling of empty action tensors.""" """Test handling of empty action tensors."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Empty 1D tensor # Empty 1D tensor
action_empty = torch.tensor([]) action_empty = torch.tensor([])
@@ -631,7 +636,7 @@ def test_empty_action_tensor():
# Task-specific tests # Task-specific tests
def test_task_string_to_list(): def test_task_string_to_list():
"""Test that string tasks get wrapped in lists to add batch dimension.""" """Test that string tasks get wrapped in lists to add batch dimension."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Create complementary data with string task # Create complementary data with string task
complementary_data = {"task": "pick_cube"} complementary_data = {"task": "pick_cube"}
@@ -648,7 +653,7 @@ def test_task_string_to_list():
def test_task_string_validation(): def test_task_string_validation():
"""Test that only string and list of strings are valid task values.""" """Test that only string and list of strings are valid task values."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Valid string task - should be converted to list # Valid string task - should be converted to list
complementary_data = {"task": "valid_task"} complementary_data = {"task": "valid_task"}
@@ -667,7 +672,7 @@ def test_task_string_validation():
def test_task_list_of_strings(): def test_task_list_of_strings():
"""Test that lists of strings remain unchanged (already batched).""" """Test that lists of strings remain unchanged (already batched)."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Test various list of strings # Test various list of strings
test_lists = [ test_lists = [
@@ -693,7 +698,7 @@ def test_task_list_of_strings():
def test_complementary_data_none(): def test_complementary_data_none():
"""Test processor handles None complementary_data gracefully.""" """Test processor handles None complementary_data gracefully."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
transition = create_transition(complementary_data=None) transition = create_transition(complementary_data=None)
result = processor(transition) result = processor(transition)
@@ -703,7 +708,7 @@ def test_complementary_data_none():
def test_complementary_data_empty(): def test_complementary_data_empty():
"""Test processor handles empty complementary_data dict.""" """Test processor handles empty complementary_data dict."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
complementary_data = {} complementary_data = {}
transition = create_transition(complementary_data=complementary_data) transition = create_transition(complementary_data=complementary_data)
@@ -715,7 +720,7 @@ def test_complementary_data_empty():
def test_complementary_data_no_task(): def test_complementary_data_no_task():
"""Test processor handles complementary_data without task field.""" """Test processor handles complementary_data without task field."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
complementary_data = { complementary_data = {
"episode_id": 123, "episode_id": 123,
@@ -733,7 +738,7 @@ def test_complementary_data_no_task():
def test_complementary_data_mixed(): def test_complementary_data_mixed():
"""Test processor with mixed complementary_data containing task and other fields.""" """Test processor with mixed complementary_data containing task and other fields."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
complementary_data = { complementary_data = {
"task": "stack_blocks", "task": "stack_blocks",
@@ -758,7 +763,7 @@ def test_complementary_data_mixed():
def test_task_with_observation_and_action(): def test_task_with_observation_and_action():
"""Test task processing together with observation and action processing.""" """Test task processing together with observation and action processing."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# All components need batching # All components need batching
observation = { observation = {
@@ -783,7 +788,7 @@ def test_task_with_observation_and_action():
def test_task_comprehensive_string_cases(): def test_task_comprehensive_string_cases():
"""Test task processing with comprehensive string cases and edge cases.""" """Test task processing with comprehensive string cases and edge cases."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Test various string formats # Test various string formats
string_tasks = [ string_tasks = [
@@ -841,7 +846,7 @@ def test_task_comprehensive_string_cases():
def test_task_preserves_other_keys(): def test_task_preserves_other_keys():
"""Test that task processing preserves other keys in complementary_data.""" """Test that task processing preserves other keys in complementary_data."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
complementary_data = { complementary_data = {
"task": "clean_table", "task": "clean_table",
@@ -869,7 +874,7 @@ def test_task_preserves_other_keys():
# Index and task_index specific tests # Index and task_index specific tests
def test_index_scalar_to_1d(): def test_index_scalar_to_1d():
"""Test that 0D index tensor gets unsqueezed to 1D.""" """Test that 0D index tensor gets unsqueezed to 1D."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Create 0D index tensor (scalar) # Create 0D index tensor (scalar)
index_0d = torch.tensor(42, dtype=torch.int64) index_0d = torch.tensor(42, dtype=torch.int64)
@@ -886,7 +891,7 @@ def test_index_scalar_to_1d():
def test_task_index_scalar_to_1d(): def test_task_index_scalar_to_1d():
"""Test that 0D task_index tensor gets unsqueezed to 1D.""" """Test that 0D task_index tensor gets unsqueezed to 1D."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Create 0D task_index tensor (scalar) # Create 0D task_index tensor (scalar)
task_index_0d = torch.tensor(7, dtype=torch.int64) task_index_0d = torch.tensor(7, dtype=torch.int64)
@@ -903,7 +908,7 @@ def test_task_index_scalar_to_1d():
def test_index_and_task_index_together(): def test_index_and_task_index_together():
"""Test processing both index and task_index together.""" """Test processing both index and task_index together."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Create 0D tensors for both # Create 0D tensors for both
index_0d = torch.tensor(100, dtype=torch.int64) index_0d = torch.tensor(100, dtype=torch.int64)
@@ -933,7 +938,7 @@ def test_index_and_task_index_together():
def test_index_already_batched(): def test_index_already_batched():
"""Test that already batched index tensors remain unchanged.""" """Test that already batched index tensors remain unchanged."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Create already batched tensors # Create already batched tensors
index_1d = torch.tensor([42], dtype=torch.int64) index_1d = torch.tensor([42], dtype=torch.int64)
@@ -954,7 +959,7 @@ def test_index_already_batched():
def test_task_index_already_batched(): def test_task_index_already_batched():
"""Test that already batched task_index tensors remain unchanged.""" """Test that already batched task_index tensors remain unchanged."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Create already batched tensors # Create already batched tensors
task_index_1d = torch.tensor([7], dtype=torch.int64) task_index_1d = torch.tensor([7], dtype=torch.int64)
@@ -975,7 +980,7 @@ def test_task_index_already_batched():
def test_index_non_tensor_unchanged(): def test_index_non_tensor_unchanged():
"""Test that non-tensor index values remain unchanged.""" """Test that non-tensor index values remain unchanged."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
complementary_data = { complementary_data = {
"index": 42, # Plain int, not tensor "index": 42, # Plain int, not tensor
@@ -992,7 +997,7 @@ def test_index_non_tensor_unchanged():
def test_index_dtype_preservation(): def test_index_dtype_preservation():
"""Test that index and task_index dtype is preserved during processing.""" """Test that index and task_index dtype is preserved during processing."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Test different dtypes # Test different dtypes
dtypes = [torch.int32, torch.int64, torch.long] dtypes = [torch.int32, torch.int64, torch.long]
@@ -1015,7 +1020,7 @@ def test_index_dtype_preservation():
def test_index_with_full_transition(): def test_index_with_full_transition():
"""Test index/task_index processing with full transition data.""" """Test index/task_index processing with full transition data."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Create full transition with all components # Create full transition with all components
observation = { observation = {
@@ -1057,7 +1062,7 @@ def test_index_with_full_transition():
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_index_device_compatibility(): def test_index_device_compatibility():
"""Test processor works with index/task_index tensors on different devices.""" """Test processor works with index/task_index tensors on different devices."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Create tensors on GPU # Create tensors on GPU
index_0d = torch.tensor(42, dtype=torch.int64, device="cuda") index_0d = torch.tensor(42, dtype=torch.int64, device="cuda")
@@ -1081,7 +1086,7 @@ def test_index_device_compatibility():
def test_empty_index_tensor(): def test_empty_index_tensor():
"""Test handling of empty index tensors.""" """Test handling of empty index tensors."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
# Empty 0D tensor doesn't make sense, but test empty 1D # Empty 0D tensor doesn't make sense, but test empty 1D
index_empty = torch.tensor([], dtype=torch.int64) index_empty = torch.tensor([], dtype=torch.int64)
@@ -1096,7 +1101,7 @@ def test_empty_index_tensor():
def test_action_processing_creates_new_transition(): def test_action_processing_creates_new_transition():
"""Test that the processor creates a new transition object with correctly processed action.""" """Test that the processor creates a new transition object with correctly processed action."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
action = torch.randn(4) action = torch.randn(4)
transition = create_transition(action=action) transition = create_transition(action=action)
@@ -1118,7 +1123,7 @@ def test_action_processing_creates_new_transition():
def test_task_processing_creates_new_transition(): def test_task_processing_creates_new_transition():
"""Test that the processor creates a new transition object with correctly processed task.""" """Test that the processor creates a new transition object with correctly processed task."""
processor = ToBatchProcessor() processor = AddBatchDimensionProcessorStep()
complementary_data = {"task": "sort_objects"} complementary_data = {"task": "sort_objects"}
transition = create_transition(complementary_data=complementary_data) transition = create_transition(complementary_data=complementary_data)
tests/processor/test_classifier_processor.py
+9 -9
View File
@@ -26,9 +26,9 @@ from lerobot.policies.sac.reward_model.configuration_classifier import RewardCla
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, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
IdentityProcessorStep, IdentityProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
TransitionKey, TransitionKey,
) )
@@ -87,13 +87,13 @@ def test_make_classifier_processor_basic():
# Check steps in preprocessor # Check steps in preprocessor
assert len(preprocessor.steps) == 3 assert len(preprocessor.steps) == 3
assert isinstance(preprocessor.steps[0], NormalizerProcessor) # For input features assert isinstance(preprocessor.steps[0], NormalizerProcessorStep) # For input features
assert isinstance(preprocessor.steps[1], NormalizerProcessor) # For output features assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) # For output features
assert isinstance(preprocessor.steps[2], DeviceProcessor) assert isinstance(preprocessor.steps[2], DeviceProcessorStep)
# Check steps in postprocessor # Check steps in postprocessor
assert len(postprocessor.steps) == 2 assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor) assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], IdentityProcessorStep) assert isinstance(postprocessor.steps[1], IdentityProcessorStep)
@@ -291,11 +291,11 @@ def test_classifier_processor_mixed_precision():
# Get the steps from the factory function # Get the steps from the factory function
factory_preprocessor, factory_postprocessor = make_classifier_processor(config, stats) factory_preprocessor, factory_postprocessor = make_classifier_processor(config, stats)
# Replace DeviceProcessor with one that uses float16 # Replace DeviceProcessorStep with one that uses float16
modified_steps = [] modified_steps = []
for step in factory_preprocessor.steps: for step in factory_preprocessor.steps:
if isinstance(step, DeviceProcessor): if isinstance(step, DeviceProcessorStep):
modified_steps.append(DeviceProcessor(device=config.device, float_dtype="float16")) modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16"))
else: else:
modified_steps.append(step) modified_steps.append(step)
tests/processor/test_device_processor.py
+68 -64
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 DataProcessorPipeline, DeviceProcessor, TransitionKey from lerobot.processor import DataProcessorPipeline, DeviceProcessorStep, TransitionKey
def create_transition( def create_transition(
@@ -46,7 +46,7 @@ def create_transition(
def test_basic_functionality(): def test_basic_functionality():
"""Test basic device processor functionality on CPU.""" """Test basic device processor functionality on CPU."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
# Create a transition with CPU tensors # Create a transition with CPU tensors
observation = {"observation.state": torch.randn(10), "observation.image": torch.randn(3, 224, 224)} observation = {"observation.state": torch.randn(10), "observation.image": torch.randn(3, 224, 224)}
@@ -73,7 +73,7 @@ def test_basic_functionality():
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_cuda_functionality(): def test_cuda_functionality():
"""Test device processor functionality on CUDA.""" """Test device processor functionality on CUDA."""
processor = DeviceProcessor(device="cuda") processor = DeviceProcessorStep(device="cuda")
# Create a transition with CPU tensors # Create a transition with CPU tensors
observation = {"observation.state": torch.randn(10), "observation.image": torch.randn(3, 224, 224)} observation = {"observation.state": torch.randn(10), "observation.image": torch.randn(3, 224, 224)}
@@ -100,7 +100,7 @@ def test_cuda_functionality():
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_specific_cuda_device(): def test_specific_cuda_device():
"""Test device processor with specific CUDA device.""" """Test device processor with specific CUDA device."""
processor = DeviceProcessor(device="cuda:0") processor = DeviceProcessorStep(device="cuda:0")
observation = {"observation.state": torch.randn(10)} observation = {"observation.state": torch.randn(10)}
action = torch.randn(5) action = torch.randn(5)
@@ -116,7 +116,7 @@ def test_specific_cuda_device():
def test_non_tensor_values(): def test_non_tensor_values():
"""Test that non-tensor values are preserved.""" """Test that non-tensor values are preserved."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
observation = { observation = {
"observation.state": torch.randn(10), "observation.state": torch.randn(10),
@@ -142,7 +142,7 @@ def test_non_tensor_values():
def test_none_values(): def test_none_values():
"""Test handling of None values.""" """Test handling of None values."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
# Test with None observation # Test with None observation
transition = create_transition(observation=None, action=torch.randn(5)) transition = create_transition(observation=None, action=torch.randn(5))
@@ -159,7 +159,7 @@ def test_none_values():
def test_empty_observation(): def test_empty_observation():
"""Test handling of empty observation dictionary.""" """Test handling of empty observation dictionary."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
transition = create_transition(observation={}, action=torch.randn(5)) transition = create_transition(observation={}, action=torch.randn(5))
result = processor(transition) result = processor(transition)
@@ -170,7 +170,7 @@ def test_empty_observation():
def test_scalar_tensors(): def test_scalar_tensors():
"""Test handling of scalar tensors.""" """Test handling of scalar tensors."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
observation = {"observation.scalar": torch.tensor(1.5)} observation = {"observation.scalar": torch.tensor(1.5)}
action = torch.tensor(2.0) action = torch.tensor(2.0)
@@ -187,7 +187,7 @@ def test_scalar_tensors():
def test_dtype_preservation(): def test_dtype_preservation():
"""Test that tensor dtypes are preserved.""" """Test that tensor dtypes are preserved."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
observation = { observation = {
"observation.float32": torch.randn(5, dtype=torch.float32), "observation.float32": torch.randn(5, dtype=torch.float32),
@@ -209,7 +209,7 @@ def test_dtype_preservation():
def test_shape_preservation(): def test_shape_preservation():
"""Test that tensor shapes are preserved.""" """Test that tensor shapes are preserved."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
observation = { observation = {
"observation.1d": torch.randn(10), "observation.1d": torch.randn(10),
@@ -232,7 +232,7 @@ def test_shape_preservation():
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_mixed_devices(): def test_mixed_devices():
"""Test handling of tensors already on different devices.""" """Test handling of tensors already on different devices."""
processor = DeviceProcessor(device="cuda") processor = DeviceProcessorStep(device="cuda")
# Create tensors on different devices # Create tensors on different devices
observation = { observation = {
@@ -253,22 +253,22 @@ def test_mixed_devices():
def test_non_blocking_flag(): def test_non_blocking_flag():
"""Test that non_blocking flag is set correctly.""" """Test that non_blocking flag is set correctly."""
# CPU processor should have non_blocking=False # CPU processor should have non_blocking=False
cpu_processor = DeviceProcessor(device="cpu") cpu_processor = DeviceProcessorStep(device="cpu")
assert cpu_processor.non_blocking is False assert cpu_processor.non_blocking is False
if torch.cuda.is_available(): if torch.cuda.is_available():
# CUDA processor should have non_blocking=True # CUDA processor should have non_blocking=True
cuda_processor = DeviceProcessor(device="cuda") cuda_processor = DeviceProcessorStep(device="cuda")
assert cuda_processor.non_blocking is True assert cuda_processor.non_blocking is True
cuda_0_processor = DeviceProcessor(device="cuda:0") cuda_0_processor = DeviceProcessorStep(device="cuda:0")
assert cuda_0_processor.non_blocking is True assert cuda_0_processor.non_blocking is True
def test_serialization_methods(): def test_serialization_methods():
"""Test get_config, state_dict, and load_state_dict methods.""" """Test get_config, state_dict, and load_state_dict methods."""
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"
processor = DeviceProcessor(device=device) processor = DeviceProcessorStep(device=device)
# Test get_config # Test get_config
config = processor.get_config() config = processor.get_config()
@@ -289,7 +289,7 @@ def test_serialization_methods():
def test_features(): def test_features():
"""Test that features returns features unchanged.""" """Test that features returns features unchanged."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
features = { features = {
"observation.state": PolicyFeature(type=FeatureType.STATE, shape=(10,)), "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(10,)),
@@ -304,11 +304,11 @@ def test_features():
def test_integration_with_robot_processor(): def test_integration_with_robot_processor():
"""Test integration with RobotProcessor.""" """Test integration with RobotProcessor."""
from lerobot.constants import OBS_STATE from lerobot.constants import OBS_STATE
from lerobot.processor import ToBatchProcessor from lerobot.processor import AddBatchDimensionProcessorStep
# Create a pipeline with DeviceProcessor # Create a pipeline with DeviceProcessorStep
device_processor = DeviceProcessor(device="cpu") device_processor = DeviceProcessorStep(device="cpu")
batch_processor = ToBatchProcessor() batch_processor = AddBatchDimensionProcessorStep()
processor = DataProcessorPipeline( processor = DataProcessorPipeline(
steps=[batch_processor, device_processor], steps=[batch_processor, device_processor],
@@ -333,9 +333,9 @@ def test_integration_with_robot_processor():
def test_save_and_load_pretrained(): def test_save_and_load_pretrained():
"""Test saving and loading processor with DeviceProcessor.""" """Test saving and loading processor with DeviceProcessorStep."""
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 = DeviceProcessorStep(device=device, float_dtype="float16")
robot_processor = DataProcessorPipeline(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:
@@ -347,7 +347,7 @@ def test_save_and_load_pretrained():
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]
assert isinstance(loaded_device_processor, DeviceProcessor) assert isinstance(loaded_device_processor, DeviceProcessorStep)
# Use getattr to access attributes safely # Use getattr to access attributes safely
assert ( assert (
getattr(loaded_device_processor, "device", None) == device.split(":")[0] getattr(loaded_device_processor, "device", None) == device.split(":")[0]
@@ -356,18 +356,18 @@ def test_save_and_load_pretrained():
def test_registry_functionality(): def test_registry_functionality():
"""Test that DeviceProcessor is properly registered.""" """Test that DeviceProcessorStep is properly registered."""
from lerobot.processor import ProcessorStepRegistry from lerobot.processor import ProcessorStepRegistry
# Check that DeviceProcessor is registered # Check that DeviceProcessorStep is registered
registered_class = ProcessorStepRegistry.get("device_processor") registered_class = ProcessorStepRegistry.get("device_processor")
assert registered_class is DeviceProcessor assert registered_class is DeviceProcessorStep
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_performance_with_large_tensors(): def test_performance_with_large_tensors():
"""Test performance with large tensors and non_blocking flag.""" """Test performance with large tensors and non_blocking flag."""
processor = DeviceProcessor(device="cuda") processor = DeviceProcessorStep(device="cuda")
# Create large tensors # Create large tensors
observation = { observation = {
@@ -389,7 +389,7 @@ def test_performance_with_large_tensors():
def test_reward_done_truncated_types(): def test_reward_done_truncated_types():
"""Test handling of different types for reward, done, and truncated.""" """Test handling of different types for reward, done, and truncated."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
# Test with scalar values (not tensors) # Test with scalar values (not tensors)
transition = create_transition( transition = create_transition(
@@ -429,7 +429,7 @@ def test_reward_done_truncated_types():
def test_complementary_data_preserved(): def test_complementary_data_preserved():
"""Test that complementary_data is preserved unchanged.""" """Test that complementary_data is preserved unchanged."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
complementary_data = { complementary_data = {
"task": "pick_object", "task": "pick_object",
@@ -449,13 +449,13 @@ def test_complementary_data_preserved():
assert result[TransitionKey.COMPLEMENTARY_DATA]["task"] == "pick_object" assert result[TransitionKey.COMPLEMENTARY_DATA]["task"] == "pick_object"
assert result[TransitionKey.COMPLEMENTARY_DATA]["episode_id"] == 42 assert result[TransitionKey.COMPLEMENTARY_DATA]["episode_id"] == 42
assert result[TransitionKey.COMPLEMENTARY_DATA]["metadata"] == {"sensor": "camera_1"} assert result[TransitionKey.COMPLEMENTARY_DATA]["metadata"] == {"sensor": "camera_1"}
# Note: Currently DeviceProcessor doesn't process tensors in complementary_data # Note: Currently DeviceProcessorStep doesn't process tensors in complementary_data
# This is intentional as complementary_data is typically metadata # This is intentional as complementary_data is typically metadata
def test_float_dtype_conversion(): def test_float_dtype_conversion():
"""Test float dtype conversion functionality.""" """Test float dtype conversion functionality."""
processor = DeviceProcessor(device="cpu", float_dtype="float16") processor = DeviceProcessorStep(device="cpu", float_dtype="float16")
# Create tensors of different types # Create tensors of different types
observation = { observation = {
@@ -485,7 +485,7 @@ def test_float_dtype_conversion():
def test_float_dtype_none(): def test_float_dtype_none():
"""Test that when float_dtype is None, no dtype conversion occurs.""" """Test that when float_dtype is None, no dtype conversion occurs."""
processor = DeviceProcessor(device="cpu", float_dtype=None) processor = DeviceProcessorStep(device="cpu", float_dtype=None)
observation = { observation = {
"observation.float32": torch.randn(5, dtype=torch.float32), "observation.float32": torch.randn(5, dtype=torch.float32),
@@ -506,7 +506,7 @@ def test_float_dtype_none():
def test_float_dtype_bfloat16(): def test_float_dtype_bfloat16():
"""Test conversion to bfloat16.""" """Test conversion to bfloat16."""
processor = DeviceProcessor(device="cpu", float_dtype="bfloat16") processor = DeviceProcessorStep(device="cpu", float_dtype="bfloat16")
observation = {"observation.state": torch.randn(5, dtype=torch.float32)} observation = {"observation.state": torch.randn(5, dtype=torch.float32)}
action = torch.randn(3, dtype=torch.float64) action = torch.randn(3, dtype=torch.float64)
@@ -520,7 +520,7 @@ def test_float_dtype_bfloat16():
def test_float_dtype_float64(): def test_float_dtype_float64():
"""Test conversion to float64.""" """Test conversion to float64."""
processor = DeviceProcessor(device="cpu", float_dtype="float64") processor = DeviceProcessorStep(device="cpu", float_dtype="float64")
observation = {"observation.state": torch.randn(5, dtype=torch.float16)} observation = {"observation.state": torch.randn(5, dtype=torch.float16)}
action = torch.randn(3, dtype=torch.float32) action = torch.randn(3, dtype=torch.float32)
@@ -535,27 +535,27 @@ def test_float_dtype_float64():
def test_float_dtype_invalid(): def test_float_dtype_invalid():
"""Test that invalid float_dtype raises ValueError.""" """Test that invalid float_dtype raises ValueError."""
with pytest.raises(ValueError, match="Invalid float_dtype 'invalid_dtype'"): with pytest.raises(ValueError, match="Invalid float_dtype 'invalid_dtype'"):
DeviceProcessor(device="cpu", float_dtype="invalid_dtype") DeviceProcessorStep(device="cpu", float_dtype="invalid_dtype")
def test_float_dtype_aliases(): def test_float_dtype_aliases():
"""Test that dtype aliases work correctly.""" """Test that dtype aliases work correctly."""
# Test 'half' alias for float16 # Test 'half' alias for float16
processor_half = DeviceProcessor(device="cpu", float_dtype="half") processor_half = DeviceProcessorStep(device="cpu", float_dtype="half")
assert processor_half._target_float_dtype == torch.float16 assert processor_half._target_float_dtype == torch.float16
# Test 'float' alias for float32 # Test 'float' alias for float32
processor_float = DeviceProcessor(device="cpu", float_dtype="float") processor_float = DeviceProcessorStep(device="cpu", float_dtype="float")
assert processor_float._target_float_dtype == torch.float32 assert processor_float._target_float_dtype == torch.float32
# Test 'double' alias for float64 # Test 'double' alias for float64
processor_double = DeviceProcessor(device="cpu", float_dtype="double") processor_double = DeviceProcessorStep(device="cpu", float_dtype="double")
assert processor_double._target_float_dtype == torch.float64 assert processor_double._target_float_dtype == torch.float64
def test_float_dtype_with_mixed_tensors(): def test_float_dtype_with_mixed_tensors():
"""Test float dtype conversion with mixed tensor types.""" """Test float dtype conversion with mixed tensor types."""
processor = DeviceProcessor(device="cpu", float_dtype="float32") processor = DeviceProcessorStep(device="cpu", float_dtype="float32")
observation = { observation = {
"observation.image": torch.randint(0, 255, (3, 64, 64), dtype=torch.uint8), # Should not convert "observation.image": torch.randint(0, 255, (3, 64, 64), dtype=torch.uint8), # Should not convert
@@ -579,13 +579,13 @@ def test_float_dtype_with_mixed_tensors():
def test_float_dtype_serialization(): def test_float_dtype_serialization():
"""Test that float_dtype is properly serialized in get_config.""" """Test that float_dtype is properly serialized in get_config."""
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"
processor = DeviceProcessor(device=device, float_dtype="float16") processor = DeviceProcessorStep(device=device, float_dtype="float16")
config = processor.get_config() config = processor.get_config()
assert config == {"device": device, "float_dtype": "float16"} assert config == {"device": device, "float_dtype": "float16"}
# Test with None float_dtype # Test with None float_dtype
processor_none = DeviceProcessor(device="cpu", float_dtype=None) processor_none = DeviceProcessorStep(device="cpu", float_dtype=None)
config_none = processor_none.get_config() config_none = processor_none.get_config()
assert config_none == {"device": "cpu", "float_dtype": None} assert config_none == {"device": "cpu", "float_dtype": None}
@@ -594,7 +594,7 @@ def test_float_dtype_serialization():
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_float_dtype_with_cuda(): def test_float_dtype_with_cuda():
"""Test float dtype conversion combined with CUDA device.""" """Test float dtype conversion combined with CUDA device."""
processor = DeviceProcessor(device="cuda", float_dtype="float16") processor = DeviceProcessorStep(device="cuda", float_dtype="float16")
# Create tensors on CPU with different dtypes # Create tensors on CPU with different dtypes
observation = { observation = {
@@ -619,7 +619,7 @@ def test_float_dtype_with_cuda():
def test_complementary_data_index_fields(): def test_complementary_data_index_fields():
"""Test processing of index and task_index fields in complementary_data.""" """Test processing of index and task_index fields in complementary_data."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
# Create transition with index and task_index in complementary_data # Create transition with index and task_index in complementary_data
complementary_data = { complementary_data = {
@@ -657,7 +657,7 @@ def test_complementary_data_index_fields():
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_complementary_data_index_fields_cuda(): def test_complementary_data_index_fields_cuda():
"""Test moving index and task_index fields to CUDA.""" """Test moving index and task_index fields to CUDA."""
processor = DeviceProcessor(device="cuda:0") processor = DeviceProcessorStep(device="cuda:0")
# Create CPU tensors # Create CPU tensors
complementary_data = { complementary_data = {
@@ -679,7 +679,7 @@ def test_complementary_data_index_fields_cuda():
def test_complementary_data_without_index_fields(): def test_complementary_data_without_index_fields():
"""Test that complementary_data without index/task_index fields works correctly.""" """Test that complementary_data without index/task_index fields works correctly."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
complementary_data = { complementary_data = {
"task": ["navigate"], "task": ["navigate"],
@@ -697,7 +697,7 @@ def test_complementary_data_without_index_fields():
def test_complementary_data_mixed_tensors(): def test_complementary_data_mixed_tensors():
"""Test complementary_data with mix of tensors and non-tensors.""" """Test complementary_data with mix of tensors and non-tensors."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
complementary_data = { complementary_data = {
"task": ["pick_and_place"], "task": ["pick_and_place"],
@@ -726,7 +726,7 @@ def test_complementary_data_mixed_tensors():
def test_complementary_data_float_dtype_conversion(): def test_complementary_data_float_dtype_conversion():
"""Test that float dtype conversion doesn't affect int tensors in complementary_data.""" """Test that float dtype conversion doesn't affect int tensors in complementary_data."""
processor = DeviceProcessor(device="cpu", float_dtype="float16") processor = DeviceProcessorStep(device="cpu", float_dtype="float16")
complementary_data = { complementary_data = {
"index": torch.tensor([42], dtype=torch.int64), "index": torch.tensor([42], dtype=torch.int64),
@@ -750,7 +750,7 @@ def test_complementary_data_float_dtype_conversion():
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_complementary_data_full_pipeline_cuda(): def test_complementary_data_full_pipeline_cuda():
"""Test full transition with complementary_data on CUDA.""" """Test full transition with complementary_data on CUDA."""
processor = DeviceProcessor(device="cuda:0", float_dtype="float16") processor = DeviceProcessorStep(device="cuda:0", float_dtype="float16")
# Create full transition with mixed CPU tensors # Create full transition with mixed CPU tensors
observation = {"observation.state": torch.randn(1, 7, dtype=torch.float32)} observation = {"observation.state": torch.randn(1, 7, dtype=torch.float32)}
@@ -796,7 +796,7 @@ def test_complementary_data_full_pipeline_cuda():
def test_complementary_data_empty(): def test_complementary_data_empty():
"""Test empty complementary_data handling.""" """Test empty complementary_data handling."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
transition = create_transition( transition = create_transition(
observation={"observation.state": torch.randn(1, 7)}, observation={"observation.state": torch.randn(1, 7)},
@@ -811,7 +811,7 @@ def test_complementary_data_empty():
def test_complementary_data_none(): def test_complementary_data_none():
"""Test None complementary_data handling.""" """Test None complementary_data handling."""
processor = DeviceProcessor(device="cpu") processor = DeviceProcessorStep(device="cpu")
transition = create_transition( transition = create_transition(
observation={"observation.state": torch.randn(1, 7)}, observation={"observation.state": torch.randn(1, 7)},
@@ -826,8 +826,8 @@ def test_complementary_data_none():
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_preserves_gpu_placement(): def test_preserves_gpu_placement():
"""Test that DeviceProcessor preserves GPU placement when tensor is already on GPU.""" """Test that DeviceProcessorStep preserves GPU placement when tensor is already on GPU."""
processor = DeviceProcessor(device="cuda:0") processor = DeviceProcessorStep(device="cuda:0")
# Create tensors already on GPU # Create tensors already on GPU
observation = { observation = {
@@ -852,9 +852,9 @@ def test_preserves_gpu_placement():
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs")
def test_multi_gpu_preservation(): def test_multi_gpu_preservation():
"""Test that DeviceProcessor preserves placement on different GPUs in multi-GPU setup.""" """Test that DeviceProcessorStep preserves placement on different GPUs in multi-GPU setup."""
# Test 1: GPU-to-GPU preservation (cuda:0 config, cuda:1 input) # Test 1: GPU-to-GPU preservation (cuda:0 config, cuda:1 input)
processor_gpu = DeviceProcessor(device="cuda:0") processor_gpu = DeviceProcessorStep(device="cuda:0")
# Create tensors on cuda:1 (simulating Accelerate placement) # Create tensors on cuda:1 (simulating Accelerate placement)
cuda1_device = torch.device("cuda:1") cuda1_device = torch.device("cuda:1")
@@ -873,7 +873,7 @@ def test_multi_gpu_preservation():
assert result[TransitionKey.ACTION].device == cuda1_device assert result[TransitionKey.ACTION].device == cuda1_device
# Test 2: GPU-to-CPU should move to CPU (not preserve GPU) # Test 2: GPU-to-CPU should move to CPU (not preserve GPU)
processor_cpu = DeviceProcessor(device="cpu") processor_cpu = DeviceProcessorStep(device="cpu")
transition_gpu = create_transition( transition_gpu = create_transition(
observation={"observation.state": torch.randn(10).cuda()}, action=torch.randn(5).cuda() observation={"observation.state": torch.randn(10).cuda()}, action=torch.randn(5).cuda()
@@ -889,7 +889,7 @@ def test_multi_gpu_preservation():
def test_multi_gpu_with_cpu_tensors(): def test_multi_gpu_with_cpu_tensors():
"""Test that CPU tensors are moved to configured device even in multi-GPU context.""" """Test that CPU tensors are moved to configured device even in multi-GPU context."""
# Processor configured for cuda:1 # Processor configured for cuda:1
processor = DeviceProcessor(device="cuda:1") processor = DeviceProcessorStep(device="cuda:1")
# Mix of CPU and GPU tensors # Mix of CPU and GPU tensors
observation = { observation = {
@@ -916,7 +916,7 @@ def test_multi_gpu_with_cpu_tensors():
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs")
def test_multi_gpu_with_float_dtype(): def test_multi_gpu_with_float_dtype():
"""Test float dtype conversion works correctly with multi-GPU preservation.""" """Test float dtype conversion works correctly with multi-GPU preservation."""
processor = DeviceProcessor(device="cuda:0", float_dtype="float16") processor = DeviceProcessorStep(device="cuda:0", float_dtype="float16")
# Create float tensors on different GPUs # Create float tensors on different GPUs
observation = { observation = {
@@ -946,7 +946,7 @@ def test_simulated_accelerate_scenario():
for gpu_id in range(min(torch.cuda.device_count(), 2)): for gpu_id in range(min(torch.cuda.device_count(), 2)):
# Each "process" has a processor configured for cuda:0 # Each "process" has a processor configured for cuda:0
# but data comes in already placed on the process's GPU # but data comes in already placed on the process's GPU
processor = DeviceProcessor(device="cuda:0") processor = DeviceProcessorStep(device="cuda:0")
# Simulate data already placed by Accelerate # Simulate data already placed by Accelerate
device = torch.device(f"cuda:{gpu_id}") device = torch.device(f"cuda:{gpu_id}")
@@ -966,7 +966,11 @@ def test_policy_processor_integration():
"""Test integration with policy processors - input on GPU, output on CPU.""" """Test integration with policy processors - input on GPU, output on CPU."""
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.constants import ACTION, OBS_STATE from lerobot.constants import ACTION, OBS_STATE
from lerobot.processor import NormalizerProcessor, ToBatchProcessor, UnnormalizerProcessor from lerobot.processor import (
AddBatchDimensionProcessorStep,
NormalizerProcessorStep,
UnnormalizerProcessorStep,
)
# Create features and stats # Create features and stats
features = { features = {
@@ -984,9 +988,9 @@ def test_policy_processor_integration():
# Create input processor (preprocessor) that moves to GPU # Create input processor (preprocessor) that moves to GPU
input_processor = DataProcessorPipeline( input_processor = DataProcessorPipeline(
steps=[ steps=[
NormalizerProcessor(features=features, norm_map=norm_map, stats=stats), NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats),
ToBatchProcessor(), AddBatchDimensionProcessorStep(),
DeviceProcessor(device="cuda"), DeviceProcessorStep(device="cuda"),
], ],
name="test_preprocessor", name="test_preprocessor",
to_transition=lambda x: x, to_transition=lambda x: x,
@@ -996,8 +1000,8 @@ def test_policy_processor_integration():
# Create output processor (postprocessor) that moves to CPU # Create output processor (postprocessor) that moves to CPU
output_processor = DataProcessorPipeline( output_processor = DataProcessorPipeline(
steps=[ steps=[
DeviceProcessor(device="cpu"), DeviceProcessorStep(device="cpu"),
UnnormalizerProcessor(features={ACTION: features[ACTION]}, norm_map=norm_map, stats=stats), UnnormalizerProcessorStep(features={ACTION: features[ACTION]}, norm_map=norm_map, stats=stats),
], ],
name="test_postprocessor", name="test_postprocessor",
to_transition=lambda x: x, to_transition=lambda x: x,
tests/processor/test_diffusion_processor.py
+14 -14
View File
@@ -25,13 +25,13 @@ 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 (
AddBatchDimensionProcessorStep,
DataProcessorPipeline, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
RenameProcessor, RenameProcessorStep,
ToBatchProcessor,
TransitionKey, TransitionKey,
UnnormalizerProcessor, UnnormalizerProcessorStep,
) )
@@ -89,15 +89,15 @@ def test_make_diffusion_processor_basic():
# Check steps in preprocessor # Check steps in preprocessor
assert len(preprocessor.steps) == 4 assert len(preprocessor.steps) == 4
assert isinstance(preprocessor.steps[0], RenameProcessor) assert isinstance(preprocessor.steps[0], RenameProcessorStep)
assert isinstance(preprocessor.steps[1], NormalizerProcessor) assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
assert isinstance(preprocessor.steps[2], ToBatchProcessor) assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep)
assert isinstance(preprocessor.steps[3], DeviceProcessor) assert isinstance(preprocessor.steps[3], DeviceProcessorStep)
# Check steps in postprocessor # Check steps in postprocessor
assert len(postprocessor.steps) == 2 assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor) assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_diffusion_processor_with_images(): def test_diffusion_processor_with_images():
@@ -294,11 +294,11 @@ def test_diffusion_processor_mixed_precision():
# Get the steps from the factory function # Get the steps from the factory function
factory_preprocessor, factory_postprocessor = make_diffusion_pre_post_processors(config, stats) factory_preprocessor, factory_postprocessor = make_diffusion_pre_post_processors(config, stats)
# Replace DeviceProcessor with one that uses float16 # Replace DeviceProcessorStep with one that uses float16
modified_steps = [] modified_steps = []
for step in factory_preprocessor.steps: for step in factory_preprocessor.steps:
if isinstance(step, DeviceProcessor): if isinstance(step, DeviceProcessorStep):
modified_steps.append(DeviceProcessor(device=config.device, float_dtype="float16")) modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16"))
else: else:
modified_steps.append(step) modified_steps.append(step)
tests/processor/test_normalize_processor.py
+61 -61
View File
@@ -23,9 +23,9 @@ from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline, DataProcessorPipeline,
IdentityProcessorStep, IdentityProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
TransitionKey, TransitionKey,
UnnormalizerProcessor, UnnormalizerProcessorStep,
hotswap_stats, hotswap_stats,
) )
from lerobot.processor.converters import to_tensor from lerobot.processor.converters import to_tensor
@@ -125,7 +125,7 @@ def _create_observation_norm_map():
} }
# Fixtures for observation normalisation tests using NormalizerProcessor # Fixtures for observation normalisation tests using NormalizerProcessorStep
@pytest.fixture @pytest.fixture
def observation_stats(): def observation_stats():
return { return {
@@ -142,10 +142,10 @@ def observation_stats():
@pytest.fixture @pytest.fixture
def observation_normalizer(observation_stats): def observation_normalizer(observation_stats):
"""Return a NormalizerProcessor that only has observation stats (no action).""" """Return a NormalizerProcessorStep that only has observation stats (no action)."""
features = _create_observation_features() features = _create_observation_features()
norm_map = _create_observation_norm_map() norm_map = _create_observation_norm_map()
return NormalizerProcessor(features=features, norm_map=norm_map, stats=observation_stats) return NormalizerProcessorStep(features=features, norm_map=norm_map, stats=observation_stats)
def test_mean_std_normalization(observation_normalizer): def test_mean_std_normalization(observation_normalizer):
@@ -182,7 +182,7 @@ def test_min_max_normalization(observation_normalizer):
def test_selective_normalization(observation_stats): def test_selective_normalization(observation_stats):
features = _create_observation_features() features = _create_observation_features()
norm_map = _create_observation_norm_map() norm_map = _create_observation_norm_map()
normalizer = NormalizerProcessor( normalizer = NormalizerProcessorStep(
features=features, features=features,
norm_map=norm_map, norm_map=norm_map,
stats=observation_stats, stats=observation_stats,
@@ -208,7 +208,7 @@ def test_selective_normalization(observation_stats):
def test_device_compatibility(observation_stats): def test_device_compatibility(observation_stats):
features = _create_observation_features() features = _create_observation_features()
norm_map = _create_observation_norm_map() norm_map = _create_observation_norm_map()
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=observation_stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=observation_stats)
observation = { observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]).cuda(), "observation.image": torch.tensor([0.7, 0.5, 0.3]).cuda(),
} }
@@ -237,7 +237,7 @@ def test_from_lerobot_dataset():
FeatureType.ACTION: NormalizationMode.MEAN_STD, FeatureType.ACTION: NormalizationMode.MEAN_STD,
} }
normalizer = NormalizerProcessor.from_lerobot_dataset(mock_dataset, features, norm_map) normalizer = NormalizerProcessorStep.from_lerobot_dataset(mock_dataset, features, norm_map)
# Both observation and action statistics should be present in tensor stats # Both observation and action statistics should be present in tensor stats
assert "observation.image" in normalizer._tensor_stats assert "observation.image" in normalizer._tensor_stats
@@ -252,7 +252,7 @@ def test_state_dict_save_load(observation_normalizer):
# Create new normalizer and load state # Create new normalizer and load state
features = _create_observation_features() features = _create_observation_features()
norm_map = _create_observation_norm_map() norm_map = _create_observation_norm_map()
new_normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats={}) new_normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats={})
new_normalizer.load_state_dict(state_dict) new_normalizer.load_state_dict(state_dict)
# Test that it works the same # Test that it works the same
@@ -303,7 +303,7 @@ def _create_action_norm_map_min_max():
def test_mean_std_unnormalization(action_stats_mean_std): def test_mean_std_unnormalization(action_stats_mean_std):
features = _create_action_features() features = _create_action_features()
norm_map = _create_action_norm_map_mean_std() norm_map = _create_action_norm_map_mean_std()
unnormalizer = UnnormalizerProcessor( unnormalizer = UnnormalizerProcessorStep(
features=features, norm_map=norm_map, stats={"action": action_stats_mean_std} features=features, norm_map=norm_map, stats={"action": action_stats_mean_std}
) )
@@ -321,7 +321,7 @@ def test_mean_std_unnormalization(action_stats_mean_std):
def test_min_max_unnormalization(action_stats_min_max): def test_min_max_unnormalization(action_stats_min_max):
features = _create_action_features() features = _create_action_features()
norm_map = _create_action_norm_map_min_max() norm_map = _create_action_norm_map_min_max()
unnormalizer = UnnormalizerProcessor( unnormalizer = UnnormalizerProcessorStep(
features=features, norm_map=norm_map, stats={"action": action_stats_min_max} features=features, norm_map=norm_map, stats={"action": action_stats_min_max}
) )
@@ -347,7 +347,7 @@ def test_min_max_unnormalization(action_stats_min_max):
def test_numpy_action_input(action_stats_mean_std): def test_numpy_action_input(action_stats_mean_std):
features = _create_action_features() features = _create_action_features()
norm_map = _create_action_norm_map_mean_std() norm_map = _create_action_norm_map_mean_std()
unnormalizer = UnnormalizerProcessor( unnormalizer = UnnormalizerProcessorStep(
features=features, norm_map=norm_map, stats={"action": action_stats_mean_std} features=features, norm_map=norm_map, stats={"action": action_stats_mean_std}
) )
@@ -365,7 +365,7 @@ def test_numpy_action_input(action_stats_mean_std):
def test_none_action(action_stats_mean_std): def test_none_action(action_stats_mean_std):
features = _create_action_features() features = _create_action_features()
norm_map = _create_action_norm_map_mean_std() norm_map = _create_action_norm_map_mean_std()
unnormalizer = UnnormalizerProcessor( unnormalizer = UnnormalizerProcessorStep(
features=features, norm_map=norm_map, stats={"action": action_stats_mean_std} features=features, norm_map=norm_map, stats={"action": action_stats_mean_std}
) )
@@ -381,11 +381,11 @@ def test_action_from_lerobot_dataset():
mock_dataset.meta.stats = {"action": {"mean": [0.0], "std": [1.0]}} mock_dataset.meta.stats = {"action": {"mean": [0.0], "std": [1.0]}}
features = {"action": PolicyFeature(FeatureType.ACTION, (1,))} features = {"action": PolicyFeature(FeatureType.ACTION, (1,))}
norm_map = {FeatureType.ACTION: NormalizationMode.MEAN_STD} norm_map = {FeatureType.ACTION: NormalizationMode.MEAN_STD}
unnormalizer = UnnormalizerProcessor.from_lerobot_dataset(mock_dataset, features, norm_map) unnormalizer = UnnormalizerProcessorStep.from_lerobot_dataset(mock_dataset, features, norm_map)
assert "mean" in unnormalizer._tensor_stats["action"] assert "mean" in unnormalizer._tensor_stats["action"]
# Fixtures for NormalizerProcessor tests # Fixtures for NormalizerProcessorStep tests
@pytest.fixture @pytest.fixture
def full_stats(): def full_stats():
return { return {
@@ -424,7 +424,7 @@ def _create_full_norm_map():
def normalizer_processor(full_stats): def normalizer_processor(full_stats):
features = _create_full_features() features = _create_full_features()
norm_map = _create_full_norm_map() norm_map = _create_full_norm_map()
return NormalizerProcessor(features=features, norm_map=norm_map, stats=full_stats) return NormalizerProcessorStep(features=features, norm_map=norm_map, stats=full_stats)
def test_combined_normalization(normalizer_processor): def test_combined_normalization(normalizer_processor):
@@ -468,7 +468,7 @@ def test_processor_from_lerobot_dataset(full_stats):
features = _create_full_features() features = _create_full_features()
norm_map = _create_full_norm_map() norm_map = _create_full_norm_map()
processor = NormalizerProcessor.from_lerobot_dataset( processor = NormalizerProcessorStep.from_lerobot_dataset(
mock_dataset, features, norm_map, normalize_observation_keys={"observation.image"} mock_dataset, features, norm_map, normalize_observation_keys={"observation.image"}
) )
@@ -480,7 +480,7 @@ def test_processor_from_lerobot_dataset(full_stats):
def test_get_config(full_stats): def test_get_config(full_stats):
features = _create_full_features() features = _create_full_features()
norm_map = _create_full_norm_map() norm_map = _create_full_norm_map()
processor = NormalizerProcessor( processor = NormalizerProcessorStep(
features=features, features=features,
norm_map=norm_map, norm_map=norm_map,
stats=full_stats, stats=full_stats,
@@ -539,7 +539,7 @@ def test_empty_observation():
stats = {"observation.image": {"mean": [0.5], "std": [0.2]}} stats = {"observation.image": {"mean": [0.5], "std": [0.2]}}
features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))} features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))}
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
transition = create_transition() transition = create_transition()
result = normalizer(transition) result = normalizer(transition)
@@ -550,7 +550,7 @@ def test_empty_observation():
def test_empty_stats(): def test_empty_stats():
features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))} features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))}
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats={}) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats={})
observation = {"observation.image": torch.tensor([0.5])} observation = {"observation.image": torch.tensor([0.5])}
transition = create_transition(observation=observation) transition = create_transition(observation=observation)
@@ -566,7 +566,7 @@ def test_partial_stats():
stats = {"observation.image": {"mean": [0.5]}} # Missing std / (min,max) stats = {"observation.image": {"mean": [0.5]}} # Missing std / (min,max)
features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))} features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))}
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
observation = {"observation.image": torch.tensor([0.7])} observation = {"observation.image": torch.tensor([0.7])}
transition = create_transition(observation=observation) transition = create_transition(observation=observation)
@@ -581,7 +581,7 @@ def test_missing_action_stats_no_error():
features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))} features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))}
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
processor = UnnormalizerProcessor.from_lerobot_dataset(mock_dataset, features, norm_map) processor = UnnormalizerProcessorStep.from_lerobot_dataset(mock_dataset, features, norm_map)
# The tensor stats should not contain the 'action' key # The tensor stats should not contain the 'action' key
assert "action" not in processor._tensor_stats assert "action" not in processor._tensor_stats
@@ -590,7 +590,7 @@ def test_serialization_roundtrip(full_stats):
"""Test that features and norm_map can be serialized and deserialized correctly.""" """Test that features and norm_map can be serialized and deserialized correctly."""
features = _create_full_features() features = _create_full_features()
norm_map = _create_full_norm_map() norm_map = _create_full_norm_map()
original_processor = NormalizerProcessor( original_processor = NormalizerProcessorStep(
features=features, features=features,
norm_map=norm_map, norm_map=norm_map,
stats=full_stats, stats=full_stats,
@@ -602,7 +602,7 @@ def test_serialization_roundtrip(full_stats):
config = original_processor.get_config() config = original_processor.get_config()
# Create a new processor from the config (deserialization) # Create a new processor from the config (deserialization)
new_processor = NormalizerProcessor( new_processor = NormalizerProcessorStep(
features=config["features"], features=config["features"],
norm_map=config["norm_map"], norm_map=config["norm_map"],
stats=full_stats, stats=full_stats,
@@ -670,7 +670,7 @@ def test_identity_normalization_observations():
"observation.state": {"mean": [0.0, 0.0], "std": [1.0, 1.0]}, "observation.state": {"mean": [0.0, 0.0], "std": [1.0, 1.0]},
} }
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
observation = { observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]), "observation.image": torch.tensor([0.7, 0.5, 0.3]),
@@ -695,7 +695,7 @@ def test_identity_normalization_actions():
norm_map = {FeatureType.ACTION: NormalizationMode.IDENTITY} norm_map = {FeatureType.ACTION: NormalizationMode.IDENTITY}
stats = {"action": {"mean": [0.0, 0.0], "std": [1.0, 2.0]}} stats = {"action": {"mean": [0.0, 0.0], "std": [1.0, 2.0]}}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
action = torch.tensor([1.0, -0.5]) action = torch.tensor([1.0, -0.5])
transition = create_transition(action=action) transition = create_transition(action=action)
@@ -721,7 +721,7 @@ def test_identity_unnormalization_observations():
"observation.state": {"min": [-1.0, -1.0], "max": [1.0, 1.0]}, "observation.state": {"min": [-1.0, -1.0], "max": [1.0, 1.0]},
} }
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
observation = { observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]), "observation.image": torch.tensor([0.7, 0.5, 0.3]),
@@ -748,7 +748,7 @@ def test_identity_unnormalization_actions():
norm_map = {FeatureType.ACTION: NormalizationMode.IDENTITY} norm_map = {FeatureType.ACTION: NormalizationMode.IDENTITY}
stats = {"action": {"min": [-1.0, -2.0], "max": [1.0, 2.0]}} stats = {"action": {"min": [-1.0, -2.0], "max": [1.0, 2.0]}}
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
action = torch.tensor([0.5, -0.8]) # Normalized values action = torch.tensor([0.5, -0.8]) # Normalized values
transition = create_transition(action=action) transition = create_transition(action=action)
@@ -771,8 +771,8 @@ def test_identity_with_missing_stats():
} }
stats = {} # No stats provided stats = {} # No stats provided
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
observation = {"observation.image": torch.tensor([0.7, 0.5, 0.3])} observation = {"observation.image": torch.tensor([0.7, 0.5, 0.3])}
action = torch.tensor([1.0, -0.5]) action = torch.tensor([1.0, -0.5])
@@ -812,7 +812,7 @@ def test_identity_mixed_with_other_modes():
"action": {"min": [-1.0, -1.0], "max": [1.0, 1.0]}, "action": {"min": [-1.0, -1.0], "max": [1.0, 1.0]},
} }
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
observation = { observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]), "observation.image": torch.tensor([0.7, 0.5, 0.3]),
@@ -854,7 +854,7 @@ def test_identity_defaults_when_not_in_norm_map():
"observation.state": {"mean": [0.0, 0.0], "std": [1.0, 1.0]}, "observation.state": {"mean": [0.0, 0.0], "std": [1.0, 1.0]},
} }
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
observation = { observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]), "observation.image": torch.tensor([0.7, 0.5, 0.3]),
@@ -888,8 +888,8 @@ def test_identity_roundtrip():
"action": {"min": [-1.0, -1.0], "max": [1.0, 1.0]}, "action": {"min": [-1.0, -1.0], "max": [1.0, 1.0]},
} }
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
original_observation = {"observation.image": torch.tensor([0.7, 0.5, 0.3])} original_observation = {"observation.image": torch.tensor([0.7, 0.5, 0.3])}
original_action = torch.tensor([0.5, -0.2]) original_action = torch.tensor([0.5, -0.2])
@@ -921,7 +921,7 @@ def test_identity_config_serialization():
"action": {"mean": [0.0, 0.0], "std": [1.0, 1.0]}, "action": {"mean": [0.0, 0.0], "std": [1.0, 1.0]},
} }
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
# Get config # Get config
config = normalizer.get_config() config = normalizer.get_config()
@@ -931,7 +931,7 @@ def test_identity_config_serialization():
assert config["norm_map"]["ACTION"] == "MEAN_STD" assert config["norm_map"]["ACTION"] == "MEAN_STD"
# Create new processor from config (simulating load) # Create new processor from config (simulating load)
new_normalizer = NormalizerProcessor( new_normalizer = NormalizerProcessorStep(
features=config["features"], features=config["features"],
norm_map=config["norm_map"], norm_map=config["norm_map"],
stats=stats, stats=stats,
@@ -969,7 +969,7 @@ def test_identity_config_serialization():
# norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD} # norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD}
# stats = {"observation.state": {"mean": [0.0, 0.0], "std": [1.0, 1.0]}} # stats = {"observation.state": {"mean": [0.0, 0.0], "std": [1.0, 1.0]}}
# normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) # normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
# # Manually inject an invalid mode to test error handling # # Manually inject an invalid mode to test error handling
# normalizer.norm_map[FeatureType.STATE] = "INVALID_MODE" # normalizer.norm_map[FeatureType.STATE] = "INVALID_MODE"
@@ -1006,8 +1006,8 @@ def test_hotswap_stats_basic_functionality():
} }
# Create processors # Create processors
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats)
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats)
identity = IdentityProcessorStep() identity = IdentityProcessorStep()
# Create robot processor # Create robot processor
@@ -1047,7 +1047,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 = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats)
original_processor = DataProcessorPipeline(steps=[normalizer]) original_processor = DataProcessorPipeline(steps=[normalizer])
# Store reference to original stats # Store reference to original stats
@@ -1072,7 +1072,7 @@ def test_hotswap_stats_deep_copy():
def test_hotswap_stats_only_affects_normalizer_steps(): def test_hotswap_stats_only_affects_normalizer_steps():
"""Test that hotswap_stats only modifies NormalizerProcessor and UnnormalizerProcessor steps.""" """Test that hotswap_stats only modifies NormalizerProcessorStep and UnnormalizerProcessorStep steps."""
stats = { stats = {
"observation.image": {"mean": np.array([0.5]), "std": np.array([0.2])}, "observation.image": {"mean": np.array([0.5]), "std": np.array([0.2])},
} }
@@ -1087,8 +1087,8 @@ def test_hotswap_stats_only_affects_normalizer_steps():
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
# Create mixed steps # Create mixed steps
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
identity = IdentityProcessorStep() identity = IdentityProcessorStep()
robot_processor = DataProcessorPipeline(steps=[normalizer, identity, unnormalizer]) robot_processor = DataProcessorPipeline(steps=[normalizer, identity, unnormalizer])
@@ -1117,7 +1117,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 = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats)
robot_processor = DataProcessorPipeline(steps=[normalizer]) robot_processor = DataProcessorPipeline(steps=[normalizer])
# Hotswap with empty stats # Hotswap with empty stats
@@ -1167,7 +1167,7 @@ def test_hotswap_stats_preserves_other_attributes():
normalize_observation_keys = {"observation.image"} normalize_observation_keys = {"observation.image"}
eps = 1e-6 eps = 1e-6
normalizer = NormalizerProcessor( normalizer = NormalizerProcessorStep(
features=features, features=features,
norm_map=norm_map, norm_map=norm_map,
stats=initial_stats, stats=initial_stats,
@@ -1212,10 +1212,10 @@ def test_hotswap_stats_multiple_normalizer_types():
} }
# Create multiple normalizers and unnormalizers # Create multiple normalizers and unnormalizers
normalizer1 = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) normalizer1 = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats)
normalizer2 = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) normalizer2 = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats)
unnormalizer1 = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) unnormalizer1 = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats)
unnormalizer2 = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) unnormalizer2 = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats)
robot_processor = DataProcessorPipeline(steps=[normalizer1, unnormalizer1, normalizer2, unnormalizer2]) robot_processor = DataProcessorPipeline(steps=[normalizer1, unnormalizer1, normalizer2, unnormalizer2])
@@ -1264,7 +1264,7 @@ def test_hotswap_stats_with_different_data_types():
FeatureType.ACTION: NormalizationMode.MEAN_STD, FeatureType.ACTION: NormalizationMode.MEAN_STD,
} }
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats)
robot_processor = DataProcessorPipeline(steps=[normalizer]) robot_processor = DataProcessorPipeline(steps=[normalizer])
# Hotswap stats # Hotswap stats
@@ -1320,7 +1320,7 @@ def test_hotswap_stats_functional_test():
} }
# Create original processor # Create original processor
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats)
original_processor = DataProcessorPipeline( original_processor = DataProcessorPipeline(
steps=[normalizer], to_transition=lambda x: x, to_output=lambda x: x steps=[normalizer], to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -1366,7 +1366,7 @@ def test_zero_std_uses_eps():
features = {"observation.state": PolicyFeature(FeatureType.STATE, (1,))} features = {"observation.state": PolicyFeature(FeatureType.STATE, (1,))}
norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD} norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD}
stats = {"observation.state": {"mean": np.array([0.5]), "std": np.array([0.0])}} stats = {"observation.state": {"mean": np.array([0.5]), "std": np.array([0.0])}}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats, eps=1e-6) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats, eps=1e-6)
observation = {"observation.state": torch.tensor([0.5])} # equals mean observation = {"observation.state": torch.tensor([0.5])} # equals mean
out = normalizer(create_transition(observation=observation)) out = normalizer(create_transition(observation=observation))
@@ -1378,7 +1378,7 @@ def test_min_equals_max_maps_to_minus_one():
features = {"observation.state": PolicyFeature(FeatureType.STATE, (1,))} features = {"observation.state": PolicyFeature(FeatureType.STATE, (1,))}
norm_map = {FeatureType.STATE: NormalizationMode.MIN_MAX} norm_map = {FeatureType.STATE: NormalizationMode.MIN_MAX}
stats = {"observation.state": {"min": np.array([2.0]), "max": np.array([2.0])}} stats = {"observation.state": {"min": np.array([2.0]), "max": np.array([2.0])}}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats, eps=1e-6) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats, eps=1e-6)
observation = {"observation.state": torch.tensor([2.0])} observation = {"observation.state": torch.tensor([2.0])}
out = normalizer(create_transition(observation=observation)) out = normalizer(create_transition(observation=observation))
@@ -1393,7 +1393,7 @@ def test_action_normalized_despite_normalize_observation_keys():
} }
norm_map = {FeatureType.STATE: NormalizationMode.IDENTITY, FeatureType.ACTION: NormalizationMode.MEAN_STD} norm_map = {FeatureType.STATE: NormalizationMode.IDENTITY, FeatureType.ACTION: NormalizationMode.MEAN_STD}
stats = {"action": {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}} stats = {"action": {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}}
normalizer = NormalizerProcessor( normalizer = NormalizerProcessorStep(
features=features, norm_map=norm_map, stats=stats, normalize_observation_keys={"observation.state"} features=features, norm_map=norm_map, stats=stats, normalize_observation_keys={"observation.state"}
) )
@@ -1411,12 +1411,12 @@ def test_unnormalize_observations_mean_std_and_min_max():
"observation.mm": PolicyFeature(FeatureType.STATE, (2,)), "observation.mm": PolicyFeature(FeatureType.STATE, (2,)),
} }
# Build two processors: one mean/std and one min/max # Build two processors: one mean/std and one min/max
unnorm_ms = UnnormalizerProcessor( unnorm_ms = UnnormalizerProcessorStep(
features={"observation.ms": features["observation.ms"]}, features={"observation.ms": features["observation.ms"]},
norm_map={FeatureType.STATE: NormalizationMode.MEAN_STD}, norm_map={FeatureType.STATE: NormalizationMode.MEAN_STD},
stats={"observation.ms": {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}}, stats={"observation.ms": {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}},
) )
unnorm_mm = UnnormalizerProcessor( unnorm_mm = UnnormalizerProcessorStep(
features={"observation.mm": features["observation.mm"]}, features={"observation.mm": features["observation.mm"]},
norm_map={FeatureType.STATE: NormalizationMode.MIN_MAX}, norm_map={FeatureType.STATE: NormalizationMode.MIN_MAX},
stats={"observation.mm": {"min": np.array([0.0, -2.0]), "max": np.array([2.0, 2.0])}}, stats={"observation.mm": {"min": np.array([0.0, -2.0]), "max": np.array([2.0, 2.0])}},
@@ -1438,7 +1438,7 @@ def test_unknown_observation_keys_ignored():
features = {"observation.state": PolicyFeature(FeatureType.STATE, (1,))} features = {"observation.state": PolicyFeature(FeatureType.STATE, (1,))}
norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD} norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD}
stats = {"observation.state": {"mean": np.array([0.0]), "std": np.array([1.0])}} stats = {"observation.state": {"mean": np.array([0.0]), "std": np.array([1.0])}}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
obs = {"observation.state": torch.tensor([1.0]), "observation.unknown": torch.tensor([5.0])} obs = {"observation.state": torch.tensor([1.0]), "observation.unknown": torch.tensor([5.0])}
tr = create_transition(observation=obs) tr = create_transition(observation=obs)
@@ -1452,7 +1452,7 @@ def test_batched_action_normalization():
features = {"action": PolicyFeature(FeatureType.ACTION, (2,))} features = {"action": PolicyFeature(FeatureType.ACTION, (2,))}
norm_map = {FeatureType.ACTION: NormalizationMode.MEAN_STD} norm_map = {FeatureType.ACTION: NormalizationMode.MEAN_STD}
stats = {"action": {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}} stats = {"action": {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
actions = torch.tensor([[1.0, -1.0], [3.0, 3.0]]) # first equals mean → zeros; second → [1, 1] actions = torch.tensor([[1.0, -1.0], [3.0, 3.0]]) # first equals mean → zeros; second → [1, 1]
out = normalizer(create_transition(action=actions))[TransitionKey.ACTION] out = normalizer(create_transition(action=actions))[TransitionKey.ACTION]
@@ -1464,7 +1464,7 @@ def test_complementary_data_preservation():
features = {"observation.state": PolicyFeature(FeatureType.STATE, (1,))} features = {"observation.state": PolicyFeature(FeatureType.STATE, (1,))}
norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD} norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD}
stats = {"observation.state": {"mean": np.array([0.0]), "std": np.array([1.0])}} stats = {"observation.state": {"mean": np.array([0.0]), "std": np.array([1.0])}}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
comp = {"existing": 123} comp = {"existing": 123}
tr = create_transition(observation={"observation.state": torch.tensor([1.0])}, complementary_data=comp) tr = create_transition(observation={"observation.state": torch.tensor([1.0])}, complementary_data=comp)
@@ -1483,8 +1483,8 @@ def test_roundtrip_normalize_unnormalize_non_identity():
"observation.state": {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}, "observation.state": {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])},
"action": {"min": np.array([-2.0, 0.0]), "max": np.array([2.0, 4.0])}, "action": {"min": np.array([-2.0, 0.0]), "max": np.array([2.0, 4.0])},
} }
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
# Add a time dimension in action for broadcasting check (B,T,D) # Add a time dimension in action for broadcasting check (B,T,D)
obs = {"observation.state": torch.tensor([[3.0, 3.0], [1.0, -1.0]])} obs = {"observation.state": torch.tensor([[3.0, 3.0], [1.0, -1.0]])}
tests/processor/test_observation_processor.py
+25 -25
View File
@@ -20,7 +20,7 @@ import torch
from lerobot.configs.types import FeatureType from lerobot.configs.types import FeatureType
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 TransitionKey, VanillaObservationProcessor from lerobot.processor import TransitionKey, VanillaObservationProcessorStep
from tests.conftest import assert_contract_is_typed from tests.conftest import assert_contract_is_typed
@@ -41,7 +41,7 @@ def create_transition(
def test_process_single_image(): def test_process_single_image():
"""Test processing a single image.""" """Test processing a single image."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
# Create a mock image (H, W, C) format, uint8 # Create a mock image (H, W, C) format, uint8
image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8) image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8)
@@ -67,7 +67,7 @@ def test_process_single_image():
def test_process_image_dict(): def test_process_image_dict():
"""Test processing multiple images in a dictionary.""" """Test processing multiple images in a dictionary."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
# Create mock images # Create mock images
image1 = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8) image1 = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8)
@@ -90,7 +90,7 @@ def test_process_image_dict():
def test_process_batched_image(): def test_process_batched_image():
"""Test processing already batched images.""" """Test processing already batched images."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
# Create a batched image (B, H, W, C) # Create a batched image (B, H, W, C)
image = np.random.randint(0, 256, size=(2, 64, 64, 3), dtype=np.uint8) image = np.random.randint(0, 256, size=(2, 64, 64, 3), dtype=np.uint8)
@@ -107,7 +107,7 @@ def test_process_batched_image():
def test_invalid_image_format(): def test_invalid_image_format():
"""Test error handling for invalid image formats.""" """Test error handling for invalid image formats."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
# Test wrong channel order (channels first) # Test wrong channel order (channels first)
image = np.random.randint(0, 256, size=(3, 64, 64), dtype=np.uint8) image = np.random.randint(0, 256, size=(3, 64, 64), dtype=np.uint8)
@@ -120,7 +120,7 @@ def test_invalid_image_format():
def test_invalid_image_dtype(): def test_invalid_image_dtype():
"""Test error handling for invalid image dtype.""" """Test error handling for invalid image dtype."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
# Test wrong dtype # Test wrong dtype
image = np.random.rand(64, 64, 3).astype(np.float32) image = np.random.rand(64, 64, 3).astype(np.float32)
@@ -133,7 +133,7 @@ def test_invalid_image_dtype():
def test_no_pixels_in_observation(): def test_no_pixels_in_observation():
"""Test processor when no pixels are in observation.""" """Test processor when no pixels are in observation."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
observation = {"other_data": np.array([1, 2, 3])} observation = {"other_data": np.array([1, 2, 3])}
transition = create_transition(observation=observation) transition = create_transition(observation=observation)
@@ -148,7 +148,7 @@ def test_no_pixels_in_observation():
def test_none_observation(): def test_none_observation():
"""Test processor with None observation.""" """Test processor with None observation."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
transition = create_transition() transition = create_transition()
result = processor(transition) result = processor(transition)
@@ -158,7 +158,7 @@ def test_none_observation():
def test_serialization_methods(): def test_serialization_methods():
"""Test serialization methods.""" """Test serialization methods."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
# Test get_config # Test get_config
config = processor.get_config() config = processor.get_config()
@@ -177,7 +177,7 @@ def test_serialization_methods():
def test_process_environment_state(): def test_process_environment_state():
"""Test processing environment_state.""" """Test processing environment_state."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32) env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32)
observation = {"environment_state": env_state} observation = {"environment_state": env_state}
@@ -198,7 +198,7 @@ def test_process_environment_state():
def test_process_agent_pos(): def test_process_agent_pos():
"""Test processing agent_pos.""" """Test processing agent_pos."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32) agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32)
observation = {"agent_pos": agent_pos} observation = {"agent_pos": agent_pos}
@@ -219,7 +219,7 @@ def test_process_agent_pos():
def test_process_batched_states(): def test_process_batched_states():
"""Test processing already batched states.""" """Test processing already batched states."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
env_state = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) env_state = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
agent_pos = np.array([[0.5, -0.5], [1.0, -1.0]], dtype=np.float32) agent_pos = np.array([[0.5, -0.5], [1.0, -1.0]], dtype=np.float32)
@@ -237,7 +237,7 @@ def test_process_batched_states():
def test_process_both_states(): def test_process_both_states():
"""Test processing both environment_state and agent_pos.""" """Test processing both environment_state and agent_pos."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
env_state = np.array([1.0, 2.0], dtype=np.float32) env_state = np.array([1.0, 2.0], dtype=np.float32)
agent_pos = np.array([0.5, -0.5], dtype=np.float32) agent_pos = np.array([0.5, -0.5], dtype=np.float32)
@@ -262,7 +262,7 @@ def test_process_both_states():
def test_no_states_in_observation(): def test_no_states_in_observation():
"""Test processor when no states are in observation.""" """Test processor when no states are in observation."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
observation = {"other_data": np.array([1, 2, 3])} observation = {"other_data": np.array([1, 2, 3])}
transition = create_transition(observation=observation) transition = create_transition(observation=observation)
@@ -276,7 +276,7 @@ def test_no_states_in_observation():
def test_complete_observation_processing(): def test_complete_observation_processing():
"""Test processing a complete observation with both images and states.""" """Test processing a complete observation with both images and states."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
# Create mock data # Create mock data
image = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8) image = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8)
@@ -313,7 +313,7 @@ def test_complete_observation_processing():
def test_image_only_processing(): def test_image_only_processing():
"""Test processing observation with only images.""" """Test processing observation with only images."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8) image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8)
observation = {"pixels": image} observation = {"pixels": image}
@@ -328,7 +328,7 @@ def test_image_only_processing():
def test_state_only_processing(): def test_state_only_processing():
"""Test processing observation with only states.""" """Test processing observation with only states."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
agent_pos = np.array([1.0, 2.0], dtype=np.float32) agent_pos = np.array([1.0, 2.0], dtype=np.float32)
observation = {"agent_pos": agent_pos} observation = {"agent_pos": agent_pos}
@@ -343,7 +343,7 @@ def test_state_only_processing():
def test_empty_observation(): def test_empty_observation():
"""Test processing empty observation.""" """Test processing empty observation."""
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
observation = {} observation = {}
transition = create_transition(observation=observation) transition = create_transition(observation=observation)
@@ -359,7 +359,7 @@ def test_equivalent_to_original_function():
# Import the original function for comparison # Import the original function for comparison
from lerobot.envs.utils import preprocess_observation from lerobot.envs.utils import preprocess_observation
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
# Create test data similar to what the original function expects # Create test data similar to what the original function expects
image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8) image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8)
@@ -386,7 +386,7 @@ def test_equivalent_with_image_dict():
"""Test equivalence with dictionary of images.""" """Test equivalence with dictionary of images."""
from lerobot.envs.utils import preprocess_observation from lerobot.envs.utils import preprocess_observation
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
# Create test data with multiple cameras # Create test data with multiple cameras
image1 = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8) image1 = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8)
@@ -410,7 +410,7 @@ def test_equivalent_with_image_dict():
def test_image_processor_features_pixels_to_image(policy_feature_factory): def test_image_processor_features_pixels_to_image(policy_feature_factory):
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
features = { features = {
"pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)), "pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
"keep": policy_feature_factory(FeatureType.ENV, (1,)), "keep": policy_feature_factory(FeatureType.ENV, (1,)),
@@ -424,7 +424,7 @@ def test_image_processor_features_pixels_to_image(policy_feature_factory):
def test_image_processor_features_observation_pixels_to_image(policy_feature_factory): def test_image_processor_features_observation_pixels_to_image(policy_feature_factory):
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
features = { features = {
"observation.pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)), "observation.pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
"keep": policy_feature_factory(FeatureType.ENV, (1,)), "keep": policy_feature_factory(FeatureType.ENV, (1,)),
@@ -438,7 +438,7 @@ def test_image_processor_features_observation_pixels_to_image(policy_feature_fac
def test_image_processor_features_multi_camera_and_prefixed(policy_feature_factory): def test_image_processor_features_multi_camera_and_prefixed(policy_feature_factory):
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
features = { features = {
"pixels.front": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)), "pixels.front": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
"pixels.wrist": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)), "pixels.wrist": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
@@ -456,7 +456,7 @@ def test_image_processor_features_multi_camera_and_prefixed(policy_feature_facto
def test_state_processor_features_environment_and_agent_pos(policy_feature_factory): def test_state_processor_features_environment_and_agent_pos(policy_feature_factory):
processor = VanillaObservationProcessor() processor = VanillaObservationProcessorStep()
features = { features = {
"environment_state": policy_feature_factory(FeatureType.STATE, (3,)), "environment_state": policy_feature_factory(FeatureType.STATE, (3,)),
"agent_pos": policy_feature_factory(FeatureType.STATE, (7,)), "agent_pos": policy_feature_factory(FeatureType.STATE, (7,)),
@@ -472,7 +472,7 @@ def test_state_processor_features_environment_and_agent_pos(policy_feature_facto
def test_state_processor_features_prefixed_inputs(policy_feature_factory): def test_state_processor_features_prefixed_inputs(policy_feature_factory):
proc = VanillaObservationProcessor() proc = VanillaObservationProcessorStep()
features = { features = {
"observation.environment_state": policy_feature_factory(FeatureType.STATE, (2,)), "observation.environment_state": policy_feature_factory(FeatureType.STATE, (2,)),
"observation.agent_pos": policy_feature_factory(FeatureType.STATE, (4,)), "observation.agent_pos": policy_feature_factory(FeatureType.STATE, (4,)),
tests/processor/test_pi0_processor.py
+20 -20
View File
@@ -25,12 +25,12 @@ from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE
from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.policies.pi0.configuration_pi0 import PI0Config
from lerobot.policies.pi0.processor_pi0 import Pi0NewLineProcessor, make_pi0_pre_post_processors from lerobot.policies.pi0.processor_pi0 import Pi0NewLineProcessor, make_pi0_pre_post_processors
from lerobot.processor import ( from lerobot.processor import (
DeviceProcessor, AddBatchDimensionProcessorStep,
NormalizerProcessor, DeviceProcessorStep,
RenameProcessor, NormalizerProcessorStep,
ToBatchProcessor, RenameProcessorStep,
TransitionKey, TransitionKey,
UnnormalizerProcessor, UnnormalizerProcessorStep,
) )
@@ -83,7 +83,7 @@ def test_make_pi0_processor_basic():
config = create_default_config() config = create_default_config()
stats = create_default_stats() stats = create_default_stats()
with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor"): with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessorStep"):
preprocessor, postprocessor = make_pi0_pre_post_processors( preprocessor, postprocessor = make_pi0_pre_post_processors(
config, config,
stats, stats,
@@ -97,17 +97,17 @@ def test_make_pi0_processor_basic():
# Check steps in preprocessor # Check steps in preprocessor
assert len(preprocessor.steps) == 6 assert len(preprocessor.steps) == 6
assert isinstance(preprocessor.steps[0], RenameProcessor) assert isinstance(preprocessor.steps[0], RenameProcessorStep)
assert isinstance(preprocessor.steps[1], NormalizerProcessor) assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
assert isinstance(preprocessor.steps[2], ToBatchProcessor) assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep)
assert isinstance(preprocessor.steps[3], Pi0NewLineProcessor) assert isinstance(preprocessor.steps[3], Pi0NewLineProcessor)
# Step 4 would be TokenizerProcessor but it's mocked # Step 4 would be TokenizerProcessorStep but it's mocked
assert isinstance(preprocessor.steps[5], DeviceProcessor) assert isinstance(preprocessor.steps[5], DeviceProcessorStep)
# Check steps in postprocessor # Check steps in postprocessor
assert len(postprocessor.steps) == 2 assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor) assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_pi0_newline_processor_single_task(): def test_pi0_newline_processor_single_task():
@@ -165,7 +165,7 @@ def test_pi0_processor_cuda():
stats = create_default_stats() stats = create_default_stats()
# Mock the tokenizer processor to act as pass-through # Mock the tokenizer processor to act as pass-through
class MockTokenizerProcessor: class MockTokenizerProcessorStep:
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
pass pass
@@ -187,7 +187,7 @@ def test_pi0_processor_cuda():
def transform_features(self, features): def transform_features(self, features):
return features return features
with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor): with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessorStep", MockTokenizerProcessorStep):
preprocessor, postprocessor = make_pi0_pre_post_processors( preprocessor, postprocessor = make_pi0_pre_post_processors(
config, config,
stats, stats,
@@ -220,7 +220,7 @@ def test_pi0_processor_accelerate_scenario():
stats = create_default_stats() stats = create_default_stats()
# Mock the tokenizer processor to act as pass-through # Mock the tokenizer processor to act as pass-through
class MockTokenizerProcessor: class MockTokenizerProcessorStep:
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
pass pass
@@ -242,7 +242,7 @@ def test_pi0_processor_accelerate_scenario():
def transform_features(self, features): def transform_features(self, features):
return features return features
with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor): with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessorStep", MockTokenizerProcessorStep):
preprocessor, postprocessor = make_pi0_pre_post_processors( preprocessor, postprocessor = make_pi0_pre_post_processors(
config, config,
stats, stats,
@@ -276,7 +276,7 @@ def test_pi0_processor_multi_gpu():
stats = create_default_stats() stats = create_default_stats()
# Mock the tokenizer processor to act as pass-through # Mock the tokenizer processor to act as pass-through
class MockTokenizerProcessor: class MockTokenizerProcessorStep:
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
pass pass
@@ -298,7 +298,7 @@ def test_pi0_processor_multi_gpu():
def transform_features(self, features): def transform_features(self, features):
return features return features
with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor): with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessorStep", MockTokenizerProcessorStep):
preprocessor, postprocessor = make_pi0_pre_post_processors( preprocessor, postprocessor = make_pi0_pre_post_processors(
config, config,
stats, stats,
@@ -329,7 +329,7 @@ def test_pi0_processor_without_stats():
config = create_default_config() config = create_default_config()
# Mock the tokenizer processor # Mock the tokenizer processor
with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor"): with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessorStep"):
preprocessor, postprocessor = make_pi0_pre_post_processors( preprocessor, postprocessor = make_pi0_pre_post_processors(
config, config,
dataset_stats=None, dataset_stats=None,
tests/processor/test_rename_processor.py
+32 -30
View File
@@ -23,7 +23,7 @@ from lerobot.configs.types import FeatureType
from lerobot.processor import ( from lerobot.processor import (
DataProcessorPipeline, DataProcessorPipeline,
ProcessorStepRegistry, ProcessorStepRegistry,
RenameProcessor, RenameProcessorStep,
TransitionKey, TransitionKey,
) )
from lerobot.processor.rename_processor import rename_stats from lerobot.processor.rename_processor import rename_stats
@@ -51,7 +51,7 @@ def test_basic_renaming():
"old_key1": "new_key1", "old_key1": "new_key1",
"old_key2": "new_key2", "old_key2": "new_key2",
} }
processor = RenameProcessor(rename_map=rename_map) processor = RenameProcessorStep(rename_map=rename_map)
observation = { observation = {
"old_key1": torch.tensor([1.0, 2.0]), "old_key1": torch.tensor([1.0, 2.0]),
@@ -79,7 +79,7 @@ def test_basic_renaming():
def test_empty_rename_map(): def test_empty_rename_map():
"""Test processor with empty rename map (should pass through unchanged).""" """Test processor with empty rename map (should pass through unchanged)."""
processor = RenameProcessor(rename_map={}) processor = RenameProcessorStep(rename_map={})
observation = { observation = {
"key1": torch.tensor([1.0]), "key1": torch.tensor([1.0]),
@@ -98,7 +98,7 @@ def test_empty_rename_map():
def test_none_observation(): def test_none_observation():
"""Test processor with None observation.""" """Test processor with None observation."""
processor = RenameProcessor(rename_map={"old": "new"}) processor = RenameProcessorStep(rename_map={"old": "new"})
transition = create_transition() transition = create_transition()
result = processor(transition) result = processor(transition)
@@ -113,7 +113,7 @@ def test_overlapping_rename():
"a": "b", "a": "b",
"b": "c", # This creates a potential conflict "b": "c", # This creates a potential conflict
} }
processor = RenameProcessor(rename_map=rename_map) processor = RenameProcessorStep(rename_map=rename_map)
observation = { observation = {
"a": 1, "a": 1,
@@ -138,7 +138,7 @@ def test_partial_rename():
"observation.state": "observation.proprio_state", "observation.state": "observation.proprio_state",
"pixels": "observation.image", "pixels": "observation.image",
} }
processor = RenameProcessor(rename_map=rename_map) processor = RenameProcessorStep(rename_map=rename_map)
observation = { observation = {
"observation.state": torch.randn(10), "observation.state": torch.randn(10),
@@ -168,15 +168,15 @@ def test_get_config():
"old1": "new1", "old1": "new1",
"old2": "new2", "old2": "new2",
} }
processor = RenameProcessor(rename_map=rename_map) processor = RenameProcessorStep(rename_map=rename_map)
config = processor.get_config() config = processor.get_config()
assert config == {"rename_map": rename_map} assert config == {"rename_map": rename_map}
def test_state_dict(): def test_state_dict():
"""Test state dict (should be empty for RenameProcessor).""" """Test state dict (should be empty for RenameProcessorStep)."""
processor = RenameProcessor(rename_map={"old": "new"}) processor = RenameProcessorStep(rename_map={"old": "new"})
state = processor.state_dict() state = processor.state_dict()
assert state == {} assert state == {}
@@ -191,7 +191,7 @@ def test_integration_with_robot_processor():
"agent_pos": "observation.state", "agent_pos": "observation.state",
"pixels": "observation.image", "pixels": "observation.image",
} }
rename_processor = RenameProcessor(rename_map=rename_map) rename_processor = RenameProcessorStep(rename_map=rename_map)
pipeline = DataProcessorPipeline([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)
@@ -225,18 +225,20 @@ def test_save_and_load_pretrained():
"old_state": "observation.state", "old_state": "observation.state",
"old_image": "observation.image", "old_image": "observation.image",
} }
processor = RenameProcessor(rename_map=rename_map) processor = RenameProcessorStep(rename_map=rename_map)
pipeline = DataProcessorPipeline([processor], name="TestRenameProcessor") pipeline = DataProcessorPipeline([processor], name="TestRenameProcessorStep")
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
# Save pipeline # Save pipeline
pipeline.save_pretrained(tmp_dir) pipeline.save_pretrained(tmp_dir)
# Check files were created # Check files were created
config_path = Path(tmp_dir) / "testrenameprocessor.json" # Based on name="TestRenameProcessor" config_path = (
Path(tmp_dir) / "testrenameprocessorstep.json"
) # Based on name="TestRenameProcessorStep"
assert config_path.exists() assert config_path.exists()
# No state files should be created for RenameProcessor # No state files should be created for RenameProcessorStep
state_files = list(Path(tmp_dir).glob("*.safetensors")) state_files = list(Path(tmp_dir).glob("*.safetensors"))
assert len(state_files) == 0 assert len(state_files) == 0
@@ -245,12 +247,12 @@ def test_save_and_load_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
) )
assert loaded_pipeline.name == "TestRenameProcessor" assert loaded_pipeline.name == "TestRenameProcessorStep"
assert len(loaded_pipeline) == 1 assert len(loaded_pipeline) == 1
# Check that loaded processor works correctly # Check that loaded processor works correctly
loaded_processor = loaded_pipeline.steps[0] loaded_processor = loaded_pipeline.steps[0]
assert isinstance(loaded_processor, RenameProcessor) assert isinstance(loaded_processor, RenameProcessorStep)
assert loaded_processor.rename_map == rename_map assert loaded_processor.rename_map == rename_map
# Test functionality after loading # Test functionality after loading
@@ -267,23 +269,23 @@ def test_save_and_load_pretrained():
def test_registry_functionality(): def test_registry_functionality():
"""Test that RenameProcessor is properly registered.""" """Test that RenameProcessorStep is properly registered."""
# Check that it's registered # Check that it's registered
assert "rename_processor" in ProcessorStepRegistry.list() assert "rename_processor" in ProcessorStepRegistry.list()
# Get from registry # Get from registry
retrieved_class = ProcessorStepRegistry.get("rename_processor") retrieved_class = ProcessorStepRegistry.get("rename_processor")
assert retrieved_class is RenameProcessor assert retrieved_class is RenameProcessorStep
# Create instance from registry # Create instance from registry
instance = retrieved_class(rename_map={"old": "new"}) instance = retrieved_class(rename_map={"old": "new"})
assert isinstance(instance, RenameProcessor) assert isinstance(instance, RenameProcessorStep)
assert instance.rename_map == {"old": "new"} assert instance.rename_map == {"old": "new"}
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 = RenameProcessorStep(rename_map={"key1": "renamed_key1"})
pipeline = DataProcessorPipeline([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:
@@ -303,14 +305,14 @@ def test_registry_based_save_load():
# Load should work # Load should work
loaded_pipeline = DataProcessorPipeline.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, RenameProcessorStep)
assert loaded_processor.rename_map == {"key1": "renamed_key1"} assert loaded_processor.rename_map == {"key1": "renamed_key1"}
def test_chained_rename_processors(): def test_chained_rename_processors():
"""Test multiple RenameProcessors in a pipeline.""" """Test multiple RenameProcessorSteps in a pipeline."""
# First processor: rename raw keys to intermediate format # First processor: rename raw keys to intermediate format
processor1 = RenameProcessor( processor1 = RenameProcessorStep(
rename_map={ rename_map={
"pos": "agent_position", "pos": "agent_position",
"img": "camera_image", "img": "camera_image",
@@ -318,7 +320,7 @@ def test_chained_rename_processors():
) )
# Second processor: rename to final format # Second processor: rename to final format
processor2 = RenameProcessor( processor2 = RenameProcessorStep(
rename_map={ rename_map={
"agent_position": "observation.state", "agent_position": "observation.state",
"camera_image": "observation.image", "camera_image": "observation.image",
@@ -363,7 +365,7 @@ def test_nested_observation_rename():
"observation.images.right": "observation.camera.right_view", "observation.images.right": "observation.camera.right_view",
"observation.proprio": "observation.proprioception", "observation.proprio": "observation.proprioception",
} }
processor = RenameProcessor(rename_map=rename_map) processor = RenameProcessorStep(rename_map=rename_map)
observation = { observation = {
"observation.images.left": torch.randn(3, 64, 64), "observation.images.left": torch.randn(3, 64, 64),
@@ -393,7 +395,7 @@ def test_nested_observation_rename():
def test_value_types_preserved(): def test_value_types_preserved():
"""Test that various value types are preserved during renaming.""" """Test that various value types are preserved during renaming."""
rename_map = {"old_tensor": "new_tensor", "old_array": "new_array", "old_scalar": "new_scalar"} rename_map = {"old_tensor": "new_tensor", "old_array": "new_array", "old_scalar": "new_scalar"}
processor = RenameProcessor(rename_map=rename_map) processor = RenameProcessorStep(rename_map=rename_map)
tensor_value = torch.randn(3, 3) tensor_value = torch.randn(3, 3)
array_value = np.random.rand(2, 2) array_value = np.random.rand(2, 2)
@@ -421,7 +423,7 @@ def test_value_types_preserved():
def test_features_basic_renaming(policy_feature_factory): def test_features_basic_renaming(policy_feature_factory):
processor = RenameProcessor(rename_map={"a": "x", "b": "y"}) processor = RenameProcessorStep(rename_map={"a": "x", "b": "y"})
features = { features = {
"a": policy_feature_factory(FeatureType.STATE, (2,)), "a": policy_feature_factory(FeatureType.STATE, (2,)),
"b": policy_feature_factory(FeatureType.ACTION, (3,)), "b": policy_feature_factory(FeatureType.ACTION, (3,)),
@@ -442,7 +444,7 @@ def test_features_basic_renaming(policy_feature_factory):
def test_features_overlapping_keys(policy_feature_factory): def test_features_overlapping_keys(policy_feature_factory):
# Overlapping renames: both 'a' and 'b' exist. 'a'->'b', 'b'->'c' # Overlapping renames: both 'a' and 'b' exist. 'a'->'b', 'b'->'c'
processor = RenameProcessor(rename_map={"a": "b", "b": "c"}) processor = RenameProcessorStep(rename_map={"a": "b", "b": "c"})
features = { features = {
"a": policy_feature_factory(FeatureType.STATE, (1,)), "a": policy_feature_factory(FeatureType.STATE, (1,)),
"b": policy_feature_factory(FeatureType.STATE, (2,)), "b": policy_feature_factory(FeatureType.STATE, (2,)),
@@ -457,8 +459,8 @@ def test_features_overlapping_keys(policy_feature_factory):
def test_features_chained_processors(policy_feature_factory): def test_features_chained_processors(policy_feature_factory):
# Chain two rename processors at the contract level # Chain two rename processors at the contract level
processor1 = RenameProcessor(rename_map={"pos": "agent_position", "img": "camera_image"}) processor1 = RenameProcessorStep(rename_map={"pos": "agent_position", "img": "camera_image"})
processor2 = RenameProcessor( processor2 = RenameProcessorStep(
rename_map={"agent_position": "observation.state", "camera_image": "observation.image"} rename_map={"agent_position": "observation.state", "camera_image": "observation.image"}
) )
pipeline = DataProcessorPipeline([processor1, processor2]) pipeline = DataProcessorPipeline([processor1, processor2])
tests/processor/test_sac_processor.py
+14 -14
View File
@@ -25,13 +25,13 @@ 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 (
AddBatchDimensionProcessorStep,
DataProcessorPipeline, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
RenameProcessor, RenameProcessorStep,
ToBatchProcessor,
TransitionKey, TransitionKey,
UnnormalizerProcessor, UnnormalizerProcessorStep,
) )
@@ -91,15 +91,15 @@ def test_make_sac_processor_basic():
# Check steps in preprocessor # Check steps in preprocessor
assert len(preprocessor.steps) == 4 assert len(preprocessor.steps) == 4
assert isinstance(preprocessor.steps[0], RenameProcessor) assert isinstance(preprocessor.steps[0], RenameProcessorStep)
assert isinstance(preprocessor.steps[1], NormalizerProcessor) assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
assert isinstance(preprocessor.steps[2], ToBatchProcessor) assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep)
assert isinstance(preprocessor.steps[3], DeviceProcessor) assert isinstance(preprocessor.steps[3], DeviceProcessorStep)
# Check steps in postprocessor # Check steps in postprocessor
assert len(postprocessor.steps) == 2 assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor) assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_sac_processor_normalization_modes(): def test_sac_processor_normalization_modes():
@@ -306,10 +306,10 @@ def test_sac_processor_mixed_precision():
postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x}, postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
) )
# Replace DeviceProcessor with one that uses float16 # Replace DeviceProcessorStep with one that uses float16
for i, step in enumerate(preprocessor.steps): for i, step in enumerate(preprocessor.steps):
if isinstance(step, DeviceProcessor): if isinstance(step, DeviceProcessorStep):
preprocessor.steps[i] = DeviceProcessor(device=config.device, float_dtype="float16") preprocessor.steps[i] = DeviceProcessorStep(device=config.device, float_dtype="float16")
# Create test data # Create test data
observation = {OBS_STATE: torch.randn(10, dtype=torch.float32)} observation = {OBS_STATE: torch.randn(10, dtype=torch.float32)}
tests/processor/test_smolvla_processor.py
+26 -20
View File
@@ -28,12 +28,12 @@ from lerobot.policies.smolvla.processor_smolvla import (
make_smolvla_pre_post_processors, make_smolvla_pre_post_processors,
) )
from lerobot.processor import ( from lerobot.processor import (
DeviceProcessor, AddBatchDimensionProcessorStep,
NormalizerProcessor, DeviceProcessorStep,
RenameProcessor, NormalizerProcessorStep,
ToBatchProcessor, RenameProcessorStep,
TransitionKey, TransitionKey,
UnnormalizerProcessor, UnnormalizerProcessorStep,
) )
@@ -88,7 +88,7 @@ def test_make_smolvla_processor_basic():
config = create_default_config() config = create_default_config()
stats = create_default_stats() stats = create_default_stats()
with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor"): with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep"):
preprocessor, postprocessor = make_smolvla_pre_post_processors( preprocessor, postprocessor = make_smolvla_pre_post_processors(
config, config,
stats, stats,
@@ -102,17 +102,17 @@ def test_make_smolvla_processor_basic():
# Check steps in preprocessor # Check steps in preprocessor
assert len(preprocessor.steps) == 6 assert len(preprocessor.steps) == 6
assert isinstance(preprocessor.steps[0], RenameProcessor) assert isinstance(preprocessor.steps[0], RenameProcessorStep)
assert isinstance(preprocessor.steps[1], NormalizerProcessor) assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
assert isinstance(preprocessor.steps[2], ToBatchProcessor) assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep)
assert isinstance(preprocessor.steps[3], SmolVLANewLineProcessor) assert isinstance(preprocessor.steps[3], SmolVLANewLineProcessor)
# Step 4 would be TokenizerProcessor but it's mocked # Step 4 would be TokenizerProcessorStep but it's mocked
assert isinstance(preprocessor.steps[5], DeviceProcessor) assert isinstance(preprocessor.steps[5], DeviceProcessorStep)
# Check steps in postprocessor # Check steps in postprocessor
assert len(postprocessor.steps) == 2 assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor) assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_smolvla_newline_processor_single_task(): def test_smolvla_newline_processor_single_task():
@@ -170,7 +170,7 @@ def test_smolvla_processor_cuda():
stats = create_default_stats() stats = create_default_stats()
# Mock the tokenizer processor to act as pass-through # Mock the tokenizer processor to act as pass-through
class MockTokenizerProcessor: class MockTokenizerProcessorStep:
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
pass pass
@@ -192,7 +192,9 @@ def test_smolvla_processor_cuda():
def transform_features(self, features): def transform_features(self, features):
return features return features
with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor): with patch(
"lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessorStep
):
preprocessor, postprocessor = make_smolvla_pre_post_processors( preprocessor, postprocessor = make_smolvla_pre_post_processors(
config, config,
stats, stats,
@@ -225,7 +227,7 @@ def test_smolvla_processor_accelerate_scenario():
stats = create_default_stats() stats = create_default_stats()
# Mock the tokenizer processor to act as pass-through # Mock the tokenizer processor to act as pass-through
class MockTokenizerProcessor: class MockTokenizerProcessorStep:
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
pass pass
@@ -247,7 +249,9 @@ def test_smolvla_processor_accelerate_scenario():
def transform_features(self, features): def transform_features(self, features):
return features return features
with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor): with patch(
"lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessorStep
):
preprocessor, postprocessor = make_smolvla_pre_post_processors( preprocessor, postprocessor = make_smolvla_pre_post_processors(
config, config,
stats, stats,
@@ -281,7 +285,7 @@ def test_smolvla_processor_multi_gpu():
stats = create_default_stats() stats = create_default_stats()
# Mock the tokenizer processor to act as pass-through # Mock the tokenizer processor to act as pass-through
class MockTokenizerProcessor: class MockTokenizerProcessorStep:
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
pass pass
@@ -303,7 +307,9 @@ def test_smolvla_processor_multi_gpu():
def transform_features(self, features): def transform_features(self, features):
return features return features
with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor): with patch(
"lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessorStep
):
preprocessor, postprocessor = make_smolvla_pre_post_processors( preprocessor, postprocessor = make_smolvla_pre_post_processors(
config, config,
stats, stats,
@@ -334,7 +340,7 @@ def test_smolvla_processor_without_stats():
config = create_default_config() config = create_default_config()
# Mock the tokenizer processor # Mock the tokenizer processor
with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor"): with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep"):
preprocessor, postprocessor = make_smolvla_pre_post_processors( preprocessor, postprocessor = make_smolvla_pre_post_processors(
config, config,
dataset_stats=None, dataset_stats=None,
tests/processor/test_tdmpc_processor.py
+14 -14
View File
@@ -25,13 +25,13 @@ 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 (
AddBatchDimensionProcessorStep,
DataProcessorPipeline, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
RenameProcessor, RenameProcessorStep,
ToBatchProcessor,
TransitionKey, TransitionKey,
UnnormalizerProcessor, UnnormalizerProcessorStep,
) )
@@ -94,15 +94,15 @@ def test_make_tdmpc_processor_basic():
# Check steps in preprocessor # Check steps in preprocessor
assert len(preprocessor.steps) == 4 assert len(preprocessor.steps) == 4
assert isinstance(preprocessor.steps[0], RenameProcessor) assert isinstance(preprocessor.steps[0], RenameProcessorStep)
assert isinstance(preprocessor.steps[1], NormalizerProcessor) assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
assert isinstance(preprocessor.steps[2], ToBatchProcessor) assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep)
assert isinstance(preprocessor.steps[3], DeviceProcessor) assert isinstance(preprocessor.steps[3], DeviceProcessorStep)
# Check steps in postprocessor # Check steps in postprocessor
assert len(postprocessor.steps) == 2 assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor) assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_tdmpc_processor_normalization(): def test_tdmpc_processor_normalization():
@@ -330,10 +330,10 @@ def test_tdmpc_processor_mixed_precision():
postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x}, postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
) )
# Replace DeviceProcessor with one that uses float16 # Replace DeviceProcessorStep with one that uses float16
for i, step in enumerate(preprocessor.steps): for i, step in enumerate(preprocessor.steps):
if isinstance(step, DeviceProcessor): if isinstance(step, DeviceProcessorStep):
preprocessor.steps[i] = DeviceProcessor(device=config.device, float_dtype="float16") preprocessor.steps[i] = DeviceProcessorStep(device=config.device, float_dtype="float16")
# Create test data # Create test data
observation = { observation = {
tests/processor/test_tokenizer_processor.py
+50 -44
View File
@@ -1,5 +1,5 @@
""" """
Tests for the TokenizerProcessor class. Tests for the TokenizerProcessorStep class.
""" """
import tempfile import tempfile
@@ -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 DataProcessorPipeline, TokenizerProcessor, TransitionKey from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep, TransitionKey
from tests.utils import require_package from tests.utils import require_package
@@ -95,7 +95,7 @@ def test_basic_tokenization(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=10) processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=10)
transition = create_transition( transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])}, observation={"state": torch.tensor([1.0, 2.0])},
@@ -127,7 +127,7 @@ def test_basic_tokenization_with_tokenizer_object():
"""Test basic string tokenization functionality using tokenizer object directly.""" """Test basic string tokenization functionality using tokenizer object directly."""
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
transition = create_transition( transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])}, observation={"state": torch.tensor([1.0, 2.0])},
@@ -161,7 +161,7 @@ def test_list_of_strings_tokenization(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=8) processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=8)
transition = create_transition( transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])}, observation={"state": torch.tensor([1.0, 2.0])},
@@ -189,7 +189,7 @@ def test_custom_keys(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(tokenizer_name="test-tokenizer", task_key="instruction", max_length=5) processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", task_key="instruction", max_length=5)
transition = create_transition( transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])}, observation={"state": torch.tensor([1.0, 2.0])},
@@ -215,7 +215,7 @@ def test_none_complementary_data(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(tokenizer_name="test-tokenizer") processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")
transition = create_transition(complementary_data=None) transition = create_transition(complementary_data=None)
@@ -230,7 +230,7 @@ def test_missing_task_key(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(tokenizer_name="test-tokenizer") processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")
transition = create_transition(complementary_data={"other_field": "some value"}) transition = create_transition(complementary_data={"other_field": "some value"})
@@ -245,7 +245,7 @@ def test_none_task_value(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(tokenizer_name="test-tokenizer") processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")
transition = create_transition(complementary_data={"task": None}) transition = create_transition(complementary_data={"task": None})
@@ -260,7 +260,7 @@ def test_unsupported_task_type(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(tokenizer_name="test-tokenizer") processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")
# Test with integer task # Test with integer task
transition = create_transition(complementary_data={"task": 123}) transition = create_transition(complementary_data={"task": 123})
@@ -279,7 +279,7 @@ def test_unsupported_task_type(mock_auto_tokenizer):
def test_no_tokenizer_error(): def test_no_tokenizer_error():
"""Test that ValueError is raised when neither tokenizer nor tokenizer_name is provided.""" """Test that ValueError is raised when neither tokenizer nor tokenizer_name is provided."""
with pytest.raises(ValueError, match="Either 'tokenizer' or 'tokenizer_name' must be provided"): with pytest.raises(ValueError, match="Either 'tokenizer' or 'tokenizer_name' must be provided"):
TokenizerProcessor() TokenizerProcessorStep()
@require_package("transformers") @require_package("transformers")
@@ -290,7 +290,7 @@ def test_invalid_tokenizer_name_error():
mock_auto_tokenizer.from_pretrained.side_effect = Exception("Model not found") mock_auto_tokenizer.from_pretrained.side_effect = Exception("Model not found")
with pytest.raises(Exception, match="Model not found"): with pytest.raises(Exception, match="Model not found"):
TokenizerProcessor(tokenizer_name="invalid-tokenizer") TokenizerProcessorStep(tokenizer_name="invalid-tokenizer")
@require_package("transformers") @require_package("transformers")
@@ -300,7 +300,7 @@ def test_get_config_with_tokenizer_name(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor( processor = TokenizerProcessorStep(
tokenizer_name="test-tokenizer", tokenizer_name="test-tokenizer",
max_length=256, max_length=256,
task_key="instruction", task_key="instruction",
@@ -327,7 +327,7 @@ def test_get_config_with_tokenizer_object():
"""Test configuration serialization when using tokenizer object.""" """Test configuration serialization when using tokenizer object."""
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor( processor = TokenizerProcessorStep(
tokenizer=mock_tokenizer, tokenizer=mock_tokenizer,
max_length=256, max_length=256,
task_key="instruction", task_key="instruction",
@@ -357,7 +357,7 @@ def test_state_dict_methods(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(tokenizer_name="test-tokenizer") processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")
# Should return empty dict # Should return empty dict
state = processor.state_dict() state = processor.state_dict()
@@ -374,7 +374,7 @@ def test_reset_method(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(tokenizer_name="test-tokenizer") processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")
# Should not raise error # Should not raise error
processor.reset() processor.reset()
@@ -387,7 +387,7 @@ def test_integration_with_robot_processor(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
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 = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=6)
robot_processor = DataProcessorPipeline( robot_processor = DataProcessorPipeline(
[tokenizer_processor], to_transition=lambda x: x, to_output=lambda x: x [tokenizer_processor], to_transition=lambda x: x, to_output=lambda x: x
) )
@@ -424,7 +424,7 @@ def test_save_and_load_pretrained_with_tokenizer_name(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
original_processor = TokenizerProcessor( original_processor = TokenizerProcessorStep(
tokenizer_name="test-tokenizer", max_length=32, task_key="instruction" tokenizer_name="test-tokenizer", max_length=32, task_key="instruction"
) )
@@ -459,7 +459,9 @@ def test_save_and_load_pretrained_with_tokenizer_object():
"""Test saving and loading processor with tokenizer object using overrides.""" """Test saving and loading processor with tokenizer object using overrides."""
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
original_processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=32, task_key="instruction") original_processor = TokenizerProcessorStep(
tokenizer=mock_tokenizer, max_length=32, task_key="instruction"
)
robot_processor = DataProcessorPipeline( robot_processor = DataProcessorPipeline(
[original_processor], to_transition=lambda x: x, to_output=lambda x: x [original_processor], to_transition=lambda x: x, to_output=lambda x: x
@@ -500,14 +502,14 @@ def test_registry_functionality():
# Check that we can retrieve it # Check that we can retrieve it
retrieved_class = ProcessorStepRegistry.get("tokenizer_processor") retrieved_class = ProcessorStepRegistry.get("tokenizer_processor")
assert retrieved_class is TokenizerProcessor assert retrieved_class is TokenizerProcessorStep
@require_package("transformers") @require_package("transformers")
def test_features_basic(): def test_features_basic():
"""Test basic feature contract functionality.""" """Test basic feature contract functionality."""
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=128) processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=128)
input_features = { input_features = {
"observation.state": PolicyFeature(type=FeatureType.STATE, shape=(10,)), "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(10,)),
@@ -538,7 +540,7 @@ def test_features_basic():
def test_features_with_custom_max_length(): def test_features_with_custom_max_length():
"""Test feature contract with custom max_length.""" """Test feature contract with custom max_length."""
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=64) processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=64)
input_features = {} input_features = {}
output_features = processor.transform_features(input_features) output_features = processor.transform_features(input_features)
@@ -558,7 +560,7 @@ def test_features_with_custom_max_length():
def test_features_existing_features(): def test_features_existing_features():
"""Test feature contract when tokenized features already exist.""" """Test feature contract when tokenized features already exist."""
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=256) processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=256)
input_features = { input_features = {
f"{OBS_LANGUAGE}.tokens": PolicyFeature(type=FeatureType.LANGUAGE, shape=(100,)), f"{OBS_LANGUAGE}.tokens": PolicyFeature(type=FeatureType.LANGUAGE, shape=(100,)),
@@ -595,7 +597,7 @@ def test_tokenization_parameters(mock_auto_tokenizer):
tracking_tokenizer = TrackingMockTokenizer() tracking_tokenizer = TrackingMockTokenizer()
mock_auto_tokenizer.from_pretrained.return_value = tracking_tokenizer mock_auto_tokenizer.from_pretrained.return_value = tracking_tokenizer
processor = TokenizerProcessor( processor = TokenizerProcessorStep(
tokenizer_name="test-tokenizer", tokenizer_name="test-tokenizer",
max_length=16, max_length=16,
padding="longest", padding="longest",
@@ -627,7 +629,7 @@ def test_preserves_other_complementary_data(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(tokenizer_name="test-tokenizer") processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")
transition = create_transition( transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])}, observation={"state": torch.tensor([1.0, 2.0])},
@@ -662,7 +664,7 @@ def test_deterministic_tokenization(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=10) processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=10)
transition = create_transition( transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])}, observation={"state": torch.tensor([1.0, 2.0])},
@@ -690,7 +692,7 @@ def test_empty_string_task(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=8) processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=8)
transition = create_transition( transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])}, observation={"state": torch.tensor([1.0, 2.0])},
@@ -714,7 +716,7 @@ def test_very_long_task(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=5, truncation=True) processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=5, truncation=True)
long_task = " ".join(["word"] * 100) # Very long task long_task = " ".join(["word"] * 100) # Very long task
transition = create_transition( transition = create_transition(
@@ -764,7 +766,9 @@ def test_custom_padding_side(mock_auto_tokenizer):
mock_auto_tokenizer.from_pretrained.return_value = tracking_tokenizer mock_auto_tokenizer.from_pretrained.return_value = tracking_tokenizer
# Test left padding # Test left padding
processor_left = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=10, padding_side="left") processor_left = TokenizerProcessorStep(
tokenizer_name="test-tokenizer", max_length=10, padding_side="left"
)
transition = create_transition( transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])}, observation={"state": torch.tensor([1.0, 2.0])},
@@ -776,7 +780,9 @@ def test_custom_padding_side(mock_auto_tokenizer):
assert tracking_tokenizer.padding_side_calls[-1] == "left" assert tracking_tokenizer.padding_side_calls[-1] == "left"
# Test right padding (default) # Test right padding (default)
processor_right = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=10, padding_side="right") processor_right = TokenizerProcessorStep(
tokenizer_name="test-tokenizer", max_length=10, padding_side="right"
)
processor_right(transition) processor_right(transition)
@@ -787,7 +793,7 @@ def test_custom_padding_side(mock_auto_tokenizer):
def test_device_detection_cpu(): def test_device_detection_cpu():
"""Test that tokenized tensors stay on CPU when other tensors are on CPU.""" """Test that tokenized tensors stay on CPU when other tensors are on CPU."""
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
# Create transition with CPU tensors # Create transition with CPU tensors
observation = {"observation.state": torch.randn(10)} # CPU tensor observation = {"observation.state": torch.randn(10)} # CPU tensor
@@ -811,7 +817,7 @@ def test_device_detection_cpu():
def test_device_detection_cuda(): def test_device_detection_cuda():
"""Test that tokenized tensors are moved to CUDA when other tensors are on CUDA.""" """Test that tokenized tensors are moved to CUDA when other tensors are on CUDA."""
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
# Create transition with CUDA tensors # Create transition with CUDA tensors
observation = {"observation.state": torch.randn(10).cuda()} # CUDA tensor observation = {"observation.state": torch.randn(10).cuda()} # CUDA tensor
@@ -836,7 +842,7 @@ def test_device_detection_cuda():
def test_device_detection_multi_gpu(): def test_device_detection_multi_gpu():
"""Test that tokenized tensors match device in multi-GPU setup.""" """Test that tokenized tensors match device in multi-GPU setup."""
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
# Test with tensors on cuda:1 # Test with tensors on cuda:1
device = torch.device("cuda:1") device = torch.device("cuda:1")
@@ -860,7 +866,7 @@ def test_device_detection_multi_gpu():
def test_device_detection_no_tensors(): def test_device_detection_no_tensors():
"""Test that tokenized tensors stay on CPU when no other tensors exist.""" """Test that tokenized tensors stay on CPU when no other tensors exist."""
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
# Create transition with no tensors # Create transition with no tensors
transition = create_transition( transition = create_transition(
@@ -882,7 +888,7 @@ def test_device_detection_no_tensors():
def test_device_detection_mixed_devices(): def test_device_detection_mixed_devices():
"""Test device detection when tensors are on different devices (uses first found).""" """Test device detection when tensors are on different devices (uses first found)."""
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
if torch.cuda.is_available(): if torch.cuda.is_available():
# Create transition with mixed devices # Create transition with mixed devices
@@ -910,7 +916,7 @@ def test_device_detection_mixed_devices():
def test_device_detection_from_action(): def test_device_detection_from_action():
"""Test that device is detected from action tensor when no observation tensors exist.""" """Test that device is detected from action tensor when no observation tensors exist."""
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
# Create transition with action on CUDA but no observation tensors # Create transition with action on CUDA but no observation tensors
observation = {"metadata": {"key": "value"}} # No tensors in observation observation = {"metadata": {"key": "value"}} # No tensors in observation
@@ -933,7 +939,7 @@ def test_device_detection_from_action():
def test_device_detection_preserves_dtype(): def test_device_detection_preserves_dtype():
"""Test that device detection doesn't affect dtype of tokenized tensors.""" """Test that device detection doesn't affect dtype of tokenized tensors."""
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
# Create transition with float tensor (to test dtype isn't affected) # Create transition with float tensor (to test dtype isn't affected)
observation = {"observation.state": torch.randn(10, dtype=torch.float16)} observation = {"observation.state": torch.randn(10, dtype=torch.float16)}
@@ -953,15 +959,15 @@ def test_device_detection_preserves_dtype():
@require_package("transformers") @require_package("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer") @patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_integration_with_device_processor(mock_auto_tokenizer): def test_integration_with_device_processor(mock_auto_tokenizer):
"""Test that TokenizerProcessor works correctly with DeviceProcessor in pipeline.""" """Test that TokenizerProcessorStep works correctly with DeviceProcessorStep in pipeline."""
from lerobot.processor import DeviceProcessor from lerobot.processor import DeviceProcessorStep
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
# Create pipeline with TokenizerProcessor then DeviceProcessor # Create pipeline with TokenizerProcessorStep then DeviceProcessorStep
tokenizer_processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=6) tokenizer_processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=6)
device_processor = DeviceProcessor(device="cuda:0") device_processor = DeviceProcessorStep(device="cuda:0")
robot_processor = DataProcessorPipeline( 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
) )
@@ -975,7 +981,7 @@ def test_integration_with_device_processor(mock_auto_tokenizer):
result = robot_processor(transition) result = robot_processor(transition)
# All tensors should end up on CUDA (moved by DeviceProcessor) # All tensors should end up on CUDA (moved by DeviceProcessorStep)
assert result[TransitionKey.OBSERVATION]["observation.state"].device.type == "cuda" assert result[TransitionKey.OBSERVATION]["observation.state"].device.type == "cuda"
assert result[TransitionKey.ACTION].device.type == "cuda" assert result[TransitionKey.ACTION].device.type == "cuda"
@@ -991,7 +997,7 @@ def test_integration_with_device_processor(mock_auto_tokenizer):
def test_simulated_accelerate_scenario(): def test_simulated_accelerate_scenario():
"""Test scenario simulating Accelerate with data already on GPU.""" """Test scenario simulating Accelerate with data already on GPU."""
mock_tokenizer = MockTokenizer(vocab_size=100) mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
# Simulate Accelerate scenario: batch already on GPU # Simulate Accelerate scenario: batch already on GPU
device = torch.device("cuda:0") device = torch.device("cuda:0")
tests/processor/test_vqbet_processor.py
+14 -14
View File
@@ -25,13 +25,13 @@ 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 (
AddBatchDimensionProcessorStep,
DataProcessorPipeline, DataProcessorPipeline,
DeviceProcessor, DeviceProcessorStep,
NormalizerProcessor, NormalizerProcessorStep,
RenameProcessor, RenameProcessorStep,
ToBatchProcessor,
TransitionKey, TransitionKey,
UnnormalizerProcessor, UnnormalizerProcessorStep,
) )
@@ -94,15 +94,15 @@ def test_make_vqbet_processor_basic():
# Check steps in preprocessor # Check steps in preprocessor
assert len(preprocessor.steps) == 4 assert len(preprocessor.steps) == 4
assert isinstance(preprocessor.steps[0], RenameProcessor) assert isinstance(preprocessor.steps[0], RenameProcessorStep)
assert isinstance(preprocessor.steps[1], NormalizerProcessor) assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
assert isinstance(preprocessor.steps[2], ToBatchProcessor) assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep)
assert isinstance(preprocessor.steps[3], DeviceProcessor) assert isinstance(preprocessor.steps[3], DeviceProcessorStep)
# Check steps in postprocessor # Check steps in postprocessor
assert len(postprocessor.steps) == 2 assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor) assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_vqbet_processor_with_images(): def test_vqbet_processor_with_images():
@@ -323,10 +323,10 @@ def test_vqbet_processor_mixed_precision():
postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x}, postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
) )
# Replace DeviceProcessor with one that uses float16 # Replace DeviceProcessorStep with one that uses float16
for i, step in enumerate(preprocessor.steps): for i, step in enumerate(preprocessor.steps):
if isinstance(step, DeviceProcessor): if isinstance(step, DeviceProcessorStep):
preprocessor.steps[i] = DeviceProcessor(device=config.device, float_dtype="float16") preprocessor.steps[i] = DeviceProcessorStep(device=config.device, float_dtype="float16")
# Create test data # Create test data
observation = { observation = {