admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-05-14 16:19:45 +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:
1. **VanillaObservationProcessor**: Converts raw robot observations into standardized format
2. **JointVelocityProcessor** (optional): Adds joint velocity information to observations
3. **MotorCurrentProcessor** (optional): Adds motor current readings to observations
1. **VanillaObservationProcessorStep**: Converts raw robot observations into standardized format
2. **JointVelocityProcessorStep** (optional): Adds joint velocity information to observations
3. **MotorCurrentProcessorStep** (optional): Adds motor current readings to observations
4. **ForwardKinematicsJointsToEE** (optional): Computes end-effector pose from joint positions
5. **ImageCropResizeProcessor** (optional): Crops and resizes camera images
6. **TimeLimitProcessor** (optional): Enforces episode time limits
7. **GripperPenaltyProcessor** (optional): Applies penalties for inappropriate gripper usage
8. **RewardClassifierProcessor** (optional): Automated reward detection using vision models
9. **ToBatchProcessor**: Converts data to batch format for neural network processing
10. **DeviceProcessor**: Moves data to the specified compute device (CPU/GPU)
5. **ImageCropResizeProcessorStep** (optional): Crops and resizes camera images
6. **TimeLimitProcessorStep** (optional): Enforces episode time limits
7. **GripperPenaltyProcessorStep** (optional): Applies penalties for inappropriate gripper usage
8. **RewardClassifierProcessorStep** (optional): Automated reward detection using vision models
9. **AddBatchDimensionProcessorStep**: Converts data to batch format for neural network processing
10. **DeviceProcessorStep**: Moves data to the specified compute device (CPU/GPU)
#### Action Processor Pipeline
The action processor (`action_processor`) handles outgoing actions and human interventions:
1. **AddTeleopActionAsComplimentaryData**: Captures teleoperator actions for logging
2. **AddTeleopEventsAsInfo**: Records intervention events and episode control signals
1. **AddTeleopActionAsComplimentaryDataStep**: Captures teleoperator actions for logging
2. **AddTeleopEventsAsInfoStep**: Records intervention events and episode control signals
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):
- **MapDeltaActionToRobotAction**: Converts delta actions to robot action format
- **MapDeltaActionToRobotActionStep**: Converts delta actions to robot action format
- **EEReferenceAndDelta**: Computes end-effector reference and delta movements
- **EEBoundsAndSafety**: Enforces workspace safety bounds
- **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.policies.act.configuration_act import ACTConfig
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DataProcessorPipeline,
DeviceProcessor,
NormalizerProcessor,
DeviceProcessorStep,
NormalizerProcessorStep,
ProcessorKwargs,
RenameProcessor,
ToBatchProcessor,
UnnormalizerProcessor,
RenameProcessorStep,
UnnormalizerProcessorStep,
)
@@ -40,18 +40,18 @@ def make_act_pre_post_processors(
postprocessor_kwargs = {}
input_steps = [
RenameProcessor(rename_map={}),
NormalizerProcessor(
RenameProcessorStep(rename_map={}),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
ToBatchProcessor(),
DeviceProcessor(device=config.device),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
]
output_steps = [
DeviceProcessor(device="cpu"),
UnnormalizerProcessor(
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
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.policies.diffusion.configuration_diffusion import DiffusionConfig
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DataProcessorPipeline,
DeviceProcessor,
NormalizerProcessor,
DeviceProcessorStep,
NormalizerProcessorStep,
ProcessorKwargs,
RenameProcessor,
ToBatchProcessor,
UnnormalizerProcessor,
RenameProcessorStep,
UnnormalizerProcessorStep,
)
@@ -41,18 +41,18 @@ def make_diffusion_pre_post_processors(
postprocessor_kwargs = {}
input_steps = [
RenameProcessor(rename_map={}),
NormalizerProcessor(
RenameProcessorStep(rename_map={}),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
ToBatchProcessor(),
DeviceProcessor(device=config.device),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
]
output_steps = [
DeviceProcessor(device="cpu"),
UnnormalizerProcessor(
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
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.policies.pi0.configuration_pi0 import PI0Config
from lerobot.processor import (
AddBatchDimensionProcessorStep,
ComplementaryDataProcessorStep,
DataProcessorPipeline,
DeviceProcessor,
NormalizerProcessor,
DeviceProcessorStep,
NormalizerProcessorStep,
ProcessorKwargs,
ProcessorStep,
ProcessorStepRegistry,
RenameProcessor,
ToBatchProcessor,
TokenizerProcessor,
UnnormalizerProcessor,
RenameProcessorStep,
TokenizerProcessorStep,
UnnormalizerProcessorStep,
)
@@ -80,26 +80,26 @@ def make_pi0_pre_post_processors(
# Add remaining processors
input_steps: list[ProcessorStep] = [
RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one
NormalizerProcessor(
RenameProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
ToBatchProcessor(),
AddBatchDimensionProcessorStep(),
Pi0NewLineProcessor(), # Add newlines before tokenization for PaliGemma
TokenizerProcessor(
TokenizerProcessorStep(
tokenizer_name="google/paligemma-3b-pt-224",
max_length=config.tokenizer_max_length,
padding_side="right",
padding="max_length",
),
DeviceProcessor(device=config.device),
DeviceProcessorStep(device=config.device),
]
output_steps: list[ProcessorStep] = [
DeviceProcessor(device="cpu"),
UnnormalizerProcessor(
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
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.policies.pi0.configuration_pi0 import PI0Config
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DataProcessorPipeline,
DeviceProcessor,
NormalizerProcessor,
DeviceProcessorStep,
NormalizerProcessorStep,
ProcessorKwargs,
RenameProcessor,
ToBatchProcessor,
UnnormalizerProcessor,
RenameProcessorStep,
UnnormalizerProcessorStep,
)
@@ -41,18 +41,18 @@ def make_pi0fast_pre_post_processors(
postprocessor_kwargs = {}
input_steps = [
RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one
NormalizerProcessor(
RenameProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
ToBatchProcessor(),
DeviceProcessor(device=config.device),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
]
output_steps = [
DeviceProcessor(device="cpu"),
UnnormalizerProcessor(
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
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.policies.sac.configuration_sac import SACConfig
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DataProcessorPipeline,
DeviceProcessor,
NormalizerProcessor,
DeviceProcessorStep,
NormalizerProcessorStep,
ProcessorKwargs,
RenameProcessor,
ToBatchProcessor,
UnnormalizerProcessor,
RenameProcessorStep,
UnnormalizerProcessorStep,
)
@@ -42,18 +42,18 @@ def make_sac_pre_post_processors(
postprocessor_kwargs = {}
input_steps = [
RenameProcessor(rename_map={}),
NormalizerProcessor(
RenameProcessorStep(rename_map={}),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
ToBatchProcessor(),
DeviceProcessor(device=config.device),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
]
output_steps = [
DeviceProcessor(device="cpu"),
UnnormalizerProcessor(
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
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.processor import (
DataProcessorPipeline,
DeviceProcessor,
DeviceProcessorStep,
IdentityProcessorStep,
NormalizerProcessor,
NormalizerProcessorStep,
ProcessorKwargs,
)
@@ -37,15 +37,15 @@ def make_classifier_processor(
postprocessor_kwargs = {}
input_steps = [
NormalizerProcessor(
NormalizerProcessorStep(
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
),
DeviceProcessor(device=config.device),
DeviceProcessorStep(device=config.device),
]
output_steps = [DeviceProcessor(device="cpu"), IdentityProcessorStep()]
output_steps = [DeviceProcessorStep(device="cpu"), IdentityProcessorStep()]
return (
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.policies.smolvla.configuration_smolvla import SmolVLAConfig
from lerobot.processor import (
AddBatchDimensionProcessorStep,
ComplementaryDataProcessorStep,
DataProcessorPipeline,
DeviceProcessor,
NormalizerProcessor,
DeviceProcessorStep,
NormalizerProcessorStep,
ProcessorKwargs,
ProcessorStepRegistry,
RenameProcessor,
ToBatchProcessor,
TokenizerProcessor,
UnnormalizerProcessor,
RenameProcessorStep,
TokenizerProcessorStep,
UnnormalizerProcessorStep,
)
@@ -45,25 +45,25 @@ def make_smolvla_pre_post_processors(
postprocessor_kwargs = {}
input_steps = [
RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one
NormalizerProcessor(
RenameProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
ToBatchProcessor(),
AddBatchDimensionProcessorStep(),
SmolVLANewLineProcessor(),
TokenizerProcessor(
TokenizerProcessorStep(
tokenizer_name=config.vlm_model_name,
padding=config.pad_language_to,
padding_side="right",
max_length=config.tokenizer_max_length,
),
DeviceProcessor(device=config.device),
DeviceProcessorStep(device=config.device),
]
output_steps = [
DeviceProcessor(device="cpu"),
UnnormalizerProcessor(
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
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.policies.tdmpc.configuration_tdmpc import TDMPCConfig
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DataProcessorPipeline,
DeviceProcessor,
NormalizerProcessor,
DeviceProcessorStep,
NormalizerProcessorStep,
ProcessorKwargs,
RenameProcessor,
ToBatchProcessor,
UnnormalizerProcessor,
RenameProcessorStep,
UnnormalizerProcessorStep,
)
@@ -41,18 +41,18 @@ def make_tdmpc_pre_post_processors(
postprocessor_kwargs = {}
input_steps = [
RenameProcessor(rename_map={}),
NormalizerProcessor(
RenameProcessorStep(rename_map={}),
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
ToBatchProcessor(),
DeviceProcessor(device=config.device),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
]
output_steps = [
DeviceProcessor(device="cpu"),
UnnormalizerProcessor(
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
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.policies.vqbet.configuration_vqbet import VQBeTConfig
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DataProcessorPipeline,
DeviceProcessor,
NormalizerProcessor,
DeviceProcessorStep,
NormalizerProcessorStep,
ProcessorKwargs,
RenameProcessor,
ToBatchProcessor,
UnnormalizerProcessor,
RenameProcessorStep,
UnnormalizerProcessorStep,
)
@@ -42,18 +42,18 @@ def make_vqbet_pre_post_processors(
postprocessor_kwargs = {}
input_steps = [
RenameProcessor(rename_map={}), # Let the possibility to the user to rename the keys
NormalizerProcessor(
RenameProcessorStep(rename_map={}), # Let the possibility to the user to rename the keys
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
ToBatchProcessor(),
DeviceProcessor(device=config.device),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=config.device),
]
output_steps = [
DeviceProcessor(device="cpu"),
UnnormalizerProcessor(
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
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
# limitations under the License.
from .batch_processor import ToBatchProcessor
from .batch_processor import AddBatchDimensionProcessorStep
from .converters import (
batch_to_transition,
create_transition,
@@ -23,21 +23,21 @@ from .converters import (
transition_to_dataset_frame,
)
from .core import EnvTransition, TransitionKey
from .delta_action_processor import MapDeltaActionToRobotAction, MapTensorToDeltaActionDict
from .device_processor import DeviceProcessor
from .gym_action_processor import Numpy2TorchActionProcessor, Torch2NumpyActionProcessor
from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep
from .device_processor import DeviceProcessorStep
from .gym_action_processor import Numpy2TorchActionProcessorStep, Torch2NumpyActionProcessorStep
from .hil_processor import (
AddTeleopActionAsComplimentaryData,
AddTeleopEventsAsInfo,
GripperPenaltyProcessor,
ImageCropResizeProcessor,
InterventionActionProcessor,
RewardClassifierProcessor,
TimeLimitProcessor,
AddTeleopActionAsComplimentaryDataStep,
AddTeleopEventsAsInfoStep,
GripperPenaltyProcessorStep,
ImageCropResizeProcessorStep,
InterventionActionProcessorStep,
RewardClassifierProcessorStep,
TimeLimitProcessorStep,
)
from .joint_observations_processor import JointVelocityProcessor, MotorCurrentProcessor
from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor, hotswap_stats
from .observation_processor import VanillaObservationProcessor
from .joint_observations_processor import JointVelocityProcessorStep, MotorCurrentProcessorStep
from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep, hotswap_stats
from .observation_processor import VanillaObservationProcessorStep
from .pipeline import (
ActionProcessorStep,
ComplementaryDataProcessorStep,
@@ -52,48 +52,48 @@ from .pipeline import (
RewardProcessorStep,
TruncatedProcessorStep,
)
from .rename_processor import RenameProcessor
from .tokenizer_processor import TokenizerProcessor
from .rename_processor import RenameProcessorStep
from .tokenizer_processor import TokenizerProcessorStep
__all__ = [
"ActionProcessorStep",
"AddTeleopActionAsComplimentaryData",
"AddTeleopEventsAsInfo",
"AddTeleopActionAsComplimentaryDataStep",
"AddTeleopEventsAsInfoStep",
"ComplementaryDataProcessorStep",
"batch_to_transition",
"create_transition",
"DeviceProcessor",
"DeviceProcessorStep",
"DoneProcessorStep",
"EnvTransition",
"GripperPenaltyProcessor",
"GripperPenaltyProcessorStep",
"hotswap_stats",
"IdentityProcessorStep",
"ImageCropResizeProcessor",
"ImageCropResizeProcessorStep",
"InfoProcessorStep",
"InterventionActionProcessor",
"JointVelocityProcessor",
"MapDeltaActionToRobotAction",
"MapTensorToDeltaActionDict",
"InterventionActionProcessorStep",
"JointVelocityProcessorStep",
"MapDeltaActionToRobotActionStep",
"MapTensorToDeltaActionDictStep",
"merge_transitions",
"MotorCurrentProcessor",
"NormalizerProcessor",
"Numpy2TorchActionProcessor",
"MotorCurrentProcessorStep",
"NormalizerProcessorStep",
"Numpy2TorchActionProcessorStep",
"ObservationProcessorStep",
"ProcessorKwargs",
"ProcessorStep",
"ProcessorStepRegistry",
"RenameProcessor",
"RewardClassifierProcessor",
"RenameProcessorStep",
"RewardClassifierProcessorStep",
"RewardProcessorStep",
"DataProcessorPipeline",
"TimeLimitProcessor",
"ToBatchProcessor",
"TokenizerProcessor",
"Torch2NumpyActionProcessor",
"TimeLimitProcessorStep",
"AddBatchDimensionProcessorStep",
"TokenizerProcessorStep",
"Torch2NumpyActionProcessorStep",
"transition_to_batch",
"transition_to_dataset_frame",
"TransitionKey",
"TruncatedProcessorStep",
"UnnormalizerProcessor",
"VanillaObservationProcessor",
"UnnormalizerProcessorStep",
"VanillaObservationProcessorStep",
]
src/lerobot/processor/batch_processor.py
+11 -9
View File
@@ -30,7 +30,7 @@ from .pipeline import (
@dataclass
@ProcessorStepRegistry.register(name="to_batch_processor_action")
class ToBatchProcessorAction(ActionProcessorStep):
class AddBatchDimensionActionStep(ActionProcessorStep):
"""Process action component in-place, adding batch dimension if needed."""
def action(self, action):
@@ -45,7 +45,7 @@ class ToBatchProcessorAction(ActionProcessorStep):
@dataclass
@ProcessorStepRegistry.register(name="to_batch_processor_observation")
class ToBatchProcessorObservation(ObservationProcessorStep):
class AddBatchDimensionObservationStep(ObservationProcessorStep):
"""Process observation component in-place, adding batch dimensions where needed."""
def observation(self, observation):
@@ -74,7 +74,7 @@ class ToBatchProcessorObservation(ObservationProcessorStep):
@dataclass
@ProcessorStepRegistry.register(name="to_batch_processor_complementary_data")
class ToBatchProcessorComplementaryData(ComplementaryDataProcessorStep):
class AddBatchDimensionComplementaryDataStep(ComplementaryDataProcessorStep):
"""Process complementary data in-place, handling task field batching."""
def complementary_data(self, complementary_data):
@@ -103,7 +103,7 @@ class ToBatchProcessorComplementaryData(ComplementaryDataProcessorStep):
@dataclass
@ProcessorStepRegistry.register(name="to_batch_processor")
class ToBatchProcessor(ProcessorStep):
class AddBatchDimensionProcessorStep(ProcessorStep):
"""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:
@@ -138,12 +138,14 @@ class ToBatchProcessor(ProcessorStep):
```
"""
to_batch_action_processor: ToBatchProcessorAction = field(default_factory=ToBatchProcessorAction)
to_batch_observation_processor: ToBatchProcessorObservation = field(
default_factory=ToBatchProcessorObservation
to_batch_action_processor: AddBatchDimensionActionStep = field(
default_factory=AddBatchDimensionActionStep
)
to_batch_complementary_data_processor: ToBatchProcessorComplementaryData = field(
default_factory=ToBatchProcessorComplementaryData
to_batch_observation_processor: AddBatchDimensionObservationStep = field(
default_factory=AddBatchDimensionObservationStep
)
to_batch_complementary_data_processor: AddBatchDimensionComplementaryDataStep = field(
default_factory=AddBatchDimensionComplementaryDataStep
)
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")
@dataclass
class MapTensorToDeltaActionDict(ActionProcessorStep):
class MapTensorToDeltaActionDictStep(ActionProcessorStep):
"""
Map a tensor to a delta action dictionary.
"""
@@ -58,7 +58,7 @@ class MapTensorToDeltaActionDict(ActionProcessorStep):
@ProcessorStepRegistry.register("map_delta_action_to_robot_action")
@dataclass
class MapDeltaActionToRobotAction(ActionProcessorStep):
class MapDeltaActionToRobotActionStep(ActionProcessorStep):
"""
Map delta actions from teleoperators (gamepad, keyboard) to robot target actions
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")
@dataclass
class DeviceProcessor(ProcessorStep):
class DeviceProcessorStep(ProcessorStep):
"""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
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")
@dataclass
class Torch2NumpyActionProcessor(ActionProcessorStep):
class Torch2NumpyActionProcessorStep(ActionProcessorStep):
"""Convert PyTorch tensor actions to NumPy arrays."""
squeeze_batch_dim: bool = True
@@ -56,7 +56,7 @@ class Torch2NumpyActionProcessor(ActionProcessorStep):
@ProcessorStepRegistry.register("numpy2torch_action_processor")
@dataclass
class Numpy2TorchActionProcessor(ActionProcessorStep):
class Numpy2TorchActionProcessorStep(ActionProcessorStep):
"""Convert NumPy array action to PyTorch 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")
@dataclass
class AddTeleopActionAsComplimentaryData(ComplementaryDataProcessorStep):
class AddTeleopActionAsComplimentaryDataStep(ComplementaryDataProcessorStep):
"""Add teleoperator action to transition complementary data."""
teleop_device: Teleoperator
@@ -45,7 +45,7 @@ class AddTeleopActionAsComplimentaryData(ComplementaryDataProcessorStep):
@ProcessorStepRegistry.register("add_teleop_action_as_info")
@dataclass
class AddTeleopEventsAsInfo(InfoProcessorStep):
class AddTeleopEventsAsInfoStep(InfoProcessorStep):
"""Add teleoperator control events to transition info."""
teleop_device: Teleoperator
@@ -62,7 +62,7 @@ class AddTeleopEventsAsInfo(InfoProcessorStep):
@ProcessorStepRegistry.register("image_crop_resize_processor")
@dataclass
class ImageCropResizeProcessor(ObservationProcessorStep):
class ImageCropResizeProcessorStep(ObservationProcessorStep):
"""Crop and resize image observations."""
crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
@@ -112,7 +112,7 @@ class ImageCropResizeProcessor(ObservationProcessorStep):
@dataclass
@ProcessorStepRegistry.register("time_limit_processor")
class TimeLimitProcessor(TruncatedProcessorStep):
class TimeLimitProcessorStep(TruncatedProcessorStep):
"""Track episode steps and enforce time limits."""
max_episode_steps: int
@@ -139,7 +139,7 @@ class TimeLimitProcessor(TruncatedProcessorStep):
@dataclass
@ProcessorStepRegistry.register("gripper_penalty_processor")
class GripperPenaltyProcessor(ComplementaryDataProcessorStep):
class GripperPenaltyProcessorStep(ComplementaryDataProcessorStep):
"""Apply penalty for inappropriate gripper usage."""
penalty: float = -0.01
@@ -188,7 +188,7 @@ class GripperPenaltyProcessor(ComplementaryDataProcessorStep):
@dataclass
@ProcessorStepRegistry.register("intervention_action_processor")
class InterventionActionProcessor(ProcessorStep):
class InterventionActionProcessorStep(ProcessorStep):
"""Handle human intervention actions and episode termination."""
use_gripper: bool = False
@@ -261,7 +261,7 @@ class InterventionActionProcessor(ProcessorStep):
@dataclass
@ProcessorStepRegistry.register("reward_classifier_processor")
class RewardClassifierProcessor(ProcessorStep):
class RewardClassifierProcessorStep(ProcessorStep):
"""Apply reward classification to image observations."""
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
@ProcessorStepRegistry.register("joint_velocity_processor")
class JointVelocityProcessor(ObservationProcessorStep):
class JointVelocityProcessorStep(ObservationProcessorStep):
"""Add joint velocity information to observations."""
dt: float = 0.1
@@ -67,7 +67,7 @@ class JointVelocityProcessor(ObservationProcessorStep):
@dataclass
@ProcessorStepRegistry.register("current_processor")
class MotorCurrentProcessor(ObservationProcessorStep):
class MotorCurrentProcessorStep(ObservationProcessorStep):
"""Add motor current information to observations."""
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 .batch_processor import ToBatchProcessor
from .device_processor import DeviceProcessor
from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor
from .batch_processor import AddBatchDimensionProcessorStep
from .device_processor import DeviceProcessorStep
from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
from .pipeline import DataProcessorPipeline
from .rename_processor import RenameProcessor
from .rename_processor import RenameProcessorStep
# Policy type to class mapping
POLICY_CLASSES = {
@@ -404,8 +404,8 @@ def main():
# Now create preprocessor and postprocessor with cleaned_config available
print("Creating preprocessor and postprocessor...")
# The pattern from existing processor factories:
# - Preprocessor has two NormalizerProcessors: one for input_features, one for output_features
# - Postprocessor has one UnnormalizerProcessor for output_features only
# - Preprocessor has two NormalizerProcessorSteps: one for input_features, one for output_features
# - Postprocessor has one UnnormalizerProcessorStep for output_features only
# Get features from cleaned_config (now they're PolicyFeature objects)
input_features = cleaned_config.get("input_features", {})
@@ -413,21 +413,21 @@ def main():
# Create preprocessor with two normalizers (following the pattern from processor factories)
preprocessor_steps = [
RenameProcessor(rename_map={}),
NormalizerProcessor(
RenameProcessorStep(rename_map={}),
NormalizerProcessorStep(
features={**input_features, **output_features},
norm_map=norm_map,
stats=stats,
),
ToBatchProcessor(),
DeviceProcessor(device=policy_config.device),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=policy_config.device),
]
preprocessor = DataProcessorPipeline(steps=preprocessor_steps, name="robot_preprocessor")
# Create postprocessor with unnormalizer for outputs only
postprocessor_steps = [
DeviceProcessor(device="cpu"),
UnnormalizerProcessor(features=output_features, norm_map=norm_map, stats=stats),
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(features=output_features, norm_map=norm_map, stats=stats),
]
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
@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.
@@ -180,7 +180,7 @@ class NormalizerProcessor(_NormalizationMixin, ProcessorStep):
normalize_observation_keys: set[str] | None = None,
eps: float = 1e-8,
device: torch.device | str | None = None,
) -> NormalizerProcessor:
) -> NormalizerProcessorStep:
return cls(
features=features,
norm_map=norm_map,
@@ -213,7 +213,7 @@ class NormalizerProcessor(_NormalizationMixin, ProcessorStep):
@dataclass
@ProcessorStepRegistry.register(name="unnormalizer_processor")
class UnnormalizerProcessor(_NormalizationMixin, ProcessorStep):
class UnnormalizerProcessorStep(_NormalizationMixin, ProcessorStep):
"""
A processor that applies unnormalization (the inverse of normalization) to
observations and actions in a transition.
@@ -230,7 +230,7 @@ class UnnormalizerProcessor(_NormalizationMixin, ProcessorStep):
norm_map: dict[FeatureType, NormalizationMode],
*,
device: torch.device | str | None = None,
) -> UnnormalizerProcessor:
) -> UnnormalizerProcessorStep:
return cls(features=features, norm_map=norm_map, stats=dataset.meta.stats, device=device)
def __call__(self, transition: EnvTransition) -> EnvTransition:
@@ -259,7 +259,7 @@ def hotswap_stats(
Replaces normalization statistics in a RobotProcessor pipeline.
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
different data distributions.
"""
src/lerobot/processor/observation_processor.py
+1 -1
View File
@@ -28,7 +28,7 @@ from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
@dataclass
@ProcessorStepRegistry.register(name="observation_processor")
class VanillaObservationProcessor(ObservationProcessorStep):
class VanillaObservationProcessorStep(ObservationProcessorStep):
"""
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
@ProcessorStepRegistry.register(name="rename_processor")
class RenameProcessor(ObservationProcessorStep):
class RenameProcessorStep(ObservationProcessorStep):
"""Rename processor that renames keys in the observation."""
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
@ProcessorStepRegistry.register(name="tokenizer_processor")
class TokenizerProcessor(ObservationProcessorStep):
class TokenizerProcessorStep(ObservationProcessorStep):
"""Tokenizes text tasks in complementary data using a huggingface tokenizer.
This processor handles tokenization of task strings found in the complementary_data
@@ -48,7 +48,7 @@ class TokenizerProcessor(ObservationProcessorStep):
Examples:
Using tokenizer name (auto-loaded):
```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:
@@ -56,7 +56,7 @@ class TokenizerProcessor(ObservationProcessorStep):
from transformers import AutoTokenizer
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:
raise ImportError(
"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:
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.model.kinematics import RobotKinematics
from lerobot.processor import (
AddTeleopActionAsComplimentaryData,
AddTeleopEventsAsInfo,
AddBatchDimensionProcessorStep,
AddTeleopActionAsComplimentaryDataStep,
AddTeleopEventsAsInfoStep,
DataProcessorPipeline,
DeviceProcessor,
DeviceProcessorStep,
EnvTransition,
GripperPenaltyProcessor,
ImageCropResizeProcessor,
InterventionActionProcessor,
JointVelocityProcessor,
MapDeltaActionToRobotAction,
MapTensorToDeltaActionDict,
MotorCurrentProcessor,
Numpy2TorchActionProcessor,
RewardClassifierProcessor,
TimeLimitProcessor,
ToBatchProcessor,
Torch2NumpyActionProcessor,
GripperPenaltyProcessorStep,
ImageCropResizeProcessorStep,
InterventionActionProcessorStep,
JointVelocityProcessorStep,
MapDeltaActionToRobotActionStep,
MapTensorToDeltaActionDictStep,
MotorCurrentProcessorStep,
Numpy2TorchActionProcessorStep,
RewardClassifierProcessorStep,
TimeLimitProcessorStep,
Torch2NumpyActionProcessorStep,
TransitionKey,
VanillaObservationProcessor,
VanillaObservationProcessorStep,
create_transition,
)
from lerobot.robots import ( # noqa: F401
@@ -362,16 +362,16 @@ def make_processors(
if cfg.name == "gym_hil":
action_pipeline_steps = [
InterventionActionProcessor(terminate_on_success=terminate_on_success),
Torch2NumpyActionProcessor(),
InterventionActionProcessorStep(terminate_on_success=terminate_on_success),
Torch2NumpyActionProcessorStep(),
]
# Minimal processor pipeline for GymHIL simulation
env_pipeline_steps = [
Numpy2TorchActionProcessor(),
VanillaObservationProcessor(),
ToBatchProcessor(),
DeviceProcessor(device=device),
Numpy2TorchActionProcessorStep(),
VanillaObservationProcessorStep(),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device=device),
]
return DataProcessorPipeline(steps=env_pipeline_steps), DataProcessorPipeline(
@@ -391,13 +391,13 @@ def make_processors(
joint_names=motor_names,
)
env_pipeline_steps = [VanillaObservationProcessor()]
env_pipeline_steps = [VanillaObservationProcessorStep()]
if cfg.processor.observation is not None:
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:
env_pipeline_steps.append(MotorCurrentProcessor(robot=env.robot))
env_pipeline_steps.append(MotorCurrentProcessorStep(robot=env.robot))
if kinematics_solver is not None:
env_pipeline_steps.append(
@@ -409,7 +409,7 @@ def make_processors(
if cfg.processor.image_preprocessing is not None:
env_pipeline_steps.append(
ImageCropResizeProcessor(
ImageCropResizeProcessorStep(
crop_params_dict=cfg.processor.image_preprocessing.crop_params_dict,
resize_size=cfg.processor.image_preprocessing.resize_size,
)
@@ -418,13 +418,13 @@ def make_processors(
# Add time limit processor if reset config exists
if cfg.processor.reset is not None:
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
if cfg.processor.gripper is not None and cfg.processor.gripper.use_gripper:
env_pipeline_steps.append(
GripperPenaltyProcessor(
GripperPenaltyProcessorStep(
penalty=cfg.processor.gripper.gripper_penalty,
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
):
env_pipeline_steps.append(
RewardClassifierProcessor(
RewardClassifierProcessorStep(
pretrained_path=cfg.processor.reward_classifier.pretrained_path,
device=device,
success_threshold=cfg.processor.reward_classifier.success_threshold,
@@ -444,14 +444,14 @@ def make_processors(
)
)
env_pipeline_steps.append(ToBatchProcessor())
env_pipeline_steps.append(DeviceProcessor(device=device))
env_pipeline_steps.append(AddBatchDimensionProcessorStep())
env_pipeline_steps.append(DeviceProcessorStep(device=device))
action_pipeline_steps = [
AddTeleopActionAsComplimentaryData(teleop_device=teleop_device),
AddTeleopEventsAsInfo(teleop_device=teleop_device),
AddTeleopActionAsComplimentaryDataStep(teleop_device=teleop_device),
AddTeleopEventsAsInfoStep(teleop_device=teleop_device),
AddRobotObservationAsComplimentaryData(robot=env.robot),
InterventionActionProcessor(
InterventionActionProcessorStep(
use_gripper=cfg.processor.gripper.use_gripper if cfg.processor.gripper is not None else False,
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:
# Add EE bounds and safety processor
inverse_kinematics_steps = [
MapTensorToDeltaActionDict(
MapTensorToDeltaActionDictStep(
use_gripper=cfg.processor.gripper.use_gripper if cfg.processor.gripper is not None else False
),
MapDeltaActionToRobotAction(),
MapDeltaActionToRobotActionStep(),
EEReferenceAndDelta(
kinematics=kinematics_solver,
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.processor_act import make_act_pre_post_processors
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DataProcessorPipeline,
DeviceProcessor,
NormalizerProcessor,
RenameProcessor,
ToBatchProcessor,
DeviceProcessorStep,
NormalizerProcessorStep,
RenameProcessorStep,
TransitionKey,
UnnormalizerProcessor,
UnnormalizerProcessorStep,
)
@@ -86,15 +86,15 @@ def test_make_act_processor_basic():
# Check steps in preprocessor
assert len(preprocessor.steps) == 4
assert isinstance(preprocessor.steps[0], RenameProcessor)
assert isinstance(preprocessor.steps[1], NormalizerProcessor)
assert isinstance(preprocessor.steps[2], ToBatchProcessor)
assert isinstance(preprocessor.steps[3], DeviceProcessor)
assert isinstance(preprocessor.steps[0], RenameProcessorStep)
assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep)
assert isinstance(preprocessor.steps[3], DeviceProcessorStep)
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessor)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
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},
)
# Replace DeviceProcessor with one that uses float16
# Replace DeviceProcessorStep with one that uses float16
modified_steps = []
for step in preprocessor.steps:
if isinstance(step, DeviceProcessor):
modified_steps.append(DeviceProcessor(device=config.device, float_dtype="float16"))
if isinstance(step, DeviceProcessorStep):
modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16"))
else:
modified_steps.append(step)
preprocessor.steps = modified_steps
tests/processor/test_batch_processor.py
+61 -56
View File
@@ -22,7 +22,12 @@ import pytest
import torch
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(
@@ -42,7 +47,7 @@ def create_transition(
def test_state_1d_to_2d():
"""Test that 1D state tensors get unsqueezed to 2D."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Test observation.state
state_1d = torch.randn(7)
@@ -58,7 +63,7 @@ def test_state_1d_to_2d():
def test_env_state_1d_to_2d():
"""Test that 1D environment state tensors get unsqueezed to 2D."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Test observation.environment_state
env_state_1d = torch.randn(10)
@@ -74,7 +79,7 @@ def test_env_state_1d_to_2d():
def test_image_3d_to_4d():
"""Test that 3D image tensors get unsqueezed to 4D."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Test observation.image
image_3d = torch.randn(224, 224, 3)
@@ -90,7 +95,7 @@ def test_image_3d_to_4d():
def test_multiple_images_3d_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
image1_3d = torch.randn(64, 64, 3)
@@ -115,7 +120,7 @@ def test_multiple_images_3d_to_4d():
def test_already_batched_tensors_unchanged():
"""Test that already batched tensors remain unchanged."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Create already batched tensors
state_2d = torch.randn(1, 7)
@@ -141,7 +146,7 @@ def test_already_batched_tensors_unchanged():
def test_higher_dimensional_tensors_unchanged():
"""Test that tensors with more dimensions than expected remain unchanged."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Create tensors with more dimensions
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():
"""Test that non-tensor values in observations remain unchanged."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
observation = {
OBS_STATE: [1, 2, 3], # List, not tensor
@@ -187,7 +192,7 @@ def test_non_tensor_values_unchanged():
def test_none_observation():
"""Test processor handles None observation gracefully."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
transition = create_transition(observation=None)
result = processor(transition)
@@ -197,7 +202,7 @@ def test_none_observation():
def test_empty_observation():
"""Test processor handles empty observation dict."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
observation = {}
transition = create_transition(observation=observation)
@@ -209,7 +214,7 @@ def test_empty_observation():
def test_mixed_observation():
"""Test processor with mixed observation containing various types and dimensions."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
state_1d = torch.randn(5)
env_state_2d = torch.randn(1, 8) # Already batched
@@ -241,8 +246,8 @@ def test_mixed_observation():
def test_integration_with_robot_processor():
"""Test ToBatchProcessor integration with RobotProcessor."""
to_batch_processor = ToBatchProcessor()
"""Test AddBatchDimensionProcessorStep integration with RobotProcessor."""
to_batch_processor = AddBatchDimensionProcessorStep()
pipeline = DataProcessorPipeline([to_batch_processor], to_transition=lambda x: x, to_output=lambda x: x)
# Create unbatched observation
@@ -261,7 +266,7 @@ def test_integration_with_robot_processor():
def test_serialization_methods():
"""Test get_config, state_dict, load_state_dict, and reset methods."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Test get_config
config = processor.get_config()
@@ -281,8 +286,8 @@ def test_serialization_methods():
def test_save_and_load_pretrained():
"""Test saving and loading ToBatchProcessor with RobotProcessor."""
processor = ToBatchProcessor()
"""Test saving and loading AddBatchDimensionProcessorStep with RobotProcessor."""
processor = AddBatchDimensionProcessorStep()
pipeline = DataProcessorPipeline(
[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 len(loaded_pipeline) == 1
assert isinstance(loaded_pipeline.steps[0], ToBatchProcessor)
assert isinstance(loaded_pipeline.steps[0], AddBatchDimensionProcessorStep)
# Test functionality of loaded processor
observation = {OBS_STATE: torch.randn(5)}
@@ -313,10 +318,10 @@ def test_save_and_load_pretrained():
def test_registry_functionality():
"""Test that ToBatchProcessor is properly registered."""
"""Test that AddBatchDimensionProcessorStep is properly registered."""
# Check that the processor is registered
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
assert "to_batch_processor" in ProcessorStepRegistry.list()
@@ -324,7 +329,7 @@ def test_registry_functionality():
def test_registry_based_save_load():
"""Test saving and loading using registry name."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
pipeline = DataProcessorPipeline([processor], to_transition=lambda x: x, to_output=lambda x: x)
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")
def test_device_compatibility():
"""Test processor works with tensors on different devices."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Create tensors on GPU
state_1d = torch.randn(7, device="cuda")
@@ -374,7 +379,7 @@ def test_device_compatibility():
def test_processor_preserves_other_transition_keys():
"""Test that processor only modifies observation and preserves other transition keys."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
action = torch.randn(5)
reward = 1.5
@@ -411,7 +416,7 @@ def test_processor_preserves_other_transition_keys():
def test_edge_case_zero_dimensional_tensors():
"""Test processor handles 0D tensors (scalars) correctly."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# 0D tensors should not be modified
scalar_tensor = torch.tensor(42.0)
@@ -433,7 +438,7 @@ def test_edge_case_zero_dimensional_tensors():
# Action-specific tests
def test_action_1d_to_2d():
"""Test that 1D action tensors get batch dimension added."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Create 1D action tensor
action_1d = torch.randn(4)
@@ -448,7 +453,7 @@ def test_action_1d_to_2d():
def test_action_already_batched():
"""Test that already batched action tensors remain unchanged."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Test various batch sizes
action_batched_1 = torch.randn(1, 4)
@@ -467,7 +472,7 @@ def test_action_already_batched():
def test_action_higher_dimensional():
"""Test that higher dimensional action tensors remain unchanged."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# 3D action tensor (e.g., sequence of actions)
action_3d = torch.randn(2, 4, 3)
@@ -484,7 +489,7 @@ def test_action_higher_dimensional():
def test_action_scalar_tensor():
"""Test that scalar (0D) action tensors remain unchanged."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
action_scalar = torch.tensor(1.5)
transition = create_transition(action=action_scalar)
@@ -497,7 +502,7 @@ def test_action_scalar_tensor():
def test_action_non_tensor():
"""Test that non-tensor actions remain unchanged."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# List action
action_list = [0.1, 0.2, 0.3, 0.4]
@@ -526,7 +531,7 @@ def test_action_non_tensor():
def test_action_none():
"""Test that None action is handled correctly."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
transition = create_transition(action=None)
result = processor(transition)
@@ -535,7 +540,7 @@ def test_action_none():
def test_action_with_observation():
"""Test action processing together with observation processing."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Both need batching
observation = {
@@ -555,7 +560,7 @@ def test_action_with_observation():
def test_action_different_sizes():
"""Test action processing with various action dimensions."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Different action sizes (robot with different DOF)
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")
def test_action_device_compatibility():
"""Test action processing on different devices."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# CUDA action
action_cuda = torch.randn(4, device="cuda")
@@ -593,7 +598,7 @@ def test_action_device_compatibility():
def test_action_dtype_preservation():
"""Test that action dtype is preserved during processing."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Different dtypes
dtypes = [torch.float32, torch.float64, torch.int32, torch.int64]
@@ -609,7 +614,7 @@ def test_action_dtype_preservation():
def test_empty_action_tensor():
"""Test handling of empty action tensors."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Empty 1D tensor
action_empty = torch.tensor([])
@@ -631,7 +636,7 @@ def test_empty_action_tensor():
# Task-specific tests
def test_task_string_to_list():
"""Test that string tasks get wrapped in lists to add batch dimension."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Create complementary data with string task
complementary_data = {"task": "pick_cube"}
@@ -648,7 +653,7 @@ def test_task_string_to_list():
def test_task_string_validation():
"""Test that only string and list of strings are valid task values."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Valid string task - should be converted to list
complementary_data = {"task": "valid_task"}
@@ -667,7 +672,7 @@ def test_task_string_validation():
def test_task_list_of_strings():
"""Test that lists of strings remain unchanged (already batched)."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Test various list of strings
test_lists = [
@@ -693,7 +698,7 @@ def test_task_list_of_strings():
def test_complementary_data_none():
"""Test processor handles None complementary_data gracefully."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
transition = create_transition(complementary_data=None)
result = processor(transition)
@@ -703,7 +708,7 @@ def test_complementary_data_none():
def test_complementary_data_empty():
"""Test processor handles empty complementary_data dict."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
complementary_data = {}
transition = create_transition(complementary_data=complementary_data)
@@ -715,7 +720,7 @@ def test_complementary_data_empty():
def test_complementary_data_no_task():
"""Test processor handles complementary_data without task field."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
complementary_data = {
"episode_id": 123,
@@ -733,7 +738,7 @@ def test_complementary_data_no_task():
def test_complementary_data_mixed():
"""Test processor with mixed complementary_data containing task and other fields."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
complementary_data = {
"task": "stack_blocks",
@@ -758,7 +763,7 @@ def test_complementary_data_mixed():
def test_task_with_observation_and_action():
"""Test task processing together with observation and action processing."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# All components need batching
observation = {
@@ -783,7 +788,7 @@ def test_task_with_observation_and_action():
def test_task_comprehensive_string_cases():
"""Test task processing with comprehensive string cases and edge cases."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Test various string formats
string_tasks = [
@@ -841,7 +846,7 @@ def test_task_comprehensive_string_cases():
def test_task_preserves_other_keys():
"""Test that task processing preserves other keys in complementary_data."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
complementary_data = {
"task": "clean_table",
@@ -869,7 +874,7 @@ def test_task_preserves_other_keys():
# Index and task_index specific tests
def test_index_scalar_to_1d():
"""Test that 0D index tensor gets unsqueezed to 1D."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Create 0D index tensor (scalar)
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():
"""Test that 0D task_index tensor gets unsqueezed to 1D."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Create 0D task_index tensor (scalar)
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():
"""Test processing both index and task_index together."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Create 0D tensors for both
index_0d = torch.tensor(100, dtype=torch.int64)
@@ -933,7 +938,7 @@ def test_index_and_task_index_together():
def test_index_already_batched():
"""Test that already batched index tensors remain unchanged."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Create already batched tensors
index_1d = torch.tensor([42], dtype=torch.int64)
@@ -954,7 +959,7 @@ def test_index_already_batched():
def test_task_index_already_batched():
"""Test that already batched task_index tensors remain unchanged."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Create already batched tensors
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():
"""Test that non-tensor index values remain unchanged."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
complementary_data = {
"index": 42, # Plain int, not tensor
@@ -992,7 +997,7 @@ def test_index_non_tensor_unchanged():
def test_index_dtype_preservation():
"""Test that index and task_index dtype is preserved during processing."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Test different dtypes
dtypes = [torch.int32, torch.int64, torch.long]
@@ -1015,7 +1020,7 @@ def test_index_dtype_preservation():
def test_index_with_full_transition():
"""Test index/task_index processing with full transition data."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Create full transition with all components
observation = {
@@ -1057,7 +1062,7 @@ def test_index_with_full_transition():
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_index_device_compatibility():
"""Test processor works with index/task_index tensors on different devices."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Create tensors on GPU
index_0d = torch.tensor(42, dtype=torch.int64, device="cuda")
@@ -1081,7 +1086,7 @@ def test_index_device_compatibility():
def test_empty_index_tensor():
"""Test handling of empty index tensors."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
# Empty 0D tensor doesn't make sense, but test empty 1D
index_empty = torch.tensor([], dtype=torch.int64)
@@ -1096,7 +1101,7 @@ def test_empty_index_tensor():
def test_action_processing_creates_new_transition():
"""Test that the processor creates a new transition object with correctly processed action."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
action = torch.randn(4)
transition = create_transition(action=action)
@@ -1118,7 +1123,7 @@ def test_action_processing_creates_new_transition():
def test_task_processing_creates_new_transition():
"""Test that the processor creates a new transition object with correctly processed task."""
processor = ToBatchProcessor()
processor = AddBatchDimensionProcessorStep()
complementary_data = {"task": "sort_objects"}
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.processor import (
DataProcessorPipeline,
DeviceProcessor,
DeviceProcessorStep,
IdentityProcessorStep,
NormalizerProcessor,
NormalizerProcessorStep,
TransitionKey,
)
@@ -87,13 +87,13 @@ def test_make_classifier_processor_basic():
# Check steps in preprocessor
assert len(preprocessor.steps) == 3
assert isinstance(preprocessor.steps[0], NormalizerProcessor) # For input features
assert isinstance(preprocessor.steps[1], NormalizerProcessor) # For output features
assert isinstance(preprocessor.steps[2], DeviceProcessor)
assert isinstance(preprocessor.steps[0], NormalizerProcessorStep) # For input features
assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) # For output features
assert isinstance(preprocessor.steps[2], DeviceProcessorStep)
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], IdentityProcessorStep)
@@ -291,11 +291,11 @@ def test_classifier_processor_mixed_precision():
# Get the steps from the factory function
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 = []
for step in factory_preprocessor.steps:
if isinstance(step, DeviceProcessor):
modified_steps.append(DeviceProcessor(device=config.device, float_dtype="float16"))
if isinstance(step, DeviceProcessorStep):
modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16"))
else:
modified_steps.append(step)
tests/processor/test_device_processor.py
+68 -64
View File
@@ -19,7 +19,7 @@ import pytest
import torch
from lerobot.configs.types import FeatureType, PolicyFeature
from lerobot.processor import DataProcessorPipeline, DeviceProcessor, TransitionKey
from lerobot.processor import DataProcessorPipeline, DeviceProcessorStep, TransitionKey
def create_transition(
@@ -46,7 +46,7 @@ def create_transition(
def test_basic_functionality():
"""Test basic device processor functionality on CPU."""
processor = DeviceProcessor(device="cpu")
processor = DeviceProcessorStep(device="cpu")
# Create a transition with CPU tensors
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")
def test_cuda_functionality():
"""Test device processor functionality on CUDA."""
processor = DeviceProcessor(device="cuda")
processor = DeviceProcessorStep(device="cuda")
# Create a transition with CPU tensors
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")
def test_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)}
action = torch.randn(5)
@@ -116,7 +116,7 @@ def test_specific_cuda_device():
def test_non_tensor_values():
"""Test that non-tensor values are preserved."""
processor = DeviceProcessor(device="cpu")
processor = DeviceProcessorStep(device="cpu")
observation = {
"observation.state": torch.randn(10),
@@ -142,7 +142,7 @@ def test_non_tensor_values():
def test_none_values():
"""Test handling of None values."""
processor = DeviceProcessor(device="cpu")
processor = DeviceProcessorStep(device="cpu")
# Test with None observation
transition = create_transition(observation=None, action=torch.randn(5))
@@ -159,7 +159,7 @@ def test_none_values():
def test_empty_observation():
"""Test handling of empty observation dictionary."""
processor = DeviceProcessor(device="cpu")
processor = DeviceProcessorStep(device="cpu")
transition = create_transition(observation={}, action=torch.randn(5))
result = processor(transition)
@@ -170,7 +170,7 @@ def test_empty_observation():
def test_scalar_tensors():
"""Test handling of scalar tensors."""
processor = DeviceProcessor(device="cpu")
processor = DeviceProcessorStep(device="cpu")
observation = {"observation.scalar": torch.tensor(1.5)}
action = torch.tensor(2.0)
@@ -187,7 +187,7 @@ def test_scalar_tensors():
def test_dtype_preservation():
"""Test that tensor dtypes are preserved."""
processor = DeviceProcessor(device="cpu")
processor = DeviceProcessorStep(device="cpu")
observation = {
"observation.float32": torch.randn(5, dtype=torch.float32),
@@ -209,7 +209,7 @@ def test_dtype_preservation():
def test_shape_preservation():
"""Test that tensor shapes are preserved."""
processor = DeviceProcessor(device="cpu")
processor = DeviceProcessorStep(device="cpu")
observation = {
"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")
def test_mixed_devices():
"""Test handling of tensors already on different devices."""
processor = DeviceProcessor(device="cuda")
processor = DeviceProcessorStep(device="cuda")
# Create tensors on different devices
observation = {
@@ -253,22 +253,22 @@ def test_mixed_devices():
def test_non_blocking_flag():
"""Test that non_blocking flag is set correctly."""
# CPU processor should have non_blocking=False
cpu_processor = DeviceProcessor(device="cpu")
cpu_processor = DeviceProcessorStep(device="cpu")
assert cpu_processor.non_blocking is False
if torch.cuda.is_available():
# CUDA processor should have non_blocking=True
cuda_processor = DeviceProcessor(device="cuda")
cuda_processor = DeviceProcessorStep(device="cuda")
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
def test_serialization_methods():
"""Test get_config, state_dict, and load_state_dict methods."""
device = "cuda" if torch.cuda.is_available() else "cpu"
processor = DeviceProcessor(device=device)
processor = DeviceProcessorStep(device=device)
# Test get_config
config = processor.get_config()
@@ -289,7 +289,7 @@ def test_serialization_methods():
def test_features():
"""Test that features returns features unchanged."""
processor = DeviceProcessor(device="cpu")
processor = DeviceProcessorStep(device="cpu")
features = {
"observation.state": PolicyFeature(type=FeatureType.STATE, shape=(10,)),
@@ -304,11 +304,11 @@ def test_features():
def test_integration_with_robot_processor():
"""Test integration with RobotProcessor."""
from lerobot.constants import OBS_STATE
from lerobot.processor import ToBatchProcessor
from lerobot.processor import AddBatchDimensionProcessorStep
# Create a pipeline with DeviceProcessor
device_processor = DeviceProcessor(device="cpu")
batch_processor = ToBatchProcessor()
# Create a pipeline with DeviceProcessorStep
device_processor = DeviceProcessorStep(device="cpu")
batch_processor = AddBatchDimensionProcessorStep()
processor = DataProcessorPipeline(
steps=[batch_processor, device_processor],
@@ -333,9 +333,9 @@ def test_integration_with_robot_processor():
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"
processor = DeviceProcessor(device=device, float_dtype="float16")
processor = DeviceProcessorStep(device=device, float_dtype="float16")
robot_processor = DataProcessorPipeline(steps=[processor], name="device_test_processor")
with tempfile.TemporaryDirectory() as tmpdir:
@@ -347,7 +347,7 @@ def test_save_and_load_pretrained():
assert len(loaded_processor.steps) == 1
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
assert (
getattr(loaded_device_processor, "device", None) == device.split(":")[0]
@@ -356,18 +356,18 @@ def test_save_and_load_pretrained():
def test_registry_functionality():
"""Test that DeviceProcessor is properly registered."""
"""Test that DeviceProcessorStep is properly registered."""
from lerobot.processor import ProcessorStepRegistry
# Check that DeviceProcessor is registered
# Check that DeviceProcessorStep is registered
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")
def test_performance_with_large_tensors():
"""Test performance with large tensors and non_blocking flag."""
processor = DeviceProcessor(device="cuda")
processor = DeviceProcessorStep(device="cuda")
# Create large tensors
observation = {
@@ -389,7 +389,7 @@ def test_performance_with_large_tensors():
def test_reward_done_truncated_types():
"""Test handling of different types for reward, done, and truncated."""
processor = DeviceProcessor(device="cpu")
processor = DeviceProcessorStep(device="cpu")
# Test with scalar values (not tensors)
transition = create_transition(
@@ -429,7 +429,7 @@ def test_reward_done_truncated_types():
def test_complementary_data_preserved():
"""Test that complementary_data is preserved unchanged."""
processor = DeviceProcessor(device="cpu")
processor = DeviceProcessorStep(device="cpu")
complementary_data = {
"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]["episode_id"] == 42
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
def test_float_dtype_conversion():
"""Test float dtype conversion functionality."""
processor = DeviceProcessor(device="cpu", float_dtype="float16")
processor = DeviceProcessorStep(device="cpu", float_dtype="float16")
# Create tensors of different types
observation = {
@@ -485,7 +485,7 @@ def test_float_dtype_conversion():
def test_float_dtype_none():
"""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.float32": torch.randn(5, dtype=torch.float32),
@@ -506,7 +506,7 @@ def test_float_dtype_none():
def test_float_dtype_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)}
action = torch.randn(3, dtype=torch.float64)
@@ -520,7 +520,7 @@ def test_float_dtype_bfloat16():
def test_float_dtype_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)}
action = torch.randn(3, dtype=torch.float32)
@@ -535,27 +535,27 @@ def test_float_dtype_float64():
def test_float_dtype_invalid():
"""Test that invalid float_dtype raises ValueError."""
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():
"""Test that dtype aliases work correctly."""
# 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
# 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
# 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
def test_float_dtype_with_mixed_tensors():
"""Test float dtype conversion with mixed tensor types."""
processor = DeviceProcessor(device="cpu", float_dtype="float32")
processor = DeviceProcessorStep(device="cpu", float_dtype="float32")
observation = {
"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():
"""Test that float_dtype is properly serialized in get_config."""
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()
assert config == {"device": device, "float_dtype": "float16"}
# 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()
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")
def test_float_dtype_with_cuda():
"""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
observation = {
@@ -619,7 +619,7 @@ def test_float_dtype_with_cuda():
def test_complementary_data_index_fields():
"""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
complementary_data = {
@@ -657,7 +657,7 @@ def test_complementary_data_index_fields():
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_complementary_data_index_fields_cuda():
"""Test moving index and task_index fields to CUDA."""
processor = DeviceProcessor(device="cuda:0")
processor = DeviceProcessorStep(device="cuda:0")
# Create CPU tensors
complementary_data = {
@@ -679,7 +679,7 @@ def test_complementary_data_index_fields_cuda():
def test_complementary_data_without_index_fields():
"""Test that complementary_data without index/task_index fields works correctly."""
processor = DeviceProcessor(device="cpu")
processor = DeviceProcessorStep(device="cpu")
complementary_data = {
"task": ["navigate"],
@@ -697,7 +697,7 @@ def test_complementary_data_without_index_fields():
def test_complementary_data_mixed_tensors():
"""Test complementary_data with mix of tensors and non-tensors."""
processor = DeviceProcessor(device="cpu")
processor = DeviceProcessorStep(device="cpu")
complementary_data = {
"task": ["pick_and_place"],
@@ -726,7 +726,7 @@ def test_complementary_data_mixed_tensors():
def test_complementary_data_float_dtype_conversion():
"""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 = {
"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")
def test_complementary_data_full_pipeline_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
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():
"""Test empty complementary_data handling."""
processor = DeviceProcessor(device="cpu")
processor = DeviceProcessorStep(device="cpu")
transition = create_transition(
observation={"observation.state": torch.randn(1, 7)},
@@ -811,7 +811,7 @@ def test_complementary_data_empty():
def test_complementary_data_none():
"""Test None complementary_data handling."""
processor = DeviceProcessor(device="cpu")
processor = DeviceProcessorStep(device="cpu")
transition = create_transition(
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")
def test_preserves_gpu_placement():
"""Test that DeviceProcessor preserves GPU placement when tensor is already on GPU."""
processor = DeviceProcessor(device="cuda:0")
"""Test that DeviceProcessorStep preserves GPU placement when tensor is already on GPU."""
processor = DeviceProcessorStep(device="cuda:0")
# Create tensors already on GPU
observation = {
@@ -852,9 +852,9 @@ def test_preserves_gpu_placement():
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs")
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)
processor_gpu = DeviceProcessor(device="cuda:0")
processor_gpu = DeviceProcessorStep(device="cuda:0")
# Create tensors on cuda:1 (simulating Accelerate placement)
cuda1_device = torch.device("cuda:1")
@@ -873,7 +873,7 @@ def test_multi_gpu_preservation():
assert result[TransitionKey.ACTION].device == cuda1_device
# 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(
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():
"""Test that CPU tensors are moved to configured device even in multi-GPU context."""
# Processor configured for cuda:1
processor = DeviceProcessor(device="cuda:1")
processor = DeviceProcessorStep(device="cuda:1")
# Mix of CPU and GPU tensors
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")
def test_multi_gpu_with_float_dtype():
"""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
observation = {
@@ -946,7 +946,7 @@ def test_simulated_accelerate_scenario():
for gpu_id in range(min(torch.cuda.device_count(), 2)):
# Each "process" has a processor configured for cuda:0
# 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
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."""
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
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
features = {
@@ -984,9 +988,9 @@ def test_policy_processor_integration():
# Create input processor (preprocessor) that moves to GPU
input_processor = DataProcessorPipeline(
steps=[
NormalizerProcessor(features=features, norm_map=norm_map, stats=stats),
ToBatchProcessor(),
DeviceProcessor(device="cuda"),
NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats),
AddBatchDimensionProcessorStep(),
DeviceProcessorStep(device="cuda"),
],
name="test_preprocessor",
to_transition=lambda x: x,
@@ -996,8 +1000,8 @@ def test_policy_processor_integration():
# Create output processor (postprocessor) that moves to CPU
output_processor = DataProcessorPipeline(
steps=[
DeviceProcessor(device="cpu"),
UnnormalizerProcessor(features={ACTION: features[ACTION]}, norm_map=norm_map, stats=stats),
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(features={ACTION: features[ACTION]}, norm_map=norm_map, stats=stats),
],
name="test_postprocessor",
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.processor_diffusion import make_diffusion_pre_post_processors
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DataProcessorPipeline,
DeviceProcessor,
NormalizerProcessor,
RenameProcessor,
ToBatchProcessor,
DeviceProcessorStep,
NormalizerProcessorStep,
RenameProcessorStep,
TransitionKey,
UnnormalizerProcessor,
UnnormalizerProcessorStep,
)
@@ -89,15 +89,15 @@ def test_make_diffusion_processor_basic():
# Check steps in preprocessor
assert len(preprocessor.steps) == 4
assert isinstance(preprocessor.steps[0], RenameProcessor)
assert isinstance(preprocessor.steps[1], NormalizerProcessor)
assert isinstance(preprocessor.steps[2], ToBatchProcessor)
assert isinstance(preprocessor.steps[3], DeviceProcessor)
assert isinstance(preprocessor.steps[0], RenameProcessorStep)
assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep)
assert isinstance(preprocessor.steps[3], DeviceProcessorStep)
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessor)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_diffusion_processor_with_images():
@@ -294,11 +294,11 @@ def test_diffusion_processor_mixed_precision():
# Get the steps from the factory function
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 = []
for step in factory_preprocessor.steps:
if isinstance(step, DeviceProcessor):
modified_steps.append(DeviceProcessor(device=config.device, float_dtype="float16"))
if isinstance(step, DeviceProcessorStep):
modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16"))
else:
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 (
DataProcessorPipeline,
IdentityProcessorStep,
NormalizerProcessor,
NormalizerProcessorStep,
TransitionKey,
UnnormalizerProcessor,
UnnormalizerProcessorStep,
hotswap_stats,
)
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
def observation_stats():
return {
@@ -142,10 +142,10 @@ def observation_stats():
@pytest.fixture
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()
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):
@@ -182,7 +182,7 @@ def test_min_max_normalization(observation_normalizer):
def test_selective_normalization(observation_stats):
features = _create_observation_features()
norm_map = _create_observation_norm_map()
normalizer = NormalizerProcessor(
normalizer = NormalizerProcessorStep(
features=features,
norm_map=norm_map,
stats=observation_stats,
@@ -208,7 +208,7 @@ def test_selective_normalization(observation_stats):
def test_device_compatibility(observation_stats):
features = _create_observation_features()
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.image": torch.tensor([0.7, 0.5, 0.3]).cuda(),
}
@@ -237,7 +237,7 @@ def test_from_lerobot_dataset():
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
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
features = _create_observation_features()
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)
# 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):
features = _create_action_features()
norm_map = _create_action_norm_map_mean_std()
unnormalizer = UnnormalizerProcessor(
unnormalizer = UnnormalizerProcessorStep(
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):
features = _create_action_features()
norm_map = _create_action_norm_map_min_max()
unnormalizer = UnnormalizerProcessor(
unnormalizer = UnnormalizerProcessorStep(
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):
features = _create_action_features()
norm_map = _create_action_norm_map_mean_std()
unnormalizer = UnnormalizerProcessor(
unnormalizer = UnnormalizerProcessorStep(
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):
features = _create_action_features()
norm_map = _create_action_norm_map_mean_std()
unnormalizer = UnnormalizerProcessor(
unnormalizer = UnnormalizerProcessorStep(
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]}}
features = {"action": PolicyFeature(FeatureType.ACTION, (1,))}
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"]
# Fixtures for NormalizerProcessor tests
# Fixtures for NormalizerProcessorStep tests
@pytest.fixture
def full_stats():
return {
@@ -424,7 +424,7 @@ def _create_full_norm_map():
def normalizer_processor(full_stats):
features = _create_full_features()
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):
@@ -468,7 +468,7 @@ def test_processor_from_lerobot_dataset(full_stats):
features = _create_full_features()
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"}
)
@@ -480,7 +480,7 @@ def test_processor_from_lerobot_dataset(full_stats):
def test_get_config(full_stats):
features = _create_full_features()
norm_map = _create_full_norm_map()
processor = NormalizerProcessor(
processor = NormalizerProcessorStep(
features=features,
norm_map=norm_map,
stats=full_stats,
@@ -539,7 +539,7 @@ def test_empty_observation():
stats = {"observation.image": {"mean": [0.5], "std": [0.2]}}
features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))}
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()
result = normalizer(transition)
@@ -550,7 +550,7 @@ def test_empty_observation():
def test_empty_stats():
features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))}
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])}
transition = create_transition(observation=observation)
@@ -566,7 +566,7 @@ def test_partial_stats():
stats = {"observation.image": {"mean": [0.5]}} # Missing std / (min,max)
features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))}
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])}
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))}
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
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."""
features = _create_full_features()
norm_map = _create_full_norm_map()
original_processor = NormalizerProcessor(
original_processor = NormalizerProcessorStep(
features=features,
norm_map=norm_map,
stats=full_stats,
@@ -602,7 +602,7 @@ def test_serialization_roundtrip(full_stats):
config = original_processor.get_config()
# Create a new processor from the config (deserialization)
new_processor = NormalizerProcessor(
new_processor = NormalizerProcessorStep(
features=config["features"],
norm_map=config["norm_map"],
stats=full_stats,
@@ -670,7 +670,7 @@ def test_identity_normalization_observations():
"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.image": torch.tensor([0.7, 0.5, 0.3]),
@@ -695,7 +695,7 @@ def test_identity_normalization_actions():
norm_map = {FeatureType.ACTION: NormalizationMode.IDENTITY}
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])
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]},
}
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]),
@@ -748,7 +748,7 @@ def test_identity_unnormalization_actions():
norm_map = {FeatureType.ACTION: NormalizationMode.IDENTITY}
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
transition = create_transition(action=action)
@@ -771,8 +771,8 @@ def test_identity_with_missing_stats():
}
stats = {} # No stats provided
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
normalizer = NormalizerProcessorStep(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])}
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]},
}
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, 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]},
}
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, 0.5, 0.3]),
@@ -888,8 +888,8 @@ def test_identity_roundtrip():
"action": {"min": [-1.0, -1.0], "max": [1.0, 1.0]},
}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
normalizer = NormalizerProcessorStep(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_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]},
}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
# Get config
config = normalizer.get_config()
@@ -931,7 +931,7 @@ def test_identity_config_serialization():
assert config["norm_map"]["ACTION"] == "MEAN_STD"
# Create new processor from config (simulating load)
new_normalizer = NormalizerProcessor(
new_normalizer = NormalizerProcessorStep(
features=config["features"],
norm_map=config["norm_map"],
stats=stats,
@@ -969,7 +969,7 @@ def test_identity_config_serialization():
# norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD}
# 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
# normalizer.norm_map[FeatureType.STATE] = "INVALID_MODE"
@@ -1006,8 +1006,8 @@ def test_hotswap_stats_basic_functionality():
}
# Create processors
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats)
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats)
normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats)
unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats)
identity = IdentityProcessorStep()
# Create robot processor
@@ -1047,7 +1047,7 @@ def test_hotswap_stats_deep_copy():
}
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])
# Store reference to original stats
@@ -1072,7 +1072,7 @@ def test_hotswap_stats_deep_copy():
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 = {
"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}
# Create mixed steps
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
identity = IdentityProcessorStep()
robot_processor = DataProcessorPipeline(steps=[normalizer, identity, unnormalizer])
@@ -1117,7 +1117,7 @@ def test_hotswap_stats_empty_stats():
}
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])
# Hotswap with empty stats
@@ -1167,7 +1167,7 @@ def test_hotswap_stats_preserves_other_attributes():
normalize_observation_keys = {"observation.image"}
eps = 1e-6
normalizer = NormalizerProcessor(
normalizer = NormalizerProcessorStep(
features=features,
norm_map=norm_map,
stats=initial_stats,
@@ -1212,10 +1212,10 @@ def test_hotswap_stats_multiple_normalizer_types():
}
# Create multiple normalizers and unnormalizers
normalizer1 = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats)
normalizer2 = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats)
unnormalizer1 = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats)
unnormalizer2 = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats)
normalizer1 = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats)
normalizer2 = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats)
unnormalizer1 = UnnormalizerProcessorStep(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])
@@ -1264,7 +1264,7 @@ def test_hotswap_stats_with_different_data_types():
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])
# Hotswap stats
@@ -1320,7 +1320,7 @@ def test_hotswap_stats_functional_test():
}
# 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(
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,))}
norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD}
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
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,))}
norm_map = {FeatureType.STATE: NormalizationMode.MIN_MAX}
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])}
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}
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"}
)
@@ -1411,12 +1411,12 @@ def test_unnormalize_observations_mean_std_and_min_max():
"observation.mm": PolicyFeature(FeatureType.STATE, (2,)),
}
# Build two processors: one mean/std and one min/max
unnorm_ms = UnnormalizerProcessor(
unnorm_ms = UnnormalizerProcessorStep(
features={"observation.ms": features["observation.ms"]},
norm_map={FeatureType.STATE: NormalizationMode.MEAN_STD},
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"]},
norm_map={FeatureType.STATE: NormalizationMode.MIN_MAX},
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,))}
norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD}
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])}
tr = create_transition(observation=obs)
@@ -1452,7 +1452,7 @@ def test_batched_action_normalization():
features = {"action": PolicyFeature(FeatureType.ACTION, (2,))}
norm_map = {FeatureType.ACTION: NormalizationMode.MEAN_STD}
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]
out = normalizer(create_transition(action=actions))[TransitionKey.ACTION]
@@ -1464,7 +1464,7 @@ def test_complementary_data_preservation():
features = {"observation.state": PolicyFeature(FeatureType.STATE, (1,))}
norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD}
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}
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])},
"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)
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
normalizer = NormalizerProcessorStep(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)
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.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
@@ -41,7 +41,7 @@ def create_transition(
def test_process_single_image():
"""Test processing a single image."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
# Create a mock image (H, W, C) format, 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():
"""Test processing multiple images in a dictionary."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
# Create mock images
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():
"""Test processing already batched images."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
# Create a batched image (B, H, W, C)
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():
"""Test error handling for invalid image formats."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
# Test wrong channel order (channels first)
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():
"""Test error handling for invalid image dtype."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
# Test wrong dtype
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():
"""Test processor when no pixels are in observation."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
observation = {"other_data": np.array([1, 2, 3])}
transition = create_transition(observation=observation)
@@ -148,7 +148,7 @@ def test_no_pixels_in_observation():
def test_none_observation():
"""Test processor with None observation."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
transition = create_transition()
result = processor(transition)
@@ -158,7 +158,7 @@ def test_none_observation():
def test_serialization_methods():
"""Test serialization methods."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
# Test get_config
config = processor.get_config()
@@ -177,7 +177,7 @@ def test_serialization_methods():
def test_process_environment_state():
"""Test processing environment_state."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32)
observation = {"environment_state": env_state}
@@ -198,7 +198,7 @@ def test_process_environment_state():
def test_process_agent_pos():
"""Test processing agent_pos."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32)
observation = {"agent_pos": agent_pos}
@@ -219,7 +219,7 @@ def test_process_agent_pos():
def test_process_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)
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():
"""Test processing both environment_state and agent_pos."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
env_state = np.array([1.0, 2.0], 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():
"""Test processor when no states are in observation."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
observation = {"other_data": np.array([1, 2, 3])}
transition = create_transition(observation=observation)
@@ -276,7 +276,7 @@ def test_no_states_in_observation():
def test_complete_observation_processing():
"""Test processing a complete observation with both images and states."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
# Create mock data
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():
"""Test processing observation with only images."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8)
observation = {"pixels": image}
@@ -328,7 +328,7 @@ def test_image_only_processing():
def test_state_only_processing():
"""Test processing observation with only states."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
agent_pos = np.array([1.0, 2.0], dtype=np.float32)
observation = {"agent_pos": agent_pos}
@@ -343,7 +343,7 @@ def test_state_only_processing():
def test_empty_observation():
"""Test processing empty observation."""
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
observation = {}
transition = create_transition(observation=observation)
@@ -359,7 +359,7 @@ def test_equivalent_to_original_function():
# Import the original function for comparison
from lerobot.envs.utils import preprocess_observation
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
# Create test data similar to what the original function expects
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."""
from lerobot.envs.utils import preprocess_observation
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
# Create test data with multiple cameras
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):
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
features = {
"pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
"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):
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
features = {
"observation.pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
"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):
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
features = {
"pixels.front": 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):
processor = VanillaObservationProcessor()
processor = VanillaObservationProcessorStep()
features = {
"environment_state": policy_feature_factory(FeatureType.STATE, (3,)),
"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):
proc = VanillaObservationProcessor()
proc = VanillaObservationProcessorStep()
features = {
"observation.environment_state": policy_feature_factory(FeatureType.STATE, (2,)),
"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.processor_pi0 import Pi0NewLineProcessor, make_pi0_pre_post_processors
from lerobot.processor import (
DeviceProcessor,
NormalizerProcessor,
RenameProcessor,
ToBatchProcessor,
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
RenameProcessorStep,
TransitionKey,
UnnormalizerProcessor,
UnnormalizerProcessorStep,
)
@@ -83,7 +83,7 @@ def test_make_pi0_processor_basic():
config = create_default_config()
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(
config,
stats,
@@ -97,17 +97,17 @@ def test_make_pi0_processor_basic():
# Check steps in preprocessor
assert len(preprocessor.steps) == 6
assert isinstance(preprocessor.steps[0], RenameProcessor)
assert isinstance(preprocessor.steps[1], NormalizerProcessor)
assert isinstance(preprocessor.steps[2], ToBatchProcessor)
assert isinstance(preprocessor.steps[0], RenameProcessorStep)
assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep)
assert isinstance(preprocessor.steps[3], Pi0NewLineProcessor)
# Step 4 would be TokenizerProcessor but it's mocked
assert isinstance(preprocessor.steps[5], DeviceProcessor)
# Step 4 would be TokenizerProcessorStep but it's mocked
assert isinstance(preprocessor.steps[5], DeviceProcessorStep)
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessor)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_pi0_newline_processor_single_task():
@@ -165,7 +165,7 @@ def test_pi0_processor_cuda():
stats = create_default_stats()
# Mock the tokenizer processor to act as pass-through
class MockTokenizerProcessor:
class MockTokenizerProcessorStep:
def __init__(self, *args, **kwargs):
pass
@@ -187,7 +187,7 @@ def test_pi0_processor_cuda():
def transform_features(self, 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(
config,
stats,
@@ -220,7 +220,7 @@ def test_pi0_processor_accelerate_scenario():
stats = create_default_stats()
# Mock the tokenizer processor to act as pass-through
class MockTokenizerProcessor:
class MockTokenizerProcessorStep:
def __init__(self, *args, **kwargs):
pass
@@ -242,7 +242,7 @@ def test_pi0_processor_accelerate_scenario():
def transform_features(self, 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(
config,
stats,
@@ -276,7 +276,7 @@ def test_pi0_processor_multi_gpu():
stats = create_default_stats()
# Mock the tokenizer processor to act as pass-through
class MockTokenizerProcessor:
class MockTokenizerProcessorStep:
def __init__(self, *args, **kwargs):
pass
@@ -298,7 +298,7 @@ def test_pi0_processor_multi_gpu():
def transform_features(self, 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(
config,
stats,
@@ -329,7 +329,7 @@ def test_pi0_processor_without_stats():
config = create_default_config()
# 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(
config,
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 (
DataProcessorPipeline,
ProcessorStepRegistry,
RenameProcessor,
RenameProcessorStep,
TransitionKey,
)
from lerobot.processor.rename_processor import rename_stats
@@ -51,7 +51,7 @@ def test_basic_renaming():
"old_key1": "new_key1",
"old_key2": "new_key2",
}
processor = RenameProcessor(rename_map=rename_map)
processor = RenameProcessorStep(rename_map=rename_map)
observation = {
"old_key1": torch.tensor([1.0, 2.0]),
@@ -79,7 +79,7 @@ def test_basic_renaming():
def test_empty_rename_map():
"""Test processor with empty rename map (should pass through unchanged)."""
processor = RenameProcessor(rename_map={})
processor = RenameProcessorStep(rename_map={})
observation = {
"key1": torch.tensor([1.0]),
@@ -98,7 +98,7 @@ def test_empty_rename_map():
def test_none_observation():
"""Test processor with None observation."""
processor = RenameProcessor(rename_map={"old": "new"})
processor = RenameProcessorStep(rename_map={"old": "new"})
transition = create_transition()
result = processor(transition)
@@ -113,7 +113,7 @@ def test_overlapping_rename():
"a": "b",
"b": "c", # This creates a potential conflict
}
processor = RenameProcessor(rename_map=rename_map)
processor = RenameProcessorStep(rename_map=rename_map)
observation = {
"a": 1,
@@ -138,7 +138,7 @@ def test_partial_rename():
"observation.state": "observation.proprio_state",
"pixels": "observation.image",
}
processor = RenameProcessor(rename_map=rename_map)
processor = RenameProcessorStep(rename_map=rename_map)
observation = {
"observation.state": torch.randn(10),
@@ -168,15 +168,15 @@ def test_get_config():
"old1": "new1",
"old2": "new2",
}
processor = RenameProcessor(rename_map=rename_map)
processor = RenameProcessorStep(rename_map=rename_map)
config = processor.get_config()
assert config == {"rename_map": rename_map}
def test_state_dict():
"""Test state dict (should be empty for RenameProcessor)."""
processor = RenameProcessor(rename_map={"old": "new"})
"""Test state dict (should be empty for RenameProcessorStep)."""
processor = RenameProcessorStep(rename_map={"old": "new"})
state = processor.state_dict()
assert state == {}
@@ -191,7 +191,7 @@ def test_integration_with_robot_processor():
"agent_pos": "observation.state",
"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)
@@ -225,18 +225,20 @@ def test_save_and_load_pretrained():
"old_state": "observation.state",
"old_image": "observation.image",
}
processor = RenameProcessor(rename_map=rename_map)
pipeline = DataProcessorPipeline([processor], name="TestRenameProcessor")
processor = RenameProcessorStep(rename_map=rename_map)
pipeline = DataProcessorPipeline([processor], name="TestRenameProcessorStep")
with tempfile.TemporaryDirectory() as tmp_dir:
# Save pipeline
pipeline.save_pretrained(tmp_dir)
# 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()
# No state files should be created for RenameProcessor
# No state files should be created for RenameProcessorStep
state_files = list(Path(tmp_dir).glob("*.safetensors"))
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
)
assert loaded_pipeline.name == "TestRenameProcessor"
assert loaded_pipeline.name == "TestRenameProcessorStep"
assert len(loaded_pipeline) == 1
# Check that loaded processor works correctly
loaded_processor = loaded_pipeline.steps[0]
assert isinstance(loaded_processor, RenameProcessor)
assert isinstance(loaded_processor, RenameProcessorStep)
assert loaded_processor.rename_map == rename_map
# Test functionality after loading
@@ -267,23 +269,23 @@ def test_save_and_load_pretrained():
def test_registry_functionality():
"""Test that RenameProcessor is properly registered."""
"""Test that RenameProcessorStep is properly registered."""
# Check that it's registered
assert "rename_processor" in ProcessorStepRegistry.list()
# Get from registry
retrieved_class = ProcessorStepRegistry.get("rename_processor")
assert retrieved_class is RenameProcessor
assert retrieved_class is RenameProcessorStep
# Create instance from registry
instance = retrieved_class(rename_map={"old": "new"})
assert isinstance(instance, RenameProcessor)
assert isinstance(instance, RenameProcessorStep)
assert instance.rename_map == {"old": "new"}
def test_registry_based_save_load():
"""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)
with tempfile.TemporaryDirectory() as tmp_dir:
@@ -303,14 +305,14 @@ def test_registry_based_save_load():
# Load should work
loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir)
loaded_processor = loaded_pipeline.steps[0]
assert isinstance(loaded_processor, RenameProcessor)
assert isinstance(loaded_processor, RenameProcessorStep)
assert loaded_processor.rename_map == {"key1": "renamed_key1"}
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
processor1 = RenameProcessor(
processor1 = RenameProcessorStep(
rename_map={
"pos": "agent_position",
"img": "camera_image",
@@ -318,7 +320,7 @@ def test_chained_rename_processors():
)
# Second processor: rename to final format
processor2 = RenameProcessor(
processor2 = RenameProcessorStep(
rename_map={
"agent_position": "observation.state",
"camera_image": "observation.image",
@@ -363,7 +365,7 @@ def test_nested_observation_rename():
"observation.images.right": "observation.camera.right_view",
"observation.proprio": "observation.proprioception",
}
processor = RenameProcessor(rename_map=rename_map)
processor = RenameProcessorStep(rename_map=rename_map)
observation = {
"observation.images.left": torch.randn(3, 64, 64),
@@ -393,7 +395,7 @@ def test_nested_observation_rename():
def test_value_types_preserved():
"""Test that various value types are preserved during renaming."""
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)
array_value = np.random.rand(2, 2)
@@ -421,7 +423,7 @@ def test_value_types_preserved():
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 = {
"a": policy_feature_factory(FeatureType.STATE, (2,)),
"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):
# 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 = {
"a": policy_feature_factory(FeatureType.STATE, (1,)),
"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):
# Chain two rename processors at the contract level
processor1 = RenameProcessor(rename_map={"pos": "agent_position", "img": "camera_image"})
processor2 = RenameProcessor(
processor1 = RenameProcessorStep(rename_map={"pos": "agent_position", "img": "camera_image"})
processor2 = RenameProcessorStep(
rename_map={"agent_position": "observation.state", "camera_image": "observation.image"}
)
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.processor_sac import make_sac_pre_post_processors
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DataProcessorPipeline,
DeviceProcessor,
NormalizerProcessor,
RenameProcessor,
ToBatchProcessor,
DeviceProcessorStep,
NormalizerProcessorStep,
RenameProcessorStep,
TransitionKey,
UnnormalizerProcessor,
UnnormalizerProcessorStep,
)
@@ -91,15 +91,15 @@ def test_make_sac_processor_basic():
# Check steps in preprocessor
assert len(preprocessor.steps) == 4
assert isinstance(preprocessor.steps[0], RenameProcessor)
assert isinstance(preprocessor.steps[1], NormalizerProcessor)
assert isinstance(preprocessor.steps[2], ToBatchProcessor)
assert isinstance(preprocessor.steps[3], DeviceProcessor)
assert isinstance(preprocessor.steps[0], RenameProcessorStep)
assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep)
assert isinstance(preprocessor.steps[3], DeviceProcessorStep)
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessor)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
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},
)
# Replace DeviceProcessor with one that uses float16
# Replace DeviceProcessorStep with one that uses float16
for i, step in enumerate(preprocessor.steps):
if isinstance(step, DeviceProcessor):
preprocessor.steps[i] = DeviceProcessor(device=config.device, float_dtype="float16")
if isinstance(step, DeviceProcessorStep):
preprocessor.steps[i] = DeviceProcessorStep(device=config.device, float_dtype="float16")
# Create test data
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,
)
from lerobot.processor import (
DeviceProcessor,
NormalizerProcessor,
RenameProcessor,
ToBatchProcessor,
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
RenameProcessorStep,
TransitionKey,
UnnormalizerProcessor,
UnnormalizerProcessorStep,
)
@@ -88,7 +88,7 @@ def test_make_smolvla_processor_basic():
config = create_default_config()
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(
config,
stats,
@@ -102,17 +102,17 @@ def test_make_smolvla_processor_basic():
# Check steps in preprocessor
assert len(preprocessor.steps) == 6
assert isinstance(preprocessor.steps[0], RenameProcessor)
assert isinstance(preprocessor.steps[1], NormalizerProcessor)
assert isinstance(preprocessor.steps[2], ToBatchProcessor)
assert isinstance(preprocessor.steps[0], RenameProcessorStep)
assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep)
assert isinstance(preprocessor.steps[3], SmolVLANewLineProcessor)
# Step 4 would be TokenizerProcessor but it's mocked
assert isinstance(preprocessor.steps[5], DeviceProcessor)
# Step 4 would be TokenizerProcessorStep but it's mocked
assert isinstance(preprocessor.steps[5], DeviceProcessorStep)
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessor)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_smolvla_newline_processor_single_task():
@@ -170,7 +170,7 @@ def test_smolvla_processor_cuda():
stats = create_default_stats()
# Mock the tokenizer processor to act as pass-through
class MockTokenizerProcessor:
class MockTokenizerProcessorStep:
def __init__(self, *args, **kwargs):
pass
@@ -192,7 +192,9 @@ def test_smolvla_processor_cuda():
def transform_features(self, 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(
config,
stats,
@@ -225,7 +227,7 @@ def test_smolvla_processor_accelerate_scenario():
stats = create_default_stats()
# Mock the tokenizer processor to act as pass-through
class MockTokenizerProcessor:
class MockTokenizerProcessorStep:
def __init__(self, *args, **kwargs):
pass
@@ -247,7 +249,9 @@ def test_smolvla_processor_accelerate_scenario():
def transform_features(self, 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(
config,
stats,
@@ -281,7 +285,7 @@ def test_smolvla_processor_multi_gpu():
stats = create_default_stats()
# Mock the tokenizer processor to act as pass-through
class MockTokenizerProcessor:
class MockTokenizerProcessorStep:
def __init__(self, *args, **kwargs):
pass
@@ -303,7 +307,9 @@ def test_smolvla_processor_multi_gpu():
def transform_features(self, 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(
config,
stats,
@@ -334,7 +340,7 @@ def test_smolvla_processor_without_stats():
config = create_default_config()
# 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(
config,
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.processor_tdmpc import make_tdmpc_pre_post_processors
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DataProcessorPipeline,
DeviceProcessor,
NormalizerProcessor,
RenameProcessor,
ToBatchProcessor,
DeviceProcessorStep,
NormalizerProcessorStep,
RenameProcessorStep,
TransitionKey,
UnnormalizerProcessor,
UnnormalizerProcessorStep,
)
@@ -94,15 +94,15 @@ def test_make_tdmpc_processor_basic():
# Check steps in preprocessor
assert len(preprocessor.steps) == 4
assert isinstance(preprocessor.steps[0], RenameProcessor)
assert isinstance(preprocessor.steps[1], NormalizerProcessor)
assert isinstance(preprocessor.steps[2], ToBatchProcessor)
assert isinstance(preprocessor.steps[3], DeviceProcessor)
assert isinstance(preprocessor.steps[0], RenameProcessorStep)
assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep)
assert isinstance(preprocessor.steps[3], DeviceProcessorStep)
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessor)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
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},
)
# Replace DeviceProcessor with one that uses float16
# Replace DeviceProcessorStep with one that uses float16
for i, step in enumerate(preprocessor.steps):
if isinstance(step, DeviceProcessor):
preprocessor.steps[i] = DeviceProcessor(device=config.device, float_dtype="float16")
if isinstance(step, DeviceProcessorStep):
preprocessor.steps[i] = DeviceProcessorStep(device=config.device, float_dtype="float16")
# Create test data
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
@@ -10,7 +10,7 @@ import torch
from lerobot.configs.types import FeatureType, PolicyFeature
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
@@ -95,7 +95,7 @@ def test_basic_tokenization(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100)
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(
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."""
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(
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_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(
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_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(
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_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)
@@ -230,7 +230,7 @@ def test_missing_task_key(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100)
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"})
@@ -245,7 +245,7 @@ def test_none_task_value(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100)
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})
@@ -260,7 +260,7 @@ def test_unsupported_task_type(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100)
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
transition = create_transition(complementary_data={"task": 123})
@@ -279,7 +279,7 @@ def test_unsupported_task_type(mock_auto_tokenizer):
def test_no_tokenizer_error():
"""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"):
TokenizerProcessor()
TokenizerProcessorStep()
@require_package("transformers")
@@ -290,7 +290,7 @@ def test_invalid_tokenizer_name_error():
mock_auto_tokenizer.from_pretrained.side_effect = Exception("Model not found")
with pytest.raises(Exception, match="Model not found"):
TokenizerProcessor(tokenizer_name="invalid-tokenizer")
TokenizerProcessorStep(tokenizer_name="invalid-tokenizer")
@require_package("transformers")
@@ -300,7 +300,7 @@ def test_get_config_with_tokenizer_name(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(
processor = TokenizerProcessorStep(
tokenizer_name="test-tokenizer",
max_length=256,
task_key="instruction",
@@ -327,7 +327,7 @@ def test_get_config_with_tokenizer_object():
"""Test configuration serialization when using tokenizer object."""
mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(
processor = TokenizerProcessorStep(
tokenizer=mock_tokenizer,
max_length=256,
task_key="instruction",
@@ -357,7 +357,7 @@ def test_state_dict_methods(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100)
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
state = processor.state_dict()
@@ -374,7 +374,7 @@ def test_reset_method(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100)
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
processor.reset()
@@ -387,7 +387,7 @@ def test_integration_with_robot_processor(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100)
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(
[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_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
original_processor = TokenizerProcessor(
original_processor = TokenizerProcessorStep(
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."""
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(
[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
retrieved_class = ProcessorStepRegistry.get("tokenizer_processor")
assert retrieved_class is TokenizerProcessor
assert retrieved_class is TokenizerProcessorStep
@require_package("transformers")
def test_features_basic():
"""Test basic feature contract functionality."""
mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=128)
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=128)
input_features = {
"observation.state": PolicyFeature(type=FeatureType.STATE, shape=(10,)),
@@ -538,7 +540,7 @@ def test_features_basic():
def test_features_with_custom_max_length():
"""Test feature contract with custom max_length."""
mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=64)
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=64)
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():
"""Test feature contract when tokenized features already exist."""
mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=256)
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=256)
input_features = {
f"{OBS_LANGUAGE}.tokens": PolicyFeature(type=FeatureType.LANGUAGE, shape=(100,)),
@@ -595,7 +597,7 @@ def test_tokenization_parameters(mock_auto_tokenizer):
tracking_tokenizer = TrackingMockTokenizer()
mock_auto_tokenizer.from_pretrained.return_value = tracking_tokenizer
processor = TokenizerProcessor(
processor = TokenizerProcessorStep(
tokenizer_name="test-tokenizer",
max_length=16,
padding="longest",
@@ -627,7 +629,7 @@ def test_preserves_other_complementary_data(mock_auto_tokenizer):
mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
processor = TokenizerProcessor(tokenizer_name="test-tokenizer")
processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")
transition = create_transition(
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_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(
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_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(
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_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
transition = create_transition(
@@ -764,7 +766,9 @@ def test_custom_padding_side(mock_auto_tokenizer):
mock_auto_tokenizer.from_pretrained.return_value = tracking_tokenizer
# 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(
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"
# 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)
@@ -787,7 +793,7 @@ def test_custom_padding_side(mock_auto_tokenizer):
def test_device_detection_cpu():
"""Test that tokenized tensors stay on CPU when other tensors are on CPU."""
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
observation = {"observation.state": torch.randn(10)} # CPU tensor
@@ -811,7 +817,7 @@ def test_device_detection_cpu():
def test_device_detection_cuda():
"""Test that tokenized tensors are moved to CUDA when other tensors are on CUDA."""
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
observation = {"observation.state": torch.randn(10).cuda()} # CUDA tensor
@@ -836,7 +842,7 @@ def test_device_detection_cuda():
def test_device_detection_multi_gpu():
"""Test that tokenized tensors match device in multi-GPU setup."""
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
device = torch.device("cuda:1")
@@ -860,7 +866,7 @@ def test_device_detection_multi_gpu():
def test_device_detection_no_tensors():
"""Test that tokenized tensors stay on CPU when no other tensors exist."""
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
transition = create_transition(
@@ -882,7 +888,7 @@ def test_device_detection_no_tensors():
def test_device_detection_mixed_devices():
"""Test device detection when tensors are on different devices (uses first found)."""
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():
# Create transition with mixed devices
@@ -910,7 +916,7 @@ def test_device_detection_mixed_devices():
def test_device_detection_from_action():
"""Test that device is detected from action tensor when no observation tensors exist."""
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
observation = {"metadata": {"key": "value"}} # No tensors in observation
@@ -933,7 +939,7 @@ def test_device_detection_from_action():
def test_device_detection_preserves_dtype():
"""Test that device detection doesn't affect dtype of tokenized tensors."""
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)
observation = {"observation.state": torch.randn(10, dtype=torch.float16)}
@@ -953,15 +959,15 @@ def test_device_detection_preserves_dtype():
@require_package("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_integration_with_device_processor(mock_auto_tokenizer):
"""Test that TokenizerProcessor works correctly with DeviceProcessor in pipeline."""
from lerobot.processor import DeviceProcessor
"""Test that TokenizerProcessorStep works correctly with DeviceProcessorStep in pipeline."""
from lerobot.processor import DeviceProcessorStep
mock_tokenizer = MockTokenizer(vocab_size=100)
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
# Create pipeline with TokenizerProcessor then DeviceProcessor
tokenizer_processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=6)
device_processor = DeviceProcessor(device="cuda:0")
# Create pipeline with TokenizerProcessorStep then DeviceProcessorStep
tokenizer_processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=6)
device_processor = DeviceProcessorStep(device="cuda:0")
robot_processor = DataProcessorPipeline(
[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)
# 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.ACTION].device.type == "cuda"
@@ -991,7 +997,7 @@ def test_integration_with_device_processor(mock_auto_tokenizer):
def test_simulated_accelerate_scenario():
"""Test scenario simulating Accelerate with data already on GPU."""
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
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.processor_vqbet import make_vqbet_pre_post_processors
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DataProcessorPipeline,
DeviceProcessor,
NormalizerProcessor,
RenameProcessor,
ToBatchProcessor,
DeviceProcessorStep,
NormalizerProcessorStep,
RenameProcessorStep,
TransitionKey,
UnnormalizerProcessor,
UnnormalizerProcessorStep,
)
@@ -94,15 +94,15 @@ def test_make_vqbet_processor_basic():
# Check steps in preprocessor
assert len(preprocessor.steps) == 4
assert isinstance(preprocessor.steps[0], RenameProcessor)
assert isinstance(preprocessor.steps[1], NormalizerProcessor)
assert isinstance(preprocessor.steps[2], ToBatchProcessor)
assert isinstance(preprocessor.steps[3], DeviceProcessor)
assert isinstance(preprocessor.steps[0], RenameProcessorStep)
assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep)
assert isinstance(preprocessor.steps[3], DeviceProcessorStep)
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], DeviceProcessor)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessor)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
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},
)
# Replace DeviceProcessor with one that uses float16
# Replace DeviceProcessorStep with one that uses float16
for i, step in enumerate(preprocessor.steps):
if isinstance(step, DeviceProcessor):
preprocessor.steps[i] = DeviceProcessor(device=config.device, float_dtype="float16")
if isinstance(step, DeviceProcessorStep):
preprocessor.steps[i] = DeviceProcessorStep(device=config.device, float_dtype="float16")
# Create test data
observation = {