mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-15 08:39:49 +00:00
refactor(processor): improve processor pipeline typing with generic type (#1810)
* refactor(processor): introduce generic type for to_output - Always return `TOutput` - Remove `_prepare_transition`, so `__call__` now always returns `TOutput` - Update tests accordingly - This refactor paves the way for adding settings for `to_transition` and `to_output` in `make_processor` and the post-processor * refactor(processor): consolidate ProcessorKwargs usage across policies - Removed the ProcessorTypes module and integrated ProcessorKwargs directly into the processor pipeline. - Updated multiple policy files to utilize the new ProcessorKwargs structure for preprocessor and postprocessor arguments. - Simplified the handling of processor kwargs by initializing them to empty dictionaries when not provided.
This commit is contained in:
@@ -311,7 +311,12 @@ def test_integration_with_robot_processor():
|
||||
device_processor = DeviceProcessor(device="cpu")
|
||||
batch_processor = ToBatchProcessor()
|
||||
|
||||
processor = RobotProcessor(steps=[batch_processor, device_processor], name="test_pipeline")
|
||||
processor = RobotProcessor(
|
||||
steps=[batch_processor, device_processor],
|
||||
name="test_pipeline",
|
||||
to_transition=lambda x: x,
|
||||
to_output=lambda x: x,
|
||||
)
|
||||
|
||||
# Create test data
|
||||
observation = {OBS_STATE: torch.randn(10)}
|
||||
@@ -985,6 +990,8 @@ def test_policy_processor_integration():
|
||||
DeviceProcessor(device="cuda"),
|
||||
],
|
||||
name="test_preprocessor",
|
||||
to_transition=lambda x: x,
|
||||
to_output=lambda x: x,
|
||||
)
|
||||
|
||||
# Create output processor (postprocessor) that moves to CPU
|
||||
@@ -994,6 +1001,8 @@ def test_policy_processor_integration():
|
||||
UnnormalizerProcessor(features={ACTION: features[ACTION]}, norm_map=norm_map, stats=stats),
|
||||
],
|
||||
name="test_postprocessor",
|
||||
to_transition=lambda x: x,
|
||||
to_output=lambda x: x,
|
||||
)
|
||||
|
||||
# Test data on CPU
|
||||
|
||||
Reference in New Issue
Block a user