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.

tests/processor/test_act_processor.py

@@ -102,7 +102,12 @@ def test_act_processor_normalization():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_act_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Create test data
     observation = {OBS_STATE: torch.randn(7)}
@@ -131,7 +136,12 @@ def test_act_processor_cuda():
     config.device = "cuda"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_act_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Create CPU data
     observation = {OBS_STATE: torch.randn(7)}
@@ -160,7 +170,12 @@ def test_act_processor_accelerate_scenario():
     config.device = "cuda:0"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_act_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Simulate Accelerate: data already on GPU
     device = torch.device("cuda:0")
@@ -223,14 +238,21 @@ def test_act_processor_save_and_load():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_act_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     with tempfile.TemporaryDirectory() as tmpdir:
         # Save preprocessor
         preprocessor.save_pretrained(tmpdir)
 
         # Load preprocessor
-        loaded_preprocessor = RobotProcessor.from_pretrained(tmpdir)
+        loaded_preprocessor = RobotProcessor.from_pretrained(
+            tmpdir, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         # Test that loaded processor works
         observation = {OBS_STATE: torch.randn(7)}
@@ -249,7 +271,12 @@ def test_act_processor_device_placement_preservation():
 
     # Test with CPU config
     config.device = "cpu"
-    preprocessor, _ = make_act_pre_post_processors(config, stats)
+    preprocessor, _ = make_act_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Process CPU data
     observation = {OBS_STATE: torch.randn(7)}
@@ -269,12 +296,21 @@ def test_act_processor_mixed_precision():
     stats = create_default_stats()
 
     # Modify the device processor to use float16
-    preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_act_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Replace DeviceProcessor with one that uses float16
-    for i, step in enumerate(preprocessor.steps):
+    modified_steps = []
+    for step in preprocessor.steps:
         if isinstance(step, DeviceProcessor):
-            preprocessor.steps[i] = DeviceProcessor(device=config.device, float_dtype="float16")
+            modified_steps.append(DeviceProcessor(device=config.device, float_dtype="float16"))
+        else:
+            modified_steps.append(step)
+    preprocessor.steps = modified_steps
 
     # Create test data
     observation = {OBS_STATE: torch.randn(7, dtype=torch.float32)}
@@ -294,7 +330,12 @@ def test_act_processor_batch_consistency():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_act_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Test single sample (unbatched)
     observation = {OBS_STATE: torch.randn(7)}

tests/processor/test_batch_processor.py

@@ -245,7 +245,7 @@ def test_mixed_observation():
 def test_integration_with_robot_processor():
     """Test ToBatchProcessor integration with RobotProcessor."""
     to_batch_processor = ToBatchProcessor()
-    pipeline = RobotProcessor([to_batch_processor])
+    pipeline = RobotProcessor([to_batch_processor], to_transition=lambda x: x, to_output=lambda x: x)
 
     # Create unbatched observation
     observation = {
@@ -285,7 +285,9 @@ def test_serialization_methods():
 def test_save_and_load_pretrained():
     """Test saving and loading ToBatchProcessor with RobotProcessor."""
     processor = ToBatchProcessor()
-    pipeline = RobotProcessor([processor], name="BatchPipeline")
+    pipeline = RobotProcessor(
+        [processor], name="BatchPipeline", to_transition=lambda x: x, to_output=lambda x: x
+    )
 
     with tempfile.TemporaryDirectory() as tmp_dir:
         # Save pipeline
@@ -296,7 +298,9 @@ def test_save_and_load_pretrained():
         assert config_path.exists()
 
         # Load pipeline
-        loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir)
+        loaded_pipeline = RobotProcessor.from_pretrained(
+            tmp_dir, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         assert loaded_pipeline.name == "BatchPipeline"
         assert len(loaded_pipeline) == 1
@@ -323,11 +327,13 @@ def test_registry_functionality():
 def test_registry_based_save_load():
     """Test saving and loading using registry name."""
     processor = ToBatchProcessor()
-    pipeline = RobotProcessor([processor])
+    pipeline = RobotProcessor([processor], to_transition=lambda x: x, to_output=lambda x: x)
 
     with tempfile.TemporaryDirectory() as tmp_dir:
         pipeline.save_pretrained(tmp_dir)
-        loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir)
+        loaded_pipeline = RobotProcessor.from_pretrained(
+            tmp_dir, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         # Verify the loaded processor works
         observation = {

tests/processor/test_classifier_processor.py

@@ -97,7 +97,12 @@ def test_classifier_processor_normalization():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_classifier_processor(config, stats)
+    preprocessor, postprocessor = make_classifier_processor(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Create test data
     observation = {
@@ -123,7 +128,12 @@ def test_classifier_processor_cuda():
     config.device = "cuda"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_classifier_processor(config, stats)
+    preprocessor, postprocessor = make_classifier_processor(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Create CPU data
     observation = {
@@ -156,7 +166,12 @@ def test_classifier_processor_accelerate_scenario():
     config.device = "cuda:0"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_classifier_processor(config, stats)
+    preprocessor, postprocessor = make_classifier_processor(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Simulate Accelerate: data already on GPU
     device = torch.device("cuda:0")
@@ -230,14 +245,22 @@ def test_classifier_processor_save_and_load():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_classifier_processor(config, stats)
+    # Get the steps from the factory function
+    factory_preprocessor, factory_postprocessor = make_classifier_processor(config, stats)
+
+    # Create new processors with EnvTransition input/output
+    preprocessor = RobotProcessor(
+        factory_preprocessor.steps, to_transition=lambda x: x, to_output=lambda x: x
+    )
 
     with tempfile.TemporaryDirectory() as tmpdir:
         # Save preprocessor
         preprocessor.save_pretrained(tmpdir)
 
         # Load preprocessor
-        loaded_preprocessor = RobotProcessor.from_pretrained(tmpdir)
+        loaded_preprocessor = RobotProcessor.from_pretrained(
+            tmpdir, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         # Test that loaded processor works
         observation = {
@@ -260,13 +283,19 @@ def test_classifier_processor_mixed_precision():
     config.device = "cuda"
     stats = create_default_stats()
 
-    # Create processor
-    preprocessor, postprocessor = make_classifier_processor(config, stats)
+    # Get the steps from the factory function
+    factory_preprocessor, factory_postprocessor = make_classifier_processor(config, stats)
 
     # Replace DeviceProcessor with one that uses float16
-    for i, step in enumerate(preprocessor.steps):
+    modified_steps = []
+    for step in factory_preprocessor.steps:
         if isinstance(step, DeviceProcessor):
-            preprocessor.steps[i] = DeviceProcessor(device=config.device, float_dtype="float16")
+            modified_steps.append(DeviceProcessor(device=config.device, float_dtype="float16"))
+        else:
+            modified_steps.append(step)
+
+    # Create new processors with EnvTransition input/output
+    preprocessor = RobotProcessor(modified_steps, to_transition=lambda x: x, to_output=lambda x: x)
 
     # Create test data
     observation = {
@@ -290,7 +319,12 @@ def test_classifier_processor_batch_data():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_classifier_processor(config, stats)
+    preprocessor, postprocessor = make_classifier_processor(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Test with batched data
     batch_size = 16
@@ -315,7 +349,12 @@ def test_classifier_processor_postprocessor_identity():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_classifier_processor(config, stats)
+    preprocessor, postprocessor = make_classifier_processor(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Create test data for postprocessor
     reward = torch.tensor([[0.8], [0.3], [0.9]])  # Batch of rewards/predictions

tests/processor/test_device_processor.py

@@ -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

tests/processor/test_diffusion_processor.py

@@ -105,7 +105,12 @@ def test_diffusion_processor_with_images():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_diffusion_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Create test data with images
     observation = {
@@ -131,7 +136,12 @@ def test_diffusion_processor_cuda():
     config.device = "cuda"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_diffusion_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Create CPU data
     observation = {
@@ -164,7 +174,12 @@ def test_diffusion_processor_accelerate_scenario():
     config.device = "cuda:0"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_diffusion_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Simulate Accelerate: data already on GPU
     device = torch.device("cuda:0")
@@ -238,14 +253,22 @@ def test_diffusion_processor_save_and_load():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
+    # Get the steps from the factory function
+    factory_preprocessor, factory_postprocessor = make_diffusion_pre_post_processors(config, stats)
+
+    # Create new processors with EnvTransition input/output
+    preprocessor = RobotProcessor(
+        factory_preprocessor.steps, to_transition=lambda x: x, to_output=lambda x: x
+    )
 
     with tempfile.TemporaryDirectory() as tmpdir:
         # Save preprocessor
         preprocessor.save_pretrained(tmpdir)
 
         # Load preprocessor
-        loaded_preprocessor = RobotProcessor.from_pretrained(tmpdir)
+        loaded_preprocessor = RobotProcessor.from_pretrained(
+            tmpdir, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         # Test that loaded processor works
         observation = {
@@ -268,13 +291,19 @@ def test_diffusion_processor_mixed_precision():
     config.device = "cuda"
     stats = create_default_stats()
 
-    # Create processor
-    preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
+    # 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
-    for i, step in enumerate(preprocessor.steps):
+    modified_steps = []
+    for step in factory_preprocessor.steps:
         if isinstance(step, DeviceProcessor):
-            preprocessor.steps[i] = DeviceProcessor(device=config.device, float_dtype="float16")
+            modified_steps.append(DeviceProcessor(device=config.device, float_dtype="float16"))
+        else:
+            modified_steps.append(step)
+
+    # Create new processors with EnvTransition input/output
+    preprocessor = RobotProcessor(modified_steps, to_transition=lambda x: x, to_output=lambda x: x)
 
     # Create test data
     observation = {
@@ -298,7 +327,12 @@ def test_diffusion_processor_identity_normalization():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_diffusion_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Create test data
     image_value = torch.rand(3, 224, 224) * 255  # Large values
@@ -322,7 +356,12 @@ def test_diffusion_processor_batch_consistency():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_diffusion_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Test with different batch sizes
     for batch_size in [1, 8, 32]:

tests/processor/test_normalize_processor.py

@@ -506,7 +506,7 @@ def test_get_config(full_stats):
 
 def test_integration_with_robot_processor(normalizer_processor):
     """Test integration with RobotProcessor pipeline"""
-    robot_processor = RobotProcessor([normalizer_processor])
+    robot_processor = RobotProcessor([normalizer_processor], to_transition=lambda x: x, to_output=lambda x: x)
 
     observation = {
         "observation.image": torch.tensor([0.7, 0.5, 0.3]),
@@ -1317,7 +1317,7 @@ def test_hotswap_stats_functional_test():
 
     # Create original processor
     normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats)
-    original_processor = RobotProcessor(steps=[normalizer])
+    original_processor = RobotProcessor(steps=[normalizer], to_transition=lambda x: x, to_output=lambda x: x)
 
     # Process with original stats
     original_result = original_processor(transition)

tests/processor/test_pi0_processor.py

@@ -84,7 +84,12 @@ def test_make_pi0_processor_basic():
     stats = create_default_stats()
 
     with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor"):
-        preprocessor, postprocessor = make_pi0_pre_post_processors(config, stats)
+        preprocessor, postprocessor = make_pi0_pre_post_processors(
+            config,
+            stats,
+            preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+            postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        )
 
     # Check processor names
     assert preprocessor.name == "robot_preprocessor"
@@ -183,7 +188,12 @@ def test_pi0_processor_cuda():
             return features
 
     with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor):
-        preprocessor, postprocessor = make_pi0_pre_post_processors(config, stats)
+        preprocessor, postprocessor = make_pi0_pre_post_processors(
+            config,
+            stats,
+            preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+            postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        )
 
     # Create CPU data
     observation = {
@@ -233,7 +243,12 @@ def test_pi0_processor_accelerate_scenario():
             return features
 
     with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor):
-        preprocessor, postprocessor = make_pi0_pre_post_processors(config, stats)
+        preprocessor, postprocessor = make_pi0_pre_post_processors(
+            config,
+            stats,
+            preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+            postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        )
 
     # Simulate Accelerate: data already on GPU and batched
     device = torch.device("cuda:0")
@@ -284,7 +299,12 @@ def test_pi0_processor_multi_gpu():
             return features
 
     with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor):
-        preprocessor, postprocessor = make_pi0_pre_post_processors(config, stats)
+        preprocessor, postprocessor = make_pi0_pre_post_processors(
+            config,
+            stats,
+            preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+            postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        )
 
     # Simulate data on different GPU
     device = torch.device("cuda:1")
@@ -310,7 +330,12 @@ def test_pi0_processor_without_stats():
 
     # Mock the tokenizer processor
     with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor"):
-        preprocessor, postprocessor = make_pi0_pre_post_processors(config, dataset_stats=None)
+        preprocessor, postprocessor = make_pi0_pre_post_processors(
+            config,
+            dataset_stats=None,
+            preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+            postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        )
 
     # Should still create processors
     assert preprocessor is not None

tests/processor/test_pipeline.py

@@ -176,7 +176,7 @@ class MockStepWithTensorState:
 
 def test_empty_pipeline():
     """Test pipeline with no steps."""
-    pipeline = RobotProcessor()
+    pipeline = RobotProcessor([], to_transition=lambda x: x, to_output=lambda x: x)
 
     transition = create_transition()
     result = pipeline(transition)
@@ -188,7 +188,7 @@ def test_empty_pipeline():
 def test_single_step_pipeline():
     """Test pipeline with a single step."""
     step = MockStep("test_step")
-    pipeline = RobotProcessor([step])
+    pipeline = RobotProcessor([step], to_transition=lambda x: x, to_output=lambda x: x)
 
     transition = create_transition()
     result = pipeline(transition)
@@ -205,7 +205,7 @@ def test_multiple_steps_pipeline():
     """Test pipeline with multiple steps."""
     step1 = MockStep("step1")
     step2 = MockStep("step2")
-    pipeline = RobotProcessor([step1, step2])
+    pipeline = RobotProcessor([step1, step2], to_transition=lambda x: x, to_output=lambda x: x)
 
     transition = create_transition()
     result = pipeline(transition)
@@ -557,7 +557,9 @@ def test_save_and_load_pretrained():
 def test_step_without_optional_methods():
     """Test pipeline with steps that don't implement optional methods."""
     step = MockStepWithoutOptionalMethods(multiplier=3.0)
-    pipeline = RobotProcessor([step])
+    pipeline = RobotProcessor(
+        [step], to_transition=lambda x: x, to_output=lambda x: x
+    )  # Identity for EnvTransition input/output
 
     transition = create_transition(reward=2.0)
     result = pipeline(transition)
@@ -878,7 +880,9 @@ def test_from_pretrained_with_overrides():
             "registered_mock_step": {"device": "cuda", "value": 200},
         }
 
-        loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides)
+        loaded_pipeline = RobotProcessor.from_pretrained(
+            tmp_dir, overrides=overrides, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         # Verify the pipeline was loaded correctly
         assert len(loaded_pipeline) == 2
@@ -914,7 +918,9 @@ def test_from_pretrained_with_partial_overrides():
 
         # The current implementation applies overrides to ALL steps with the same class name
         # Both steps will get the override
-        loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides)
+        loaded_pipeline = RobotProcessor.from_pretrained(
+            tmp_dir, overrides=overrides, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         transition = create_transition(reward=1.0)
         result = loaded_pipeline(transition)
@@ -971,7 +977,9 @@ def test_from_pretrained_registered_step_override():
         # Override using registry name
         overrides = {"registered_mock_step": {"value": 999, "device": "cuda"}}
 
-        loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides)
+        loaded_pipeline = RobotProcessor.from_pretrained(
+            tmp_dir, overrides=overrides, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         # Test that overrides were applied
         transition = create_transition()
@@ -999,7 +1007,9 @@ def test_from_pretrained_mixed_registered_and_unregistered():
             "registered_mock_step": {"value": 777},
         }
 
-        loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides)
+        loaded_pipeline = RobotProcessor.from_pretrained(
+            tmp_dir, overrides=overrides, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         # Test both steps
         transition = create_transition(reward=2.0)
@@ -1020,7 +1030,9 @@ def test_from_pretrained_no_overrides():
         pipeline.save_pretrained(tmp_dir)
 
         # Load without overrides
-        loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir)
+        loaded_pipeline = RobotProcessor.from_pretrained(
+            tmp_dir, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         assert len(loaded_pipeline) == 1
 
@@ -1040,7 +1052,9 @@ def test_from_pretrained_empty_overrides():
         pipeline.save_pretrained(tmp_dir)
 
         # Load with empty overrides
-        loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides={})
+        loaded_pipeline = RobotProcessor.from_pretrained(
+            tmp_dir, overrides={}, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         assert len(loaded_pipeline) == 1
 
@@ -1455,6 +1469,8 @@ def test_override_with_nested_config():
             loaded = RobotProcessor.from_pretrained(
                 tmp_dir,
                 overrides={"complex_config_step": {"nested_config": {"level1": {"level2": "overridden"}}}},
+                to_transition=lambda x: x,
+                to_output=lambda x: x,
             )
 
             # Test that override worked
@@ -1553,7 +1569,10 @@ def test_override_with_callables():
 
             # Load with callable override
             loaded = RobotProcessor.from_pretrained(
-                tmp_dir, overrides={"callable_step": {"transform_fn": double_values}}
+                tmp_dir,
+                overrides={"callable_step": {"transform_fn": double_values}},
+                to_transition=lambda x: x,
+                to_output=lambda x: x,
             )
 
             # Test it works
@@ -1857,7 +1876,8 @@ def test_save_load_with_custom_converter_functions():
 
         # Should work with standard format (wouldn't work with custom converter)
         result = loaded(batch)
-        assert "observation.image" in result  # Standard format preserved
+        # With new behavior, default to_output is _default_transition_to_batch, so result is batch dict
+        assert "observation.image" in result
 
 
 class NonCompliantStep:

tests/processor/test_rename_processor.py

@@ -188,7 +188,7 @@ def test_integration_with_robot_processor():
     }
     rename_processor = RenameProcessor(rename_map=rename_map)
 
-    pipeline = RobotProcessor([rename_processor])
+    pipeline = RobotProcessor([rename_processor], to_transition=lambda x: x, to_output=lambda x: x)
 
     observation = {
         "agent_pos": np.array([1.0, 2.0, 3.0]),
@@ -236,7 +236,9 @@ def test_save_and_load_pretrained():
         assert len(state_files) == 0
 
         # Load pipeline
-        loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir)
+        loaded_pipeline = RobotProcessor.from_pretrained(
+            tmp_dir, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         assert loaded_pipeline.name == "TestRenameProcessor"
         assert len(loaded_pipeline) == 1
@@ -277,7 +279,7 @@ def test_registry_functionality():
 def test_registry_based_save_load():
     """Test save/load using registry name instead of module path."""
     processor = RenameProcessor(rename_map={"key1": "renamed_key1"})
-    pipeline = RobotProcessor([processor])
+    pipeline = RobotProcessor([processor], to_transition=lambda x: x, to_output=lambda x: x)
 
     with tempfile.TemporaryDirectory() as tmp_dir:
         # Save and load
@@ -318,7 +320,7 @@ def test_chained_rename_processors():
         }
     )
 
-    pipeline = RobotProcessor([processor1, processor2])
+    pipeline = RobotProcessor([processor1, processor2], to_transition=lambda x: x, to_output=lambda x: x)
 
     observation = {
         "pos": np.array([1.0, 2.0]),

tests/processor/test_sac_processor.py

@@ -78,7 +78,12 @@ def test_make_sac_processor_basic():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_sac_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Check processor names
     assert preprocessor.name == "robot_preprocessor"
@@ -102,7 +107,12 @@ def test_sac_processor_normalization_modes():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_sac_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Create test data
     observation = {OBS_STATE: torch.randn(10) * 2}  # Larger values to test normalization
@@ -133,7 +143,12 @@ def test_sac_processor_cuda():
     config.device = "cuda"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_sac_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Create CPU data
     observation = {OBS_STATE: torch.randn(10)}
@@ -162,7 +177,12 @@ def test_sac_processor_accelerate_scenario():
     config.device = "cuda:0"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_sac_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Simulate Accelerate: data already on GPU
     device = torch.device("cuda:0")
@@ -185,7 +205,12 @@ def test_sac_processor_multi_gpu():
     config.device = "cuda:0"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_sac_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Simulate data on different GPU
     device = torch.device("cuda:1")
@@ -205,7 +230,22 @@ def test_sac_processor_without_stats():
     """Test SAC processor creation without dataset statistics."""
     config = create_default_config()
 
-    preprocessor, postprocessor = make_sac_pre_post_processors(config, dataset_stats=None)
+    # Get the steps from the factory function
+    factory_preprocessor, factory_postprocessor = make_sac_pre_post_processors(config, dataset_stats=None)
+
+    # Create new processors with EnvTransition input/output
+    preprocessor = RobotProcessor(
+        factory_preprocessor.steps,
+        name=factory_preprocessor.name,
+        to_transition=lambda x: x,
+        to_output=lambda x: x,
+    )
+    postprocessor = RobotProcessor(
+        factory_postprocessor.steps,
+        name=factory_postprocessor.name,
+        to_transition=lambda x: x,
+        to_output=lambda x: x,
+    )
 
     # Should still create processors
     assert preprocessor is not None
@@ -225,14 +265,21 @@ def test_sac_processor_save_and_load():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_sac_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     with tempfile.TemporaryDirectory() as tmpdir:
         # Save preprocessor
         preprocessor.save_pretrained(tmpdir)
 
         # Load preprocessor
-        loaded_preprocessor = RobotProcessor.from_pretrained(tmpdir)
+        loaded_preprocessor = RobotProcessor.from_pretrained(
+            tmpdir, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         # Test that loaded processor works
         observation = {OBS_STATE: torch.randn(10)}
@@ -252,7 +299,12 @@ def test_sac_processor_mixed_precision():
     stats = create_default_stats()
 
     # Create processor
-    preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_sac_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Replace DeviceProcessor with one that uses float16
     for i, step in enumerate(preprocessor.steps):
@@ -277,7 +329,12 @@ def test_sac_processor_batch_data():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_sac_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Test with batched data
     batch_size = 32
@@ -298,7 +355,12 @@ def test_sac_processor_edge_cases():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_sac_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Test with empty observation
     transition = create_transition(observation={}, action=torch.randn(5))

tests/processor/test_smolvla_processor.py

@@ -89,7 +89,12 @@ def test_make_smolvla_processor_basic():
     stats = create_default_stats()
 
     with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor"):
-        preprocessor, postprocessor = make_smolvla_pre_post_processors(config, stats)
+        preprocessor, postprocessor = make_smolvla_pre_post_processors(
+            config,
+            stats,
+            preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+            postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        )
 
     # Check processor names
     assert preprocessor.name == "robot_preprocessor"
@@ -188,7 +193,12 @@ def test_smolvla_processor_cuda():
             return features
 
     with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor):
-        preprocessor, postprocessor = make_smolvla_pre_post_processors(config, stats)
+        preprocessor, postprocessor = make_smolvla_pre_post_processors(
+            config,
+            stats,
+            preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+            postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        )
 
     # Create CPU data
     observation = {
@@ -238,7 +248,12 @@ def test_smolvla_processor_accelerate_scenario():
             return features
 
     with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor):
-        preprocessor, postprocessor = make_smolvla_pre_post_processors(config, stats)
+        preprocessor, postprocessor = make_smolvla_pre_post_processors(
+            config,
+            stats,
+            preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+            postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        )
 
     # Simulate Accelerate: data already on GPU and batched
     device = torch.device("cuda:0")
@@ -289,7 +304,12 @@ def test_smolvla_processor_multi_gpu():
             return features
 
     with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor):
-        preprocessor, postprocessor = make_smolvla_pre_post_processors(config, stats)
+        preprocessor, postprocessor = make_smolvla_pre_post_processors(
+            config,
+            stats,
+            preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+            postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        )
 
     # Simulate data on different GPU
     device = torch.device("cuda:1")
@@ -315,7 +335,12 @@ def test_smolvla_processor_without_stats():
 
     # Mock the tokenizer processor
     with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor"):
-        preprocessor, postprocessor = make_smolvla_pre_post_processors(config, dataset_stats=None)
+        preprocessor, postprocessor = make_smolvla_pre_post_processors(
+            config,
+            dataset_stats=None,
+            preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+            postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        )
 
     # Should still create processors
     assert preprocessor is not None

tests/processor/test_tdmpc_processor.py

@@ -81,7 +81,12 @@ def test_make_tdmpc_processor_basic():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Check processor names
     assert preprocessor.name == "robot_preprocessor"
@@ -105,7 +110,12 @@ def test_tdmpc_processor_normalization():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Create test data
     observation = {
@@ -138,7 +148,12 @@ def test_tdmpc_processor_cuda():
     config.device = "cuda"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Create CPU data
     observation = {
@@ -171,7 +186,12 @@ def test_tdmpc_processor_accelerate_scenario():
     config.device = "cuda:0"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Simulate Accelerate: data already on GPU
     device = torch.device("cuda:0")
@@ -198,7 +218,12 @@ def test_tdmpc_processor_multi_gpu():
     config.device = "cuda:0"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Simulate data on different GPU
     device = torch.device("cuda:1")
@@ -222,7 +247,22 @@ def test_tdmpc_processor_without_stats():
     """Test TDMPC processor creation without dataset statistics."""
     config = create_default_config()
 
-    preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, dataset_stats=None)
+    # Get the steps from the factory function
+    factory_preprocessor, factory_postprocessor = make_tdmpc_pre_post_processors(config, dataset_stats=None)
+
+    # Create new processors with EnvTransition input/output
+    preprocessor = RobotProcessor(
+        factory_preprocessor.steps,
+        name=factory_preprocessor.name,
+        to_transition=lambda x: x,
+        to_output=lambda x: x,
+    )
+    postprocessor = RobotProcessor(
+        factory_postprocessor.steps,
+        name=factory_postprocessor.name,
+        to_transition=lambda x: x,
+        to_output=lambda x: x,
+    )
 
     # Should still create processors
     assert preprocessor is not None
@@ -245,14 +285,21 @@ def test_tdmpc_processor_save_and_load():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     with tempfile.TemporaryDirectory() as tmpdir:
         # Save preprocessor
         preprocessor.save_pretrained(tmpdir)
 
         # Load preprocessor
-        loaded_preprocessor = RobotProcessor.from_pretrained(tmpdir)
+        loaded_preprocessor = RobotProcessor.from_pretrained(
+            tmpdir, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         # Test that loaded processor works
         observation = {
@@ -276,7 +323,12 @@ def test_tdmpc_processor_mixed_precision():
     stats = create_default_stats()
 
     # Create processor
-    preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Replace DeviceProcessor with one that uses float16
     for i, step in enumerate(preprocessor.steps):
@@ -305,7 +357,12 @@ def test_tdmpc_processor_batch_data():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Test with batched data
     batch_size = 64
@@ -330,7 +387,12 @@ def test_tdmpc_processor_edge_cases():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Test with only state observation (no image)
     observation = {OBS_STATE: torch.randn(12)}

tests/processor/test_tokenizer_processor.py

@@ -389,7 +389,7 @@ def test_integration_with_robot_processor(mock_auto_tokenizer):
     mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
 
     tokenizer_processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=6)
-    robot_processor = RobotProcessor([tokenizer_processor])
+    robot_processor = RobotProcessor([tokenizer_processor], to_transition=lambda x: x, to_output=lambda x: x)
 
     transition = create_transition(
         observation={"state": torch.tensor([1.0, 2.0])},
@@ -427,14 +427,16 @@ def test_save_and_load_pretrained_with_tokenizer_name(mock_auto_tokenizer):
         tokenizer_name="test-tokenizer", max_length=32, task_key="instruction"
     )
 
-    robot_processor = RobotProcessor([original_processor])
+    robot_processor = RobotProcessor([original_processor], to_transition=lambda x: x, to_output=lambda x: x)
 
     with tempfile.TemporaryDirectory() as temp_dir:
         # Save processor
         robot_processor.save_pretrained(temp_dir)
 
         # Load processor - tokenizer will be recreated from saved config
-        loaded_processor = RobotProcessor.from_pretrained(temp_dir)
+        loaded_processor = RobotProcessor.from_pretrained(
+            temp_dir, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         # Test that loaded processor works
         transition = create_transition(
@@ -456,7 +458,7 @@ def test_save_and_load_pretrained_with_tokenizer_object():
 
     original_processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=32, task_key="instruction")
 
-    robot_processor = RobotProcessor([original_processor])
+    robot_processor = RobotProcessor([original_processor], to_transition=lambda x: x, to_output=lambda x: x)
 
     with tempfile.TemporaryDirectory() as temp_dir:
         # Save processor
@@ -464,7 +466,10 @@ def test_save_and_load_pretrained_with_tokenizer_object():
 
         # Load processor with tokenizer override (since tokenizer object wasn't saved)
         loaded_processor = RobotProcessor.from_pretrained(
-            temp_dir, overrides={"tokenizer_processor": {"tokenizer": mock_tokenizer}}
+            temp_dir,
+            overrides={"tokenizer_processor": {"tokenizer": mock_tokenizer}},
+            to_transition=lambda x: x,
+            to_output=lambda x: x,
         )
 
         # Test that loaded processor works
@@ -952,7 +957,9 @@ def test_integration_with_device_processor(mock_auto_tokenizer):
     # Create pipeline with TokenizerProcessor then DeviceProcessor
     tokenizer_processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=6)
     device_processor = DeviceProcessor(device="cuda:0")
-    robot_processor = RobotProcessor([tokenizer_processor, device_processor])
+    robot_processor = RobotProcessor(
+        [tokenizer_processor, device_processor], to_transition=lambda x: x, to_output=lambda x: x
+    )
 
     # Start with CPU tensors
     transition = create_transition(

tests/processor/test_vqbet_processor.py

@@ -81,7 +81,12 @@ def test_make_vqbet_processor_basic():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Check processor names
     assert preprocessor.name == "robot_preprocessor"
@@ -105,7 +110,12 @@ def test_vqbet_processor_with_images():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Create test data with images and states
     observation = {
@@ -131,7 +141,12 @@ def test_vqbet_processor_cuda():
     config.device = "cuda"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Create CPU data
     observation = {
@@ -164,7 +179,12 @@ def test_vqbet_processor_accelerate_scenario():
     config.device = "cuda:0"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Simulate Accelerate: data already on GPU and batched
     device = torch.device("cuda:0")
@@ -191,7 +211,12 @@ def test_vqbet_processor_multi_gpu():
     config.device = "cuda:0"
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Simulate data on different GPU
     device = torch.device("cuda:1")
@@ -215,7 +240,22 @@ def test_vqbet_processor_without_stats():
     """Test VQBeT processor creation without dataset statistics."""
     config = create_default_config()
 
-    preprocessor, postprocessor = make_vqbet_pre_post_processors(config, dataset_stats=None)
+    # Get the steps from the factory function
+    factory_preprocessor, factory_postprocessor = make_vqbet_pre_post_processors(config, dataset_stats=None)
+
+    # Create new processors with EnvTransition input/output
+    preprocessor = RobotProcessor(
+        factory_preprocessor.steps,
+        name=factory_preprocessor.name,
+        to_transition=lambda x: x,
+        to_output=lambda x: x,
+    )
+    postprocessor = RobotProcessor(
+        factory_postprocessor.steps,
+        name=factory_postprocessor.name,
+        to_transition=lambda x: x,
+        to_output=lambda x: x,
+    )
 
     # Should still create processors
     assert preprocessor is not None
@@ -238,14 +278,21 @@ def test_vqbet_processor_save_and_load():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     with tempfile.TemporaryDirectory() as tmpdir:
         # Save preprocessor
         preprocessor.save_pretrained(tmpdir)
 
         # Load preprocessor
-        loaded_preprocessor = RobotProcessor.from_pretrained(tmpdir)
+        loaded_preprocessor = RobotProcessor.from_pretrained(
+            tmpdir, to_transition=lambda x: x, to_output=lambda x: x
+        )
 
         # Test that loaded processor works
         observation = {
@@ -269,7 +316,12 @@ def test_vqbet_processor_mixed_precision():
     stats = create_default_stats()
 
     # Create processor
-    preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Replace DeviceProcessor with one that uses float16
     for i, step in enumerate(preprocessor.steps):
@@ -298,7 +350,12 @@ def test_vqbet_processor_large_batch():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Test with large batch
     batch_size = 128
@@ -323,7 +380,12 @@ def test_vqbet_processor_sequential_processing():
     config = create_default_config()
     stats = create_default_stats()
 
-    preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+        preprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+        postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x},
+    )
 
     # Process multiple samples sequentially
     results = []