test(processor): all processors use now the same create_transition (#1906)

* test(processor): all processors use now the same create_transition

* test(processor): use identity instead of lambda for transition in pipelines

tests/processor/test_vqbet_processor.py

@@ -33,19 +33,7 @@ from lerobot.processor import (
     TransitionKey,
     UnnormalizerProcessorStep,
 )
-
-
-def create_transition(observation=None, action=None, **kwargs):
-    """Helper function to create a transition dictionary."""
-    transition = {}
-    if observation is not None:
-        transition[TransitionKey.OBSERVATION] = observation
-    if action is not None:
-        transition[TransitionKey.ACTION] = action
-    for key, value in kwargs.items():
-        if hasattr(TransitionKey, key.upper()):
-            transition[getattr(TransitionKey, key.upper())] = value
-    return transition
+from lerobot.processor.converters import create_transition, identity_transition
 
 
 def create_default_config():
@@ -84,8 +72,8 @@ def test_make_vqbet_processor_basic():
     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},
+        preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
+        postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
     )
 
     # Check processor names
@@ -113,8 +101,8 @@ def test_vqbet_processor_with_images():
     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},
+        preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
+        postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
     )
 
     # Create test data with images and states
@@ -144,8 +132,8 @@ def test_vqbet_processor_cuda():
     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},
+        preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
+        postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
     )
 
     # Create CPU data
@@ -182,8 +170,8 @@ def test_vqbet_processor_accelerate_scenario():
     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},
+        preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
+        postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
     )
 
     # Simulate Accelerate: data already on GPU and batched
@@ -214,8 +202,8 @@ def test_vqbet_processor_multi_gpu():
     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},
+        preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
+        postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
     )
 
     # Simulate data on different GPU
@@ -247,14 +235,14 @@ def test_vqbet_processor_without_stats():
     preprocessor = DataProcessorPipeline(
         factory_preprocessor.steps,
         name=factory_preprocessor.name,
-        to_transition=lambda x: x,
-        to_output=lambda x: x,
+        to_transition=identity_transition,
+        to_output=identity_transition,
     )
     postprocessor = DataProcessorPipeline(
         factory_postprocessor.steps,
         name=factory_postprocessor.name,
-        to_transition=lambda x: x,
-        to_output=lambda x: x,
+        to_transition=identity_transition,
+        to_output=identity_transition,
     )
 
     # Should still create processors
@@ -281,8 +269,8 @@ def test_vqbet_processor_save_and_load():
     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},
+        preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
+        postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
     )
 
     with tempfile.TemporaryDirectory() as tmpdir:
@@ -291,7 +279,7 @@ def test_vqbet_processor_save_and_load():
 
         # Load preprocessor
         loaded_preprocessor = DataProcessorPipeline.from_pretrained(
-            tmpdir, to_transition=lambda x: x, to_output=lambda x: x
+            tmpdir, to_transition=identity_transition, to_output=identity_transition
         )
 
         # Test that loaded processor works
@@ -319,8 +307,8 @@ def test_vqbet_processor_mixed_precision():
     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},
+        preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
+        postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
     )
 
     # Replace DeviceProcessorStep with one that uses float16
@@ -368,8 +356,8 @@ def test_vqbet_processor_large_batch():
     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},
+        preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
+        postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
     )
 
     # Test with large batch
@@ -398,8 +386,8 @@ def test_vqbet_processor_sequential_processing():
     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},
+        preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
+        postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
     )
 
     # Process multiple samples sequentially
@@ -432,8 +420,8 @@ def test_vqbet_processor_bfloat16_device_float32_normalizer():
     preprocessor, _ = 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},
+        preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
+        postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
     )
 
     # Modify the pipeline to use bfloat16 device processor with float32 normalizer