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:
Adil Zouitine
2025-09-02 12:57:14 +02:00
committed by GitHub
parent 08fb310eaa
commit d32b76cc66
26 changed files with 847 additions and 220 deletions
+10 -1
View File
@@ -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