chore: import upstream snapshot with attribution

tests/processor/test_act_processor.py

@@ -0,0 +1,412 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for ACT policy processor."""
+
+import tempfile
+
+import pytest
+import torch
+
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+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,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    RenameObservationsProcessorStep,
+    TransitionKey,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import create_transition, transition_to_batch
+from lerobot.utils.constants import ACTION, OBS_STATE
+
+
+def create_default_config():
+    """Create a default ACT configuration for testing."""
+    config = ACTConfig()
+    config.input_features = {
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(7,)),
+    }
+    config.output_features = {
+        ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(4,)),
+    }
+    config.normalization_mapping = {
+        FeatureType.STATE: NormalizationMode.MEAN_STD,
+        FeatureType.ACTION: NormalizationMode.MEAN_STD,
+    }
+    config.device = "cpu"
+    return config
+
+
+def create_default_stats():
+    """Create default dataset statistics for testing."""
+    return {
+        OBS_STATE: {"mean": torch.zeros(7), "std": torch.ones(7)},
+        ACTION: {"mean": torch.zeros(4), "std": torch.ones(4)},
+    }
+
+
+def test_make_act_processor_basic():
+    """Test basic creation of ACT processor."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
+
+    # Check processor names
+    assert preprocessor.name == "policy_preprocessor"
+    assert postprocessor.name == "policy_postprocessor"
+
+    # Check steps in preprocessor
+    assert len(preprocessor.steps) == 4
+    assert isinstance(preprocessor.steps[0], RenameObservationsProcessorStep)
+    assert isinstance(preprocessor.steps[1], AddBatchDimensionProcessorStep)
+    assert isinstance(preprocessor.steps[2], DeviceProcessorStep)
+    assert isinstance(preprocessor.steps[3], NormalizerProcessorStep)
+
+    # Check steps in postprocessor
+    assert len(postprocessor.steps) == 2
+    assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
+    assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
+
+
+def test_act_processor_normalization():
+    """Test that ACT processor correctly normalizes and unnormalizes data."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_act_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Create test data
+    observation = {OBS_STATE: torch.randn(7)}
+    action = torch.randn(4)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+    processed = preprocessor(batch)
+
+    # Check that data is normalized and batched
+    assert processed[OBS_STATE].shape == (1, 7)
+    assert processed[TransitionKey.ACTION.value].shape == (1, 4)
+
+    # Process action through postprocessor
+    postprocessed = postprocessor(processed[TransitionKey.ACTION.value])
+
+    # Check that action is unnormalized
+    assert postprocessed.shape == (1, 4)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_act_processor_cuda():
+    """Test ACT processor with CUDA device."""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_act_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Create CPU data
+    observation = {OBS_STATE: torch.randn(7)}
+    action = torch.randn(4)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+    processed = preprocessor(batch)
+
+    # Check that data is on CUDA
+    assert processed[OBS_STATE].device.type == "cuda"
+    assert processed[TransitionKey.ACTION.value].device.type == "cuda"
+
+    # Process through postprocessor
+    postprocessed = postprocessor(processed[TransitionKey.ACTION.value])
+
+    # Check that action is back on CPU
+    assert postprocessed.device.type == "cpu"
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_act_processor_accelerate_scenario():
+    """Test ACT processor in simulated Accelerate scenario (data already on GPU)."""
+    config = create_default_config()
+    config.device = "cuda:0"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_act_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Simulate Accelerate: data already on GPU
+    device = torch.device("cuda:0")
+    observation = {OBS_STATE: torch.randn(1, 7).to(device)}  # Already batched and on GPU
+    action = torch.randn(1, 4).to(device)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+    processed = preprocessor(batch)
+
+    # Check that data stays on same GPU (not moved unnecessarily)
+    assert processed[OBS_STATE].device == device
+    assert processed[TransitionKey.ACTION.value].device == device
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs")
+def test_act_processor_multi_gpu():
+    """Test ACT processor with multi-GPU setup."""
+    config = create_default_config()
+    config.device = "cuda:0"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_act_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Simulate data on different GPU (like in multi-GPU training)
+    device = torch.device("cuda:1")
+    observation = {OBS_STATE: torch.randn(1, 7).to(device)}
+    action = torch.randn(1, 4).to(device)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+    processed = preprocessor(batch)
+
+    # Check that data stays on cuda:1 (not moved to cuda:0)
+    assert processed[OBS_STATE].device == device
+    assert processed[TransitionKey.ACTION.value].device == device
+
+
+def test_act_processor_without_stats():
+    """Test ACT processor creation without dataset statistics."""
+    config = create_default_config()
+
+    preprocessor, postprocessor = make_act_pre_post_processors(
+        config,
+        dataset_stats=None,
+    )
+
+    # Should still create processors, but normalization won't have stats
+    assert preprocessor is not None
+    assert postprocessor is not None
+
+    # Process should still work (but won't normalize without stats)
+    observation = {OBS_STATE: torch.randn(7)}
+    action = torch.randn(4)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    processed = preprocessor(batch)
+    assert processed is not None
+
+
+def test_act_processor_save_and_load():
+    """Test saving and loading ACT processor."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_act_pre_post_processors(
+        config,
+        stats,
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Save preprocessor
+        preprocessor.save_pretrained(tmpdir)
+
+        # Load preprocessor
+        loaded_preprocessor = DataProcessorPipeline.from_pretrained(
+            tmpdir, config_filename="policy_preprocessor.json"
+        )
+
+        # Test that loaded processor works
+        observation = {OBS_STATE: torch.randn(7)}
+        action = torch.randn(4)
+        transition = create_transition(observation, action)
+        batch = transition_to_batch(transition)
+
+        processed = loaded_preprocessor(batch)
+        assert processed[OBS_STATE].shape == (1, 7)
+        assert processed[TransitionKey.ACTION.value].shape == (1, 4)
+
+
+def test_act_processor_device_placement_preservation():
+    """Test that ACT processor preserves device placement correctly."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    # Test with CPU config
+    config.device = "cpu"
+    preprocessor, _ = make_act_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Process CPU data
+    observation = {OBS_STATE: torch.randn(7)}
+    action = torch.randn(4)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    processed = preprocessor(batch)
+    assert processed[OBS_STATE].device.type == "cpu"
+    assert processed[TransitionKey.ACTION.value].device.type == "cpu"
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_act_processor_mixed_precision():
+    """Test ACT processor with mixed precision (float16)."""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    # Modify the device processor to use float16
+    preprocessor, postprocessor = make_act_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Replace DeviceProcessorStep with one that uses float16
+    modified_steps = []
+    for step in preprocessor.steps:
+        if isinstance(step, DeviceProcessorStep):
+            modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16"))
+        elif isinstance(step, NormalizerProcessorStep):
+            # Update normalizer to use the same device as the device processor
+            norm_step = step  # Now type checker knows this is NormalizerProcessorStep
+            modified_steps.append(
+                NormalizerProcessorStep(
+                    features=norm_step.features,
+                    norm_map=norm_step.norm_map,
+                    stats=norm_step.stats,
+                    device=config.device,
+                    dtype=torch.float16,  # Match the float16 dtype
+                )
+            )
+        else:
+            modified_steps.append(step)
+    preprocessor.steps = modified_steps
+
+    # Create test data
+    observation = {OBS_STATE: torch.randn(7, dtype=torch.float32)}
+    action = torch.randn(4, dtype=torch.float32)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+    processed = preprocessor(batch)
+
+    # Check that data is converted to float16
+    assert processed[OBS_STATE].dtype == torch.float16
+    assert processed[TransitionKey.ACTION.value].dtype == torch.float16
+
+
+def test_act_processor_batch_consistency():
+    """Test that ACT processor handles different batch sizes correctly."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_act_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Test single sample (unbatched)
+    observation = {OBS_STATE: torch.randn(7)}
+    action = torch.randn(4)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    processed = preprocessor(batch)
+    assert processed[OBS_STATE].shape[0] == 1  # Batched
+
+    # Test already batched data
+    observation_batched = {OBS_STATE: torch.randn(8, 7)}  # Batch of 8
+    action_batched = torch.randn(8, 4)
+    transition_batched = create_transition(observation_batched, action_batched)
+    batch_batched = transition_to_batch(transition_batched)
+
+    processed_batched = preprocessor(batch_batched)
+    assert processed_batched[OBS_STATE].shape[0] == 8
+    assert processed_batched[TransitionKey.ACTION.value].shape[0] == 8
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_act_processor_bfloat16_device_float32_normalizer():
+    """Test: DeviceProcessor(bfloat16) + NormalizerProcessor(float32) → output bfloat16 via automatic adaptation"""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    preprocessor, _ = make_act_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Modify the pipeline to use bfloat16 device processor with float32 normalizer
+    modified_steps = []
+    for step in preprocessor.steps:
+        if isinstance(step, DeviceProcessorStep):
+            # Device processor converts to bfloat16
+            modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="bfloat16"))
+        elif isinstance(step, NormalizerProcessorStep):
+            # Normalizer stays configured as float32 (will auto-adapt to bfloat16)
+            norm_step = step  # Now type checker knows this is NormalizerProcessorStep
+            modified_steps.append(
+                NormalizerProcessorStep(
+                    features=norm_step.features,
+                    norm_map=norm_step.norm_map,
+                    stats=norm_step.stats,
+                    device=config.device,
+                    dtype=torch.float32,  # Deliberately configured as float32
+                )
+            )
+        else:
+            modified_steps.append(step)
+    preprocessor.steps = modified_steps
+
+    # Verify initial normalizer configuration
+    normalizer_step = preprocessor.steps[3]  # NormalizerProcessorStep
+    assert normalizer_step.dtype == torch.float32
+
+    # Create test data
+    observation = {OBS_STATE: torch.randn(7, dtype=torch.float32)}  # Start with float32
+    action = torch.randn(4, dtype=torch.float32)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    # Process through full pipeline
+    processed = preprocessor(batch)
+
+    # Verify: DeviceProcessor → bfloat16, NormalizerProcessor adapts → final output is bfloat16
+    assert processed[OBS_STATE].dtype == torch.bfloat16
+    assert processed[TransitionKey.ACTION.value].dtype == torch.bfloat16
+
+    # Verify normalizer automatically adapted its internal state
+    assert normalizer_step.dtype == torch.bfloat16
+    for stat_tensor in normalizer_step._tensor_stats[OBS_STATE].values():
+        assert stat_tensor.dtype == torch.bfloat16

tests/processor/test_batch_conversion.py

@@ -0,0 +1,296 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from lerobot.processor import DataProcessorPipeline
+from lerobot.processor.converters import batch_to_transition, transition_to_batch
+from lerobot.types import TransitionKey
+from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_PREFIX, OBS_STATE, REWARD, TRUNCATED
+
+
+def _dummy_batch():
+    """Create a dummy batch using the new format with observation.* and next.* keys."""
+    return {
+        f"{OBS_IMAGE}.left": torch.randn(1, 3, 128, 128),
+        f"{OBS_IMAGE}.right": torch.randn(1, 3, 128, 128),
+        OBS_STATE: torch.tensor([[0.1, 0.2, 0.3, 0.4]]),
+        ACTION: torch.tensor([[0.5]]),
+        REWARD: 1.0,
+        DONE: False,
+        TRUNCATED: False,
+        "info": {"key": "value"},
+    }
+
+
+def test_observation_grouping_roundtrip():
+    """Test that observation.* keys are properly grouped and ungrouped."""
+    proc = DataProcessorPipeline([])
+    batch_in = _dummy_batch()
+    batch_out = proc(batch_in)
+
+    # Check that all observation.* keys are preserved
+    original_obs_keys = {k: v for k, v in batch_in.items() if k.startswith(OBS_PREFIX)}
+    reconstructed_obs_keys = {k: v for k, v in batch_out.items() if k.startswith(OBS_PREFIX)}
+
+    assert set(original_obs_keys.keys()) == set(reconstructed_obs_keys.keys())
+
+    # Check tensor values
+    assert torch.allclose(batch_out[f"{OBS_IMAGE}.left"], batch_in[f"{OBS_IMAGE}.left"])
+    assert torch.allclose(batch_out[f"{OBS_IMAGE}.right"], batch_in[f"{OBS_IMAGE}.right"])
+    assert torch.allclose(batch_out[OBS_STATE], batch_in[OBS_STATE])
+
+    # Check other fields
+    assert torch.allclose(batch_out[ACTION], batch_in[ACTION])
+    assert batch_out[REWARD] == batch_in[REWARD]
+    assert batch_out[DONE] == batch_in[DONE]
+    assert batch_out[TRUNCATED] == batch_in[TRUNCATED]
+    assert batch_out["info"] == batch_in["info"]
+
+
+def test_batch_to_transition_observation_grouping():
+    """Test that batch_to_transition correctly groups observation.* keys."""
+    batch = {
+        f"{OBS_IMAGE}.top": torch.randn(1, 3, 128, 128),
+        f"{OBS_IMAGE}.left": torch.randn(1, 3, 128, 128),
+        OBS_STATE: [1, 2, 3, 4],
+        ACTION: torch.tensor([0.1, 0.2, 0.3, 0.4]),
+        REWARD: 1.5,
+        DONE: True,
+        TRUNCATED: False,
+        "info": {"episode": 42},
+    }
+
+    transition = batch_to_transition(batch)
+
+    # Check observation is a dict with all observation.* keys
+    assert isinstance(transition[TransitionKey.OBSERVATION], dict)
+    assert f"{OBS_IMAGE}.top" in transition[TransitionKey.OBSERVATION]
+    assert f"{OBS_IMAGE}.left" in transition[TransitionKey.OBSERVATION]
+    assert OBS_STATE in transition[TransitionKey.OBSERVATION]
+
+    # Check values are preserved
+    assert torch.allclose(
+        transition[TransitionKey.OBSERVATION][f"{OBS_IMAGE}.top"], batch[f"{OBS_IMAGE}.top"]
+    )
+    assert torch.allclose(
+        transition[TransitionKey.OBSERVATION][f"{OBS_IMAGE}.left"], batch[f"{OBS_IMAGE}.left"]
+    )
+    assert transition[TransitionKey.OBSERVATION][OBS_STATE] == [1, 2, 3, 4]
+
+    # Check other fields
+    assert torch.allclose(transition[TransitionKey.ACTION], torch.tensor([0.1, 0.2, 0.3, 0.4]))
+    assert transition[TransitionKey.REWARD] == 1.5
+    assert transition[TransitionKey.DONE]
+    assert not transition[TransitionKey.TRUNCATED]
+    assert transition[TransitionKey.INFO] == {"episode": 42}
+    assert transition[TransitionKey.COMPLEMENTARY_DATA] == {}
+
+
+def test_transition_to_batch_observation_flattening():
+    """Test that transition_to_batch correctly flattens observation dict."""
+    observation_dict = {
+        f"{OBS_IMAGE}.top": torch.randn(1, 3, 128, 128),
+        f"{OBS_IMAGE}.left": torch.randn(1, 3, 128, 128),
+        OBS_STATE: [1, 2, 3, 4],
+    }
+
+    transition = {
+        TransitionKey.OBSERVATION: observation_dict,
+        TransitionKey.ACTION: "action_data",
+        TransitionKey.REWARD: 1.5,
+        TransitionKey.DONE: True,
+        TransitionKey.TRUNCATED: False,
+        TransitionKey.INFO: {"episode": 42},
+        TransitionKey.COMPLEMENTARY_DATA: {},
+    }
+
+    batch = transition_to_batch(transition)
+
+    # Check that observation.* keys are flattened back to batch
+    assert f"{OBS_IMAGE}.top" in batch
+    assert f"{OBS_IMAGE}.left" in batch
+    assert OBS_STATE in batch
+
+    # Check values are preserved
+    assert torch.allclose(batch[f"{OBS_IMAGE}.top"], observation_dict[f"{OBS_IMAGE}.top"])
+    assert torch.allclose(batch[f"{OBS_IMAGE}.left"], observation_dict[f"{OBS_IMAGE}.left"])
+    assert batch[OBS_STATE] == [1, 2, 3, 4]
+
+    # Check other fields are mapped to next.* format
+    assert batch[ACTION] == "action_data"
+    assert batch[REWARD] == 1.5
+    assert batch[DONE]
+    assert not batch[TRUNCATED]
+    assert batch["info"] == {"episode": 42}
+
+
+def test_no_observation_keys():
+    """Test behavior when there are no observation.* keys."""
+    batch = {
+        ACTION: torch.tensor([1.0, 2.0]),
+        REWARD: 2.0,
+        DONE: False,
+        TRUNCATED: True,
+        "info": {"test": "no_obs"},
+    }
+
+    transition = batch_to_transition(batch)
+
+    # Observation should be None when no observation.* keys
+    assert transition[TransitionKey.OBSERVATION] is None
+
+    # Check other fields
+    assert torch.allclose(transition[TransitionKey.ACTION], torch.tensor([1.0, 2.0]))
+    assert transition[TransitionKey.REWARD] == 2.0
+    assert not transition[TransitionKey.DONE]
+    assert transition[TransitionKey.TRUNCATED]
+    assert transition[TransitionKey.INFO] == {"test": "no_obs"}
+
+    # Round trip should work
+    reconstructed_batch = transition_to_batch(transition)
+    assert torch.allclose(reconstructed_batch[ACTION], torch.tensor([1.0, 2.0]))
+    assert reconstructed_batch[REWARD] == 2.0
+    assert not reconstructed_batch[DONE]
+    assert reconstructed_batch[TRUNCATED]
+    assert reconstructed_batch["info"] == {"test": "no_obs"}
+
+
+def test_minimal_batch():
+    """Test with minimal batch containing only observation.* and action."""
+    batch = {OBS_STATE: "minimal_state", ACTION: torch.tensor([0.5])}
+
+    transition = batch_to_transition(batch)
+
+    # Check observation
+    assert transition[TransitionKey.OBSERVATION] == {OBS_STATE: "minimal_state"}
+    assert torch.allclose(transition[TransitionKey.ACTION], torch.tensor([0.5]))
+
+    # Check defaults
+    assert transition[TransitionKey.REWARD] == 0.0
+    assert not transition[TransitionKey.DONE]
+    assert not transition[TransitionKey.TRUNCATED]
+    assert transition[TransitionKey.INFO] == {}
+    assert transition[TransitionKey.COMPLEMENTARY_DATA] == {}
+
+    # Round trip
+    reconstructed_batch = transition_to_batch(transition)
+    assert reconstructed_batch[OBS_STATE] == "minimal_state"
+    assert torch.allclose(reconstructed_batch[ACTION], torch.tensor([0.5]))
+    assert reconstructed_batch[REWARD] == 0.0
+    assert not reconstructed_batch[DONE]
+    assert not reconstructed_batch[TRUNCATED]
+    assert reconstructed_batch["info"] == {}
+
+
+def test_empty_batch():
+    """Test behavior with empty batch."""
+    batch = {}
+
+    transition = batch_to_transition(batch)
+
+    # All fields should have defaults
+    assert transition[TransitionKey.OBSERVATION] is None
+    assert transition[TransitionKey.ACTION] is None
+    assert transition[TransitionKey.REWARD] == 0.0
+    assert not transition[TransitionKey.DONE]
+    assert not transition[TransitionKey.TRUNCATED]
+    assert transition[TransitionKey.INFO] == {}
+    assert transition[TransitionKey.COMPLEMENTARY_DATA] == {}
+
+    # Round trip
+    reconstructed_batch = transition_to_batch(transition)
+    assert reconstructed_batch[ACTION] is None
+    assert reconstructed_batch[REWARD] == 0.0
+    assert not reconstructed_batch[DONE]
+    assert not reconstructed_batch[TRUNCATED]
+    assert reconstructed_batch["info"] == {}
+
+
+def test_complex_nested_observation():
+    """Test with complex nested observation data."""
+    batch = {
+        f"{OBS_IMAGE}.top": {"image": torch.randn(1, 3, 128, 128), "timestamp": 1234567890},
+        f"{OBS_IMAGE}.left": {"image": torch.randn(1, 3, 128, 128), "timestamp": 1234567891},
+        OBS_STATE: torch.randn(7),
+        ACTION: torch.randn(8),
+        REWARD: 3.14,
+        DONE: False,
+        TRUNCATED: True,
+        "info": {"episode_length": 200, "success": True},
+    }
+
+    transition = batch_to_transition(batch)
+    reconstructed_batch = transition_to_batch(transition)
+
+    # Check that all observation keys are preserved
+    original_obs_keys = {k for k in batch if k.startswith(OBS_PREFIX)}
+    reconstructed_obs_keys = {k for k in reconstructed_batch if k.startswith(OBS_PREFIX)}
+
+    assert original_obs_keys == reconstructed_obs_keys
+
+    # Check tensor values
+    assert torch.allclose(batch[OBS_STATE], reconstructed_batch[OBS_STATE])
+
+    # Check nested dict with tensors
+    assert torch.allclose(
+        batch[f"{OBS_IMAGE}.top"]["image"], reconstructed_batch[f"{OBS_IMAGE}.top"]["image"]
+    )
+    assert torch.allclose(
+        batch[f"{OBS_IMAGE}.left"]["image"], reconstructed_batch[f"{OBS_IMAGE}.left"]["image"]
+    )
+
+    # Check action tensor
+    assert torch.allclose(batch[ACTION], reconstructed_batch[ACTION])
+
+    # Check other fields
+    assert batch[REWARD] == reconstructed_batch[REWARD]
+    assert batch[DONE] == reconstructed_batch[DONE]
+    assert batch[TRUNCATED] == reconstructed_batch[TRUNCATED]
+    assert batch["info"] == reconstructed_batch["info"]
+
+
+def test_custom_converter():
+    """Test that custom converters can still be used."""
+
+    def to_tr(batch):
+        # Custom converter that modifies the reward
+        tr = batch_to_transition(batch)
+        # Double the reward
+        reward = tr.get(TransitionKey.REWARD, 0.0)
+        new_tr = tr.copy()
+        new_tr[TransitionKey.REWARD] = reward * 2 if reward is not None else 0.0
+        return new_tr
+
+    def to_batch(tr):
+        batch = transition_to_batch(tr)
+        return batch
+
+    processor = DataProcessorPipeline(steps=[], to_transition=to_tr, to_output=to_batch)
+
+    batch = {
+        OBS_STATE: torch.randn(1, 4),
+        ACTION: torch.randn(1, 2),
+        REWARD: 1.0,
+        DONE: False,
+    }
+
+    result = processor(batch)
+
+    # Check the reward was doubled by our custom converter
+    assert result[REWARD] == 2.0
+    assert torch.allclose(result[OBS_STATE], batch[OBS_STATE])
+    assert torch.allclose(result[ACTION], batch[ACTION])

tests/processor/test_converters.py

@@ -0,0 +1,309 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pytest
+import torch
+
+from lerobot.processor.converters import (
+    batch_to_transition,
+    create_transition,
+    to_tensor,
+    transition_to_batch,
+)
+from lerobot.types import TransitionKey
+from lerobot.utils.constants import ACTION, DONE, OBS_STATE, OBS_STR, REWARD
+
+
+# Tests for the unified to_tensor function
+def test_to_tensor_numpy_arrays():
+    """Test to_tensor with various numpy arrays."""
+    # Regular numpy array
+    arr = np.array([1.0, 2.0, 3.0])
+    result = to_tensor(arr)
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype == torch.float32
+    assert torch.allclose(result, torch.tensor([1.0, 2.0, 3.0]))
+
+    # Different numpy dtypes should convert to float32 by default
+    int_arr = np.array([1, 2, 3], dtype=np.int64)
+    result = to_tensor(int_arr)
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype == torch.float32
+    assert torch.allclose(result, torch.tensor([1.0, 2.0, 3.0]))
+
+    # uint8 arrays (previously "preserved") should now convert
+    uint8_arr = np.array([100, 150, 200], dtype=np.uint8)
+    result = to_tensor(uint8_arr)
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype == torch.float32
+    assert torch.allclose(result, torch.tensor([100.0, 150.0, 200.0]))
+
+
+def test_to_tensor_numpy_scalars():
+    """Test to_tensor with numpy scalars (0-dimensional arrays)."""
+    # numpy float32 scalar
+    scalar = np.float32(3.14)
+    result = to_tensor(scalar)
+    assert isinstance(result, torch.Tensor)
+    assert result.ndim == 0  # Should be 0-dimensional tensor
+    assert result.dtype == torch.float32
+    assert result.item() == pytest.approx(3.14)
+
+    # numpy int32 scalar
+    int_scalar = np.int32(42)
+    result = to_tensor(int_scalar)
+    assert isinstance(result, torch.Tensor)
+    assert result.ndim == 0
+    assert result.dtype == torch.float32
+    assert result.item() == pytest.approx(42.0)
+
+
+def test_to_tensor_python_scalars():
+    """Test to_tensor with Python scalars."""
+    # Python int
+    result = to_tensor(42)
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype == torch.float32
+    assert result.item() == pytest.approx(42.0)
+
+    # Python float
+    result = to_tensor(3.14)
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype == torch.float32
+    assert result.item() == pytest.approx(3.14)
+
+
+def test_to_tensor_sequences():
+    """Test to_tensor with lists and tuples."""
+    # List
+    result = to_tensor([1, 2, 3])
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype == torch.float32
+    assert torch.allclose(result, torch.tensor([1.0, 2.0, 3.0]))
+
+    # Tuple
+    result = to_tensor((4.5, 5.5, 6.5))
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype == torch.float32
+    assert torch.allclose(result, torch.tensor([4.5, 5.5, 6.5]))
+
+
+def test_to_tensor_existing_tensors():
+    """Test to_tensor with existing PyTorch tensors."""
+    # Tensor with same dtype should pass through with potential device change
+    tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
+    result = to_tensor(tensor)
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype == torch.float32
+    assert torch.allclose(result, tensor)
+
+    # Tensor with different dtype should convert
+    int_tensor = torch.tensor([1, 2, 3], dtype=torch.int64)
+    result = to_tensor(int_tensor)
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype == torch.float32
+    assert torch.allclose(result, torch.tensor([1.0, 2.0, 3.0]))
+
+
+def test_to_tensor_dictionaries():
+    """Test to_tensor with nested dictionaries."""
+    # Simple dictionary
+    data = {"mean": [0.1, 0.2], "std": np.array([1.0, 2.0]), "count": 42}
+    result = to_tensor(data)
+    assert isinstance(result, dict)
+    assert isinstance(result["mean"], torch.Tensor)
+    assert isinstance(result["std"], torch.Tensor)
+    assert isinstance(result["count"], torch.Tensor)
+    assert torch.allclose(result["mean"], torch.tensor([0.1, 0.2]))
+    assert torch.allclose(result["std"], torch.tensor([1.0, 2.0]))
+    assert result["count"].item() == pytest.approx(42.0)
+
+    # Nested dictionary
+    nested = {
+        ACTION: {"mean": [0.1, 0.2], "std": [1.0, 2.0]},
+        OBS_STR: {"mean": np.array([0.5, 0.6]), "count": 10},
+    }
+    result = to_tensor(nested)
+    assert isinstance(result, dict)
+    assert isinstance(result[ACTION], dict)
+    assert isinstance(result[OBS_STR], dict)
+    assert isinstance(result[ACTION]["mean"], torch.Tensor)
+    assert isinstance(result[OBS_STR]["mean"], torch.Tensor)
+    assert torch.allclose(result[ACTION]["mean"], torch.tensor([0.1, 0.2]))
+    assert torch.allclose(result[OBS_STR]["mean"], torch.tensor([0.5, 0.6]))
+
+
+def test_to_tensor_none_filtering():
+    """Test that None values are filtered out from dictionaries."""
+    data = {"valid": [1, 2, 3], "none_value": None, "nested": {"valid": [4, 5], "also_none": None}}
+    result = to_tensor(data)
+    assert "none_value" not in result
+    assert "also_none" not in result["nested"]
+    assert "valid" in result
+    assert "valid" in result["nested"]
+    assert torch.allclose(result["valid"], torch.tensor([1.0, 2.0, 3.0]))
+
+
+def test_to_tensor_dtype_parameter():
+    """Test to_tensor with different dtype parameters."""
+    arr = np.array([1, 2, 3])
+
+    # Default dtype (float32)
+    result = to_tensor(arr)
+    assert result.dtype == torch.float32
+
+    # Explicit float32
+    result = to_tensor(arr, dtype=torch.float32)
+    assert result.dtype == torch.float32
+
+    # Float64
+    result = to_tensor(arr, dtype=torch.float64)
+    assert result.dtype == torch.float64
+
+    # Preserve original dtype
+    float64_arr = np.array([1.0, 2.0, 3.0], dtype=np.float64)
+    result = to_tensor(float64_arr, dtype=None)
+    assert result.dtype == torch.float64
+
+
+def test_to_tensor_device_parameter():
+    """Test to_tensor with device parameter."""
+    arr = np.array([1.0, 2.0, 3.0])
+
+    # CPU device (default)
+    result = to_tensor(arr, device="cpu")
+    assert result.device.type == "cpu"
+
+    # CUDA device (if available)
+    if torch.cuda.is_available():
+        result = to_tensor(arr, device="cuda")
+        assert result.device.type == "cuda"
+
+
+def test_to_tensor_empty_dict():
+    """Test to_tensor with empty dictionary."""
+    result = to_tensor({})
+    assert isinstance(result, dict)
+    assert len(result) == 0
+
+
+def test_to_tensor_unsupported_type():
+    """Test to_tensor with unsupported types raises TypeError."""
+    with pytest.raises(TypeError, match="Unsupported type for tensor conversion"):
+        to_tensor("unsupported_string")
+
+    with pytest.raises(TypeError, match="Unsupported type for tensor conversion"):
+        to_tensor(object())
+
+
+def test_batch_to_transition_with_index_fields():
+    """Test that batch_to_transition handles index and task_index fields correctly."""
+
+    # Create batch with index and task_index fields
+    batch = {
+        OBS_STATE: torch.randn(1, 7),
+        ACTION: torch.randn(1, 4),
+        REWARD: 1.5,
+        DONE: False,
+        "task": ["pick_cube"],
+        "index": torch.tensor([42], dtype=torch.int64),
+        "task_index": torch.tensor([3], dtype=torch.int64),
+    }
+
+    transition = batch_to_transition(batch)
+
+    # Check basic transition structure
+    assert TransitionKey.OBSERVATION in transition
+    assert TransitionKey.ACTION in transition
+    assert TransitionKey.COMPLEMENTARY_DATA in transition
+
+    # Check that index and task_index are in complementary_data
+    comp_data = transition[TransitionKey.COMPLEMENTARY_DATA]
+    assert "index" in comp_data
+    assert "task_index" in comp_data
+    assert "task" in comp_data
+
+    # Verify values
+    assert torch.equal(comp_data["index"], batch["index"])
+    assert torch.equal(comp_data["task_index"], batch["task_index"])
+    assert comp_data["task"] == batch["task"]
+
+
+def testtransition_to_batch_with_index_fields():
+    """Test that transition_to_batch handles index and task_index fields correctly."""
+
+    # Create transition with index and task_index in complementary_data
+    transition = create_transition(
+        observation={OBS_STATE: torch.randn(1, 7)},
+        action=torch.randn(1, 4),
+        reward=1.5,
+        done=False,
+        complementary_data={
+            "task": ["navigate"],
+            "index": torch.tensor([100], dtype=torch.int64),
+            "task_index": torch.tensor([5], dtype=torch.int64),
+        },
+    )
+
+    batch = transition_to_batch(transition)
+
+    # Check that index and task_index are in the batch
+    assert "index" in batch
+    assert "task_index" in batch
+    assert "task" in batch
+
+    # Verify values
+    assert torch.equal(batch["index"], transition[TransitionKey.COMPLEMENTARY_DATA]["index"])
+    assert torch.equal(batch["task_index"], transition[TransitionKey.COMPLEMENTARY_DATA]["task_index"])
+    assert batch["task"] == transition[TransitionKey.COMPLEMENTARY_DATA]["task"]
+
+
+def test_batch_to_transition_without_index_fields():
+    """Test that conversion works without index and task_index fields."""
+
+    # Batch without index/task_index
+    batch = {
+        OBS_STATE: torch.randn(1, 7),
+        ACTION: torch.randn(1, 4),
+        "task": ["pick_cube"],
+    }
+
+    transition = batch_to_transition(batch)
+    comp_data = transition[TransitionKey.COMPLEMENTARY_DATA]
+
+    # Should have task but not index/task_index
+    assert "task" in comp_data
+    assert "index" not in comp_data
+    assert "task_index" not in comp_data
+
+
+def test_transition_to_batch_without_index_fields():
+    """Test that conversion works without index and task_index fields."""
+
+    # Transition without index/task_index
+    transition = create_transition(
+        observation={OBS_STATE: torch.randn(1, 7)},
+        action=torch.randn(1, 4),
+        complementary_data={"task": ["navigate"]},
+    )
+
+    batch = transition_to_batch(transition)
+
+    # Should have task but not index/task_index
+    assert "task" in batch
+    assert "index" not in batch
+    assert "task_index" not in batch

tests/processor/test_diffusion_processor.py

@@ -0,0 +1,398 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for Diffusion policy processor."""
+
+import tempfile
+
+import pytest
+import torch
+
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+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,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    RenameObservationsProcessorStep,
+    TransitionKey,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import create_transition, transition_to_batch
+from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE
+
+
+def create_default_config():
+    """Create a default Diffusion configuration for testing."""
+    config = DiffusionConfig()
+    config.input_features = {
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(7,)),
+        OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
+    }
+    config.output_features = {
+        ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(6,)),
+    }
+    config.normalization_mapping = {
+        FeatureType.STATE: NormalizationMode.MEAN_STD,
+        FeatureType.VISUAL: NormalizationMode.IDENTITY,
+        FeatureType.ACTION: NormalizationMode.MIN_MAX,
+    }
+    config.device = "cpu"
+    return config
+
+
+def create_default_stats():
+    """Create default dataset statistics for testing."""
+    return {
+        OBS_STATE: {"mean": torch.zeros(7), "std": torch.ones(7)},
+        OBS_IMAGE: {},  # No normalization for images
+        ACTION: {"min": torch.full((6,), -1.0), "max": torch.ones(6)},
+    }
+
+
+def test_make_diffusion_processor_basic():
+    """Test basic creation of Diffusion processor."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
+
+    # Check processor names
+    assert preprocessor.name == "policy_preprocessor"
+    assert postprocessor.name == "policy_postprocessor"
+
+    # Check steps in preprocessor
+    assert len(preprocessor.steps) == 4
+    assert isinstance(preprocessor.steps[0], RenameObservationsProcessorStep)
+    assert isinstance(preprocessor.steps[1], AddBatchDimensionProcessorStep)
+    assert isinstance(preprocessor.steps[2], DeviceProcessorStep)
+    assert isinstance(preprocessor.steps[3], NormalizerProcessorStep)
+
+    # Check steps in postprocessor
+    assert len(postprocessor.steps) == 2
+    assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
+    assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
+
+
+def test_diffusion_processor_with_images():
+    """Test Diffusion processor with image observations."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_diffusion_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Create test data with images
+    observation = {
+        OBS_STATE: torch.randn(7),
+        OBS_IMAGE: torch.randn(3, 224, 224),
+    }
+    action = torch.randn(6)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data is batched
+    assert processed[OBS_STATE].shape == (1, 7)
+    assert processed[OBS_IMAGE].shape == (1, 3, 224, 224)
+    assert processed[TransitionKey.ACTION.value].shape == (1, 6)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_diffusion_processor_cuda():
+    """Test Diffusion processor with CUDA device."""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_diffusion_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Create CPU data
+    observation = {
+        OBS_STATE: torch.randn(7),
+        OBS_IMAGE: torch.randn(3, 224, 224),
+    }
+    action = torch.randn(6)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data is on CUDA
+    assert processed[OBS_STATE].device.type == "cuda"
+    assert processed[OBS_IMAGE].device.type == "cuda"
+    assert processed[TransitionKey.ACTION.value].device.type == "cuda"
+
+    # Process through postprocessor
+    postprocessed = postprocessor(processed[TransitionKey.ACTION.value])
+
+    # Check that action is back on CPU
+    assert postprocessed.device.type == "cpu"
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_diffusion_processor_accelerate_scenario():
+    """Test Diffusion processor in simulated Accelerate scenario."""
+    config = create_default_config()
+    config.device = "cuda:0"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_diffusion_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Simulate Accelerate: data already on GPU
+    device = torch.device("cuda:0")
+    observation = {
+        OBS_STATE: torch.randn(1, 7).to(device),
+        OBS_IMAGE: torch.randn(1, 3, 224, 224).to(device),
+    }
+    action = torch.randn(1, 6).to(device)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data stays on same GPU
+    assert processed[OBS_STATE].device == device
+    assert processed[OBS_IMAGE].device == device
+    assert processed[TransitionKey.ACTION.value].device == device
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs")
+def test_diffusion_processor_multi_gpu():
+    """Test Diffusion processor with multi-GPU setup."""
+    config = create_default_config()
+    config.device = "cuda:0"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
+
+    # Simulate data on different GPU
+    device = torch.device("cuda:1")
+    observation = {
+        OBS_STATE: torch.randn(1, 7).to(device),
+        OBS_IMAGE: torch.randn(1, 3, 224, 224).to(device),
+    }
+    action = torch.randn(1, 6).to(device)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data stays on cuda:1
+    assert processed[OBS_STATE].device == device
+    assert processed[OBS_IMAGE].device == device
+    assert processed[TransitionKey.ACTION.value].device == device
+
+
+def test_diffusion_processor_without_stats():
+    """Test Diffusion processor creation without dataset statistics."""
+    config = create_default_config()
+
+    preprocessor, postprocessor = make_diffusion_pre_post_processors(
+        config,
+        dataset_stats=None,
+    )
+
+    # Should still create processors
+    assert preprocessor is not None
+    assert postprocessor is not None
+
+    # Process should still work
+    observation = {
+        OBS_STATE: torch.randn(7),
+        OBS_IMAGE: torch.randn(3, 224, 224),
+    }
+    action = torch.randn(6)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    processed = preprocessor(batch)
+    assert processed is not None
+
+
+def test_diffusion_processor_save_and_load():
+    """Test saving and loading Diffusion processor."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Save preprocessor
+        preprocessor.save_pretrained(tmpdir)
+
+        # Load preprocessor
+        loaded_preprocessor = DataProcessorPipeline.from_pretrained(
+            tmpdir, config_filename="policy_preprocessor.json"
+        )
+
+        # Test that loaded processor works
+        observation = {
+            OBS_STATE: torch.randn(7),
+            OBS_IMAGE: torch.randn(3, 224, 224),
+        }
+        action = torch.randn(6)
+        transition = create_transition(observation, action)
+        batch = transition_to_batch(transition)
+
+        processed = loaded_preprocessor(batch)
+        assert processed[OBS_STATE].shape == (1, 7)
+        assert processed[OBS_IMAGE].shape == (1, 3, 224, 224)
+        assert processed[TransitionKey.ACTION.value].shape == (1, 6)
+
+
+def test_diffusion_processor_identity_normalization():
+    """Test that images with IDENTITY normalization are not normalized."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_diffusion_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Create test data
+    image_value = torch.rand(3, 224, 224) * 255  # Large values
+    observation = {
+        OBS_STATE: torch.randn(7),
+        OBS_IMAGE: image_value.clone(),
+    }
+    action = torch.randn(6)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Image should not be normalized (IDENTITY mode)
+    # Just batched
+    assert torch.allclose(processed[OBS_IMAGE][0], image_value, rtol=1e-5)
+
+
+def test_diffusion_processor_batch_consistency():
+    """Test Diffusion processor with different batch sizes."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_diffusion_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Test with different batch sizes
+    for batch_size in [1, 8, 32]:
+        observation = {
+            OBS_STATE: torch.randn(batch_size, 7) if batch_size > 1 else torch.randn(7),
+            OBS_IMAGE: torch.randn(batch_size, 3, 224, 224) if batch_size > 1 else torch.randn(3, 224, 224),
+        }
+        action = torch.randn(batch_size, 6) if batch_size > 1 else torch.randn(6)
+        transition = create_transition(observation, action)
+
+        batch = transition_to_batch(transition)
+
+        processed = preprocessor(batch)
+
+        # Check correct batch size
+        expected_batch = batch_size if batch_size > 1 else 1
+        assert processed[OBS_STATE].shape[0] == expected_batch
+        assert processed[OBS_IMAGE].shape[0] == expected_batch
+        assert processed[TransitionKey.ACTION.value].shape[0] == expected_batch
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_diffusion_processor_bfloat16_device_float32_normalizer():
+    """Test: DeviceProcessor(bfloat16) + NormalizerProcessor(float32) → output bfloat16 via automatic adaptation"""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    preprocessor, _ = make_diffusion_pre_post_processors(config, stats)
+
+    # Modify the pipeline to use bfloat16 device processor with float32 normalizer
+    modified_steps = []
+    for step in preprocessor.steps:
+        if isinstance(step, DeviceProcessorStep):
+            # Device processor converts to bfloat16
+            modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="bfloat16"))
+        elif isinstance(step, NormalizerProcessorStep):
+            # Normalizer stays configured as float32 (will auto-adapt to bfloat16)
+            norm_step = step  # Now type checker knows this is NormalizerProcessorStep
+            modified_steps.append(
+                NormalizerProcessorStep(
+                    features=norm_step.features,
+                    norm_map=norm_step.norm_map,
+                    stats=norm_step.stats,
+                    device=config.device,
+                    dtype=torch.float32,  # Deliberately configured as float32
+                )
+            )
+        else:
+            modified_steps.append(step)
+    preprocessor.steps = modified_steps
+
+    # Verify initial normalizer configuration
+    normalizer_step = preprocessor.steps[3]  # NormalizerProcessorStep
+    assert normalizer_step.dtype == torch.float32
+
+    # Create test data with both state and visual observations
+    observation = {
+        OBS_STATE: torch.randn(7, dtype=torch.float32),
+        OBS_IMAGE: torch.randn(3, 224, 224, dtype=torch.float32),
+    }
+    action = torch.randn(6, dtype=torch.float32)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through full pipeline
+    processed = preprocessor(batch)
+
+    # Verify: DeviceProcessor → bfloat16, NormalizerProcessor adapts → final output is bfloat16
+    assert processed[OBS_STATE].dtype == torch.bfloat16
+    assert processed[OBS_IMAGE].dtype == torch.bfloat16  # IDENTITY normalization still gets dtype conversion
+    assert processed[TransitionKey.ACTION.value].dtype == torch.bfloat16
+
+    # Verify normalizer automatically adapted its internal state
+    assert normalizer_step.dtype == torch.bfloat16
+    # Check state stats (has normalization)
+    for stat_tensor in normalizer_step._tensor_stats[OBS_STATE].values():
+        assert stat_tensor.dtype == torch.bfloat16
+    # OBS_IMAGE uses IDENTITY normalization, so no stats to check

tests/processor/test_gaussian_actor_processor.py

@@ -0,0 +1,414 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for SAC policy processor."""
+
+import tempfile
+
+import pytest
+import torch
+
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.policies.gaussian_actor.configuration_gaussian_actor import GaussianActorConfig
+from lerobot.policies.gaussian_actor.processor_gaussian_actor import make_gaussian_actor_pre_post_processors
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DataProcessorPipeline,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    RenameObservationsProcessorStep,
+    TransitionKey,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import create_transition, transition_to_batch
+from lerobot.utils.constants import ACTION, OBS_STATE
+
+
+def create_default_config():
+    """Create a default SAC configuration for testing."""
+    config = GaussianActorConfig()
+    config.input_features = {
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(10,)),
+    }
+    config.output_features = {
+        ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(5,)),
+    }
+    config.normalization_mapping = {
+        FeatureType.STATE: NormalizationMode.MEAN_STD,
+        FeatureType.ACTION: NormalizationMode.MIN_MAX,
+    }
+    config.device = "cpu"
+    return config
+
+
+def create_default_stats():
+    """Create default dataset statistics for testing."""
+    return {
+        OBS_STATE: {"mean": torch.zeros(10), "std": torch.ones(10)},
+        ACTION: {"min": torch.full((5,), -1.0), "max": torch.ones(5)},
+    }
+
+
+def test_make_sac_processor_basic():
+    """Test basic creation of SAC processor."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_gaussian_actor_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Check processor names
+    assert preprocessor.name == "policy_preprocessor"
+    assert postprocessor.name == "policy_postprocessor"
+
+    # Check steps in preprocessor
+    assert len(preprocessor.steps) == 4
+    assert isinstance(preprocessor.steps[0], RenameObservationsProcessorStep)
+    assert isinstance(preprocessor.steps[1], AddBatchDimensionProcessorStep)
+    assert isinstance(preprocessor.steps[2], DeviceProcessorStep)
+    assert isinstance(preprocessor.steps[3], NormalizerProcessorStep)
+
+    # Check steps in postprocessor
+    assert len(postprocessor.steps) == 2
+    assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
+    assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
+
+
+def test_gaussian_actor_processor_normalization_modes():
+    """Test that SAC processor correctly handles different normalization modes."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_gaussian_actor_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Create test data
+    observation = {OBS_STATE: torch.randn(10) * 2}  # Larger values to test normalization
+    action = torch.rand(5) * 2 - 1  # Range [-1, 1]
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+    processed = preprocessor(batch)
+
+    # Check that data is normalized and batched
+    # State should be mean-std normalized
+    # Action should be min-max normalized to [-1, 1]
+    assert processed[OBS_STATE].shape == (1, 10)
+    assert processed[TransitionKey.ACTION.value].shape == (1, 5)
+
+    # Process action through postprocessor
+    postprocessed = postprocessor(processed[TransitionKey.ACTION.value])
+
+    # Check that action is unnormalized (but still batched)
+    assert postprocessed.shape == (1, 5)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_gaussian_actor_processor_cuda():
+    """Test SAC processor with CUDA device."""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_gaussian_actor_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Create CPU data
+    observation = {OBS_STATE: torch.randn(10)}
+    action = torch.randn(5)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+    processed = preprocessor(batch)
+
+    # Check that data is on CUDA
+    assert processed[OBS_STATE].device.type == "cuda"
+    assert processed[TransitionKey.ACTION.value].device.type == "cuda"
+
+    # Process through postprocessor
+    postprocessed = postprocessor(processed[TransitionKey.ACTION.value])
+
+    # Check that action is back on CPU
+    assert postprocessed.device.type == "cpu"
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_gaussian_actor_processor_accelerate_scenario():
+    """Test SAC processor in simulated Accelerate scenario."""
+    config = create_default_config()
+    config.device = "cuda:0"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_gaussian_actor_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Simulate Accelerate: data already on GPU
+    device = torch.device("cuda:0")
+    observation = {OBS_STATE: torch.randn(10).to(device)}
+    action = torch.randn(5).to(device)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+    processed = preprocessor(batch)
+
+    # Check that data stays on same GPU
+    assert processed[OBS_STATE].device == device
+    assert processed[TransitionKey.ACTION.value].device == device
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs")
+def test_gaussian_actor_processor_multi_gpu():
+    """Test SAC processor with multi-GPU setup."""
+    config = create_default_config()
+    config.device = "cuda:0"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_gaussian_actor_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Simulate data on different GPU
+    device = torch.device("cuda:1")
+    observation = {OBS_STATE: torch.randn(10).to(device)}
+    action = torch.randn(5).to(device)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+    processed = preprocessor(batch)
+
+    # Check that data stays on cuda:1
+    assert processed[OBS_STATE].device == device
+    assert processed[TransitionKey.ACTION.value].device == device
+
+
+def test_gaussian_actor_processor_without_stats():
+    """Test SAC processor creation without dataset statistics."""
+    config = create_default_config()
+
+    preprocessor, postprocessor = make_gaussian_actor_pre_post_processors(config, dataset_stats=None)
+
+    # Should still create processors
+    assert preprocessor is not None
+    assert postprocessor is not None
+
+    # Process should still work
+    observation = {OBS_STATE: torch.randn(10)}
+    action = torch.randn(5)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    processed = preprocessor(batch)
+    assert processed is not None
+
+
+def test_gaussian_actor_processor_save_and_load():
+    """Test saving and loading SAC processor."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_gaussian_actor_pre_post_processors(
+        config,
+        stats,
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Save preprocessor
+        preprocessor.save_pretrained(tmpdir)
+
+        # Load preprocessor
+        loaded_preprocessor = DataProcessorPipeline.from_pretrained(
+            tmpdir, config_filename="policy_preprocessor.json"
+        )
+
+        # Test that loaded processor works
+        observation = {OBS_STATE: torch.randn(10)}
+        action = torch.randn(5)
+        transition = create_transition(observation, action)
+        batch = transition_to_batch(transition)
+
+        processed = loaded_preprocessor(batch)
+        assert processed[OBS_STATE].shape == (1, 10)
+        assert processed[TransitionKey.ACTION.value].shape == (1, 5)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_gaussian_actor_processor_mixed_precision():
+    """Test SAC processor with mixed precision."""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    # Create processor
+    preprocessor, postprocessor = make_gaussian_actor_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Replace DeviceProcessorStep with one that uses float16
+    modified_steps = []
+    for step in preprocessor.steps:
+        if isinstance(step, DeviceProcessorStep):
+            modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16"))
+        elif isinstance(step, NormalizerProcessorStep):
+            # Update normalizer to use the same device as the device processor
+            norm_step = step  # Now type checker knows this is NormalizerProcessorStep
+            modified_steps.append(
+                NormalizerProcessorStep(
+                    features=norm_step.features,
+                    norm_map=norm_step.norm_map,
+                    stats=norm_step.stats,
+                    device=config.device,
+                    dtype=torch.float16,  # Match the float16 dtype
+                )
+            )
+        else:
+            modified_steps.append(step)
+    preprocessor.steps = modified_steps
+
+    # Create test data
+    observation = {OBS_STATE: torch.randn(10, dtype=torch.float32)}
+    action = torch.randn(5, dtype=torch.float32)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+    processed = preprocessor(batch)
+
+    # Check that data is converted to float16
+    assert processed[OBS_STATE].dtype == torch.float16
+    assert processed[TransitionKey.ACTION.value].dtype == torch.float16
+
+
+def test_gaussian_actor_processor_batch_data():
+    """Test SAC processor with batched data."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_gaussian_actor_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Test with batched data
+    batch_size = 32
+    observation = {OBS_STATE: torch.randn(batch_size, 10)}
+    action = torch.randn(batch_size, 5)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+    processed = preprocessor(batch)
+
+    # Check that batch dimension is preserved
+    assert processed[OBS_STATE].shape == (batch_size, 10)
+    assert processed[TransitionKey.ACTION.value].shape == (batch_size, 5)
+
+
+def test_gaussian_actor_processor_edge_cases():
+    """Test SAC processor with edge cases."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_gaussian_actor_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Test with observation that has no state key but still exists
+    observation = {"observation.dummy": torch.randn(1)}  # Some dummy observation to pass validation
+    action = torch.randn(5)
+    batch = {TransitionKey.ACTION.value: action, **observation}
+    processed = preprocessor(batch)
+    # observation.state wasn't in original, so it won't be in processed
+    assert OBS_STATE not in processed
+    assert processed[TransitionKey.ACTION.value].shape == (1, 5)
+
+    # Test with zero action (representing "null" action)
+    transition = create_transition(observation={OBS_STATE: torch.randn(10)}, action=torch.zeros(5))
+    batch = transition_to_batch(transition)
+    processed = preprocessor(batch)
+    assert processed[OBS_STATE].shape == (1, 10)
+    # Action should be present and batched, even if it's zeros
+    assert processed[TransitionKey.ACTION.value].shape == (1, 5)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_gaussian_actor_processor_bfloat16_device_float32_normalizer():
+    """Test: DeviceProcessor(bfloat16) + NormalizerProcessor(float32) → output bfloat16 via automatic adaptation"""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    preprocessor, _ = make_gaussian_actor_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Modify the pipeline to use bfloat16 device processor with float32 normalizer
+    modified_steps = []
+    for step in preprocessor.steps:
+        if isinstance(step, DeviceProcessorStep):
+            # Device processor converts to bfloat16
+            modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="bfloat16"))
+        elif isinstance(step, NormalizerProcessorStep):
+            # Normalizer stays configured as float32 (will auto-adapt to bfloat16)
+            norm_step = step  # Now type checker knows this is NormalizerProcessorStep
+            modified_steps.append(
+                NormalizerProcessorStep(
+                    features=norm_step.features,
+                    norm_map=norm_step.norm_map,
+                    stats=norm_step.stats,
+                    device=config.device,
+                    dtype=torch.float32,  # Deliberately configured as float32
+                )
+            )
+        else:
+            modified_steps.append(step)
+    preprocessor.steps = modified_steps
+
+    # Verify initial normalizer configuration
+    normalizer_step = preprocessor.steps[3]  # NormalizerProcessorStep
+    assert normalizer_step.dtype == torch.float32
+
+    # Create test data
+    observation = {OBS_STATE: torch.randn(10, dtype=torch.float32)}  # Start with float32
+    action = torch.randn(5, dtype=torch.float32)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    # Process through full pipeline
+    processed = preprocessor(batch)
+
+    # Verify: DeviceProcessor → bfloat16, NormalizerProcessor adapts → final output is bfloat16
+    assert processed[OBS_STATE].dtype == torch.bfloat16
+    assert processed[TransitionKey.ACTION.value].dtype == torch.bfloat16
+
+    # Verify normalizer automatically adapted its internal state
+    assert normalizer_step.dtype == torch.bfloat16
+    for stat_tensor in normalizer_step._tensor_stats[OBS_STATE].values():
+        assert stat_tensor.dtype == torch.bfloat16

tests/processor/test_libero_processor.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import torch
+
+from lerobot.envs.utils import preprocess_observation
+from lerobot.processor.env_processor import LiberoProcessorStep
+from lerobot.processor.pipeline import PolicyProcessorPipeline
+
+seed = 42
+np.random.seed(seed)
+
+B = 5
+obs1 = {
+    "pixels": {
+        "image": (np.random.rand(B, 256, 256, 3) * 255).astype(np.uint8),
+        "image2": (np.random.rand(B, 256, 256, 3) * 255).astype(np.uint8),
+    },
+    "robot_state": {
+        "eef": {
+            "pos": np.random.randn(B, 3),
+            "quat": np.random.randn(B, 4),
+            "mat": np.random.randn(B, 3, 3),
+        },
+        "gripper": {
+            "qpos": np.random.randn(B, 2),
+            "qvel": np.random.randn(B, 2),
+        },
+        "joints": {
+            "pos": np.random.randn(B, 7),
+            "vel": np.random.randn(B, 7),
+        },
+    },
+}
+
+observation = preprocess_observation(obs1)
+libero_preprocessor = PolicyProcessorPipeline(
+    steps=[
+        LiberoProcessorStep(),
+    ]
+)
+processed_obs = libero_preprocessor(observation)
+assert "observation.state" in processed_obs
+state = processed_obs["observation.state"]
+assert isinstance(state, torch.Tensor)
+assert state.dtype == torch.float32
+
+assert state.shape[0] == B
+assert state.shape[1] == 8
+
+assert "observation.images.image" in processed_obs
+assert "observation.images.image2" in processed_obs
+
+assert isinstance(processed_obs["observation.images.image"], torch.Tensor)
+assert isinstance(processed_obs["observation.images.image2"], torch.Tensor)
+
+assert processed_obs["observation.images.image"].shape == (B, 3, 256, 256)
+assert processed_obs["observation.images.image2"].shape == (B, 3, 256, 256)

tests/processor/test_migration_detection.py

@@ -0,0 +1,342 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Tests for processor migration detection functionality.
+"""
+
+import json
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from lerobot.processor.pipeline import DataProcessorPipeline, ProcessorMigrationError
+from lerobot.utils.constants import ACTION, OBS_STATE
+
+
+def test_is_processor_config_valid_configs():
+    """Test processor config detection with valid configurations."""
+    valid_configs = [
+        {"steps": []},  # Empty steps
+        {"steps": [{"class": "MyClass"}]},  # Class-based step
+        {"steps": [{"registry_name": "my_step"}]},  # Registry-based step
+        {"steps": [{"class": "A"}, {"registry_name": "B"}]},  # Mixed
+        {"name": "Test", "steps": [{"class": "MyClass"}]},  # With name
+    ]
+
+    for i, config in enumerate(valid_configs):
+        assert DataProcessorPipeline._is_processor_config(config), (
+            f"Valid config {i} should be detected as processor config: {config}"
+        )
+
+
+def test_is_processor_config_invalid_configs():
+    """Test processor config detection with invalid configurations."""
+    invalid_configs = [
+        {},  # No steps field
+        {"steps": "not a list"},  # Steps is not a list
+        {"steps": [{}]},  # Step without class or registry_name
+        {"steps": ["not a dict"]},  # Step is not a dict
+        {"steps": [{"other_field": "value"}]},  # Step with wrong fields
+        {"other_field": "value"},  # Completely different structure
+    ]
+
+    for i, config in enumerate(invalid_configs):
+        assert not DataProcessorPipeline._is_processor_config(config), (
+            f"Invalid config {i} should not be detected as processor config: {config}"
+        )
+
+
+def test_should_suggest_migration_with_processor_config():
+    """Test that migration is NOT suggested when processor config exists."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create a valid processor config
+        processor_config = {
+            "name": "TestProcessor",
+            "steps": [
+                {
+                    "class": "lerobot.processor.normalize.NormalizeStep",
+                    "config": {"mean": 0.0, "std": 1.0},
+                }
+            ],
+        }
+
+        with open(tmp_path / "processor.json", "w") as f:
+            json.dump(processor_config, f)
+
+        # Should NOT suggest migration (processor config exists)
+        result = DataProcessorPipeline._should_suggest_migration(tmp_path)
+        assert not result
+
+
+def test_should_suggest_migration_with_empty_processor_config():
+    """Test that migration is NOT suggested when empty processor config exists."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create an empty processor config
+        empty_processor_config = {
+            "name": "EmptyProcessor",
+            "steps": [],  # Empty steps is valid
+        }
+
+        with open(tmp_path / "empty_processor.json", "w") as f:
+            json.dump(empty_processor_config, f)
+
+        # Should NOT suggest migration (processor config exists, even if empty)
+        result = DataProcessorPipeline._should_suggest_migration(tmp_path)
+        assert not result
+
+
+def test_should_suggest_migration_with_model_config_only():
+    """Test that migration IS suggested when only model config exists."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create a model config (like old LeRobot format)
+        model_config = {
+            "type": "act",
+            "input_features": {OBS_STATE: {"shape": [7]}},
+            "output_features": {ACTION: {"shape": [7]}},
+            "hidden_dim": 256,
+            "n_obs_steps": 1,
+            "n_action_steps": 1,
+        }
+
+        with open(tmp_path / "config.json", "w") as f:
+            json.dump(model_config, f)
+
+        # SHOULD suggest migration (model config exists but no processor)
+        result = DataProcessorPipeline._should_suggest_migration(tmp_path)
+        assert result
+
+
+def test_should_suggest_migration_no_json_files():
+    """Test that migration is NOT suggested when no JSON files exist."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create some non-JSON files
+        with open(tmp_path / "model.safetensors", "w") as f:
+            f.write("fake model data")
+
+        with open(tmp_path / "README.md", "w") as f:
+            f.write("# Model README")
+
+        # Should NOT suggest migration (no JSON files)
+        result = DataProcessorPipeline._should_suggest_migration(tmp_path)
+        assert not result
+
+
+def test_should_suggest_migration_random_json_files():
+    """Test that migration IS suggested when JSON files exist but none are processor configs."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create some random JSON file (not a processor config)
+        random_config = {"some_field": "some_value", "another_field": 123}
+
+        with open(tmp_path / "random.json", "w") as f:
+            json.dump(random_config, f)
+
+        # SHOULD suggest migration (JSON files exist but none are processor configs)
+        result = DataProcessorPipeline._should_suggest_migration(tmp_path)
+        assert result
+
+
+def test_should_suggest_migration_mixed_configs():
+    """Test that migration is NOT suggested when processor config exists alongside other configs."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create both a processor config and a model config
+        processor_config = {"name": "TestProcessor", "steps": [{"registry_name": "normalize_step"}]}
+
+        model_config = {"type": "diffusion", "hidden_dim": 512}
+
+        with open(tmp_path / "processor.json", "w") as f:
+            json.dump(processor_config, f)
+
+        with open(tmp_path / "config.json", "w") as f:
+            json.dump(model_config, f)
+
+        # Should NOT suggest migration (processor config exists)
+        result = DataProcessorPipeline._should_suggest_migration(tmp_path)
+        assert not result
+
+
+def test_should_suggest_migration_invalid_json():
+    """Test that invalid JSON is handled gracefully."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create an invalid JSON file
+        with open(tmp_path / "invalid.json", "w") as f:
+            f.write("{ invalid json")
+
+        # Create a valid non-processor config
+        model_config = {"type": "act"}
+        with open(tmp_path / "model.json", "w") as f:
+            json.dump(model_config, f)
+
+        # SHOULD suggest migration (invalid JSON is ignored, but we have non-processor JSON)
+        result = DataProcessorPipeline._should_suggest_migration(tmp_path)
+        assert result
+
+
+def test_from_pretrained_multiple_json_files_migration_error():
+    """Test that multiple JSON files trigger ProcessorMigrationError."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create multiple non-processor configs
+        model_config = {"type": "act", "hidden_dim": 128}
+        train_config = {"batch_size": 32, "lr": 0.001}
+
+        with open(tmp_path / "config.json", "w") as f:
+            json.dump(model_config, f)
+
+        with open(tmp_path / "train_config.json", "w") as f:
+            json.dump(train_config, f)
+
+        # Should raise ProcessorMigrationError
+        with pytest.raises(ProcessorMigrationError) as exc_info:
+            DataProcessorPipeline.from_pretrained(tmp_path, config_filename="config.json")
+
+        # Check the error details
+        error = exc_info.value
+        assert str(tmp_path) in str(error.model_path)
+        assert "migrate_policy_normalization.py" in error.migration_command
+        assert "not a valid processor configuration" in error.original_error
+
+
+def test_from_pretrained_no_processor_config_migration_error():
+    """Test that missing processor config triggers ProcessorMigrationError."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create a model config but no processor
+        model_config = {"type": "diffusion", "hidden_dim": 256}
+
+        with open(tmp_path / "config.json", "w") as f:
+            json.dump(model_config, f)
+
+        # Should raise ProcessorMigrationError
+        with pytest.raises(ProcessorMigrationError) as exc_info:
+            DataProcessorPipeline.from_pretrained(tmp_path, config_filename="config.json")
+
+        # Check the error details
+        error = exc_info.value
+        assert str(tmp_path) in str(error.model_path)
+        assert "migrate_policy_normalization.py" in error.migration_command
+        assert "not a valid processor configuration" in error.original_error
+
+
+def test_from_pretrained_valid_processor_no_migration_error():
+    """Test that valid processor config does NOT trigger migration error."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create a valid processor config
+        processor_config = {
+            "name": "TestProcessor",
+            "steps": [],  # Empty is valid
+        }
+
+        with open(tmp_path / "processor.json", "w") as f:
+            json.dump(processor_config, f)
+
+        # Should succeed and create pipeline
+        pipeline = DataProcessorPipeline.from_pretrained(tmp_path, config_filename="processor.json")
+        assert pipeline is not None
+        assert pipeline.name == "TestProcessor"
+        assert len(pipeline) == 0
+
+
+def test_from_pretrained_no_json_files_no_migration_error():
+    """Test that directories with no JSON files don't trigger migration errors."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create some non-JSON files
+        with open(tmp_path / "model.safetensors", "w") as f:
+            f.write("fake model data")
+
+        # Should raise FileNotFoundError (config file not found)
+        with pytest.raises(FileNotFoundError, match="not found in directory"):
+            DataProcessorPipeline.from_pretrained(tmp_path, config_filename="processor.json")
+
+
+def test_processor_migration_error_creation():
+    """Test that ProcessorMigrationError is created correctly."""
+    model_path = "/path/to/model"
+    migration_command = "python migrate.py --path /path/to/model"
+    original_error = "Config not found"
+
+    error = ProcessorMigrationError(model_path, migration_command, original_error)
+
+    assert error.model_path == model_path
+    assert error.migration_command == migration_command
+    assert error.original_error == original_error
+    assert model_path in str(error)
+    assert migration_command in str(error)
+    assert original_error in str(error)
+
+
+def test_processor_migration_error_attributes():
+    """Test that ProcessorMigrationError has correct attributes."""
+    model_path = Path("/test/path")
+    migration_command = "python test.py"
+    original_error = "Test error"
+
+    error = ProcessorMigrationError(model_path, migration_command, original_error)
+
+    # Test that attributes are accessible
+    assert hasattr(error, "model_path")
+    assert hasattr(error, "migration_command")
+    assert hasattr(error, "original_error")
+
+    # Test that it's still an Exception
+    assert isinstance(error, Exception)
+
+
+def test_migration_suggestion_raises_error():
+    """Test that migration suggestion always raises ProcessorMigrationError."""
+    with pytest.raises(ProcessorMigrationError) as exc_info:
+        DataProcessorPipeline._suggest_processor_migration("/test/path", "Test error")
+
+    error = exc_info.value
+    assert "/test/path" in str(error.model_path)
+    assert "Test error" in error.original_error
+    assert "migrate_policy_normalization.py" in error.migration_command
+
+
+def test_migration_error_always_raised_for_invalid_configs():
+    """Test that ProcessorMigrationError is always raised for invalid configs."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create a model config
+        model_config = {"type": "test", "param": "value"}
+        with open(tmp_path / "config.json", "w") as f:
+            json.dump(model_config, f)
+
+        # Should always raise ProcessorMigrationError
+        with pytest.raises(ProcessorMigrationError):
+            DataProcessorPipeline.from_pretrained(tmp_path, config_filename="config.json")

tests/processor/test_observation_processor.py

@@ -0,0 +1,528 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pytest
+import torch
+
+from lerobot.configs.types import FeatureType, PipelineFeatureType
+from lerobot.processor import VanillaObservationProcessorStep
+from lerobot.processor.converters import create_transition
+from lerobot.types import TransitionKey
+from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
+from tests.conftest import assert_contract_is_typed
+
+
+def test_process_single_image():
+    """Test processing a single image."""
+    processor = VanillaObservationProcessorStep()
+
+    # Create a mock image (H, W, C) format, uint8
+    image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8)
+
+    observation = {"pixels": image}
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Check that the image was processed correctly
+    assert OBS_IMAGE in processed_obs
+    processed_img = processed_obs[OBS_IMAGE]
+
+    # Check shape: should be (1, 3, 64, 64) - batch, channels, height, width
+    assert processed_img.shape == (1, 3, 64, 64)
+
+    # Check dtype and range
+    assert processed_img.dtype == torch.float32
+    assert processed_img.min() >= 0.0
+    assert processed_img.max() <= 1.0
+
+
+def test_process_image_dict():
+    """Test processing multiple images in a dictionary."""
+    processor = VanillaObservationProcessorStep()
+
+    # Create mock images
+    image1 = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8)
+    image2 = np.random.randint(0, 256, size=(48, 48, 3), dtype=np.uint8)
+
+    observation = {"pixels": {"camera1": image1, "camera2": image2}}
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Check that both images were processed
+    assert f"{OBS_IMAGES}.camera1" in processed_obs
+    assert f"{OBS_IMAGES}.camera2" in processed_obs
+
+    # Check shapes
+    assert processed_obs[f"{OBS_IMAGES}.camera1"].shape == (1, 3, 32, 32)
+    assert processed_obs[f"{OBS_IMAGES}.camera2"].shape == (1, 3, 48, 48)
+
+
+def test_process_batched_image():
+    """Test processing already batched images."""
+    processor = VanillaObservationProcessorStep()
+
+    # Create a batched image (B, H, W, C)
+    image = np.random.randint(0, 256, size=(2, 64, 64, 3), dtype=np.uint8)
+
+    observation = {"pixels": image}
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Check that batch dimension is preserved
+    assert processed_obs[OBS_IMAGE].shape == (2, 3, 64, 64)
+
+
+def test_invalid_image_format():
+    """Test error handling for invalid image formats."""
+    processor = VanillaObservationProcessorStep()
+
+    # Test wrong channel order (channels first)
+    image = np.random.randint(0, 256, size=(3, 64, 64), dtype=np.uint8)
+    observation = {"pixels": image}
+    transition = create_transition(observation=observation)
+
+    with pytest.raises(ValueError, match="Expected channel-last images"):
+        processor(transition)
+
+
+def test_invalid_image_dtype():
+    """Test error handling for invalid image dtype."""
+    processor = VanillaObservationProcessorStep()
+
+    # Test wrong dtype
+    image = np.random.rand(64, 64, 3).astype(np.float32)
+    observation = {"pixels": image}
+    transition = create_transition(observation=observation)
+
+    with pytest.raises(ValueError, match="Expected torch.uint8 images"):
+        processor(transition)
+
+
+def test_no_pixels_in_observation():
+    """Test processor when no pixels are in observation."""
+    processor = VanillaObservationProcessorStep()
+
+    observation = {"other_data": np.array([1, 2, 3])}
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Should preserve other data unchanged
+    assert "other_data" in processed_obs
+    np.testing.assert_array_equal(processed_obs["other_data"], np.array([1, 2, 3]))
+
+
+def test_none_observation():
+    """Test processor with None observation."""
+    processor = VanillaObservationProcessorStep()
+
+    transition = create_transition(observation={})
+    result = processor(transition)
+
+    assert result == transition
+
+
+def test_serialization_methods():
+    """Test serialization methods."""
+    processor = VanillaObservationProcessorStep()
+
+    # Test get_config
+    config = processor.get_config()
+    assert isinstance(config, dict)
+
+    # Test state_dict
+    state = processor.state_dict()
+    assert isinstance(state, dict)
+
+    # Test load_state_dict (should not raise)
+    processor.load_state_dict(state)
+
+    # Test reset (should not raise)
+    processor.reset()
+
+
+def test_process_environment_state():
+    """Test processing environment_state."""
+    processor = VanillaObservationProcessorStep()
+
+    env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32)
+    observation = {"environment_state": env_state}
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Check that environment_state was renamed and processed
+    assert OBS_ENV_STATE in processed_obs
+    assert "environment_state" not in processed_obs
+
+    processed_state = processed_obs[OBS_ENV_STATE]
+    assert processed_state.shape == (1, 3)  # Batch dimension added
+    assert processed_state.dtype == torch.float32
+    torch.testing.assert_close(processed_state, torch.tensor([[1.0, 2.0, 3.0]]))
+
+
+def test_process_agent_pos():
+    """Test processing agent_pos."""
+    processor = VanillaObservationProcessorStep()
+
+    agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32)
+    observation = {"agent_pos": agent_pos}
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Check that agent_pos was renamed and processed
+    assert OBS_STATE in processed_obs
+    assert "agent_pos" not in processed_obs
+
+    processed_state = processed_obs[OBS_STATE]
+    assert processed_state.shape == (1, 3)  # Batch dimension added
+    assert processed_state.dtype == torch.float32
+    torch.testing.assert_close(processed_state, torch.tensor([[0.5, -0.5, 1.0]]))
+
+
+def test_process_batched_states():
+    """Test processing already batched states."""
+    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)
+
+    observation = {"environment_state": env_state, "agent_pos": agent_pos}
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Check that batch dimensions are preserved
+    assert processed_obs[OBS_ENV_STATE].shape == (2, 2)
+    assert processed_obs[OBS_STATE].shape == (2, 2)
+
+
+def test_process_both_states():
+    """Test processing both environment_state and agent_pos."""
+    processor = VanillaObservationProcessorStep()
+
+    env_state = np.array([1.0, 2.0], dtype=np.float32)
+    agent_pos = np.array([0.5, -0.5], dtype=np.float32)
+
+    observation = {"environment_state": env_state, "agent_pos": agent_pos, "other_data": "keep_me"}
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Check that both states were processed
+    assert OBS_ENV_STATE in processed_obs
+    assert OBS_STATE in processed_obs
+
+    # Check that original keys were removed
+    assert "environment_state" not in processed_obs
+    assert "agent_pos" not in processed_obs
+
+    # Check that other data was preserved
+    assert processed_obs["other_data"] == "keep_me"
+
+
+def test_no_states_in_observation():
+    """Test processor when no states are in observation."""
+    processor = VanillaObservationProcessorStep()
+
+    observation = {"other_data": np.array([1, 2, 3])}
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Should preserve data unchanged
+    np.testing.assert_array_equal(processed_obs, observation)
+
+
+def test_complete_observation_processing():
+    """Test processing a complete observation with both images and states."""
+    processor = VanillaObservationProcessorStep()
+
+    # Create mock data
+    image = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8)
+    env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32)
+    agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32)
+
+    observation = {
+        "pixels": image,
+        "environment_state": env_state,
+        "agent_pos": agent_pos,
+        "other_data": "preserve_me",
+    }
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Check that image was processed
+    assert OBS_IMAGE in processed_obs
+    assert processed_obs[OBS_IMAGE].shape == (1, 3, 32, 32)
+
+    # Check that states were processed
+    assert OBS_ENV_STATE in processed_obs
+    assert OBS_STATE in processed_obs
+
+    # Check that original keys were removed
+    assert "pixels" not in processed_obs
+    assert "environment_state" not in processed_obs
+    assert "agent_pos" not in processed_obs
+
+    # Check that other data was preserved
+    assert processed_obs["other_data"] == "preserve_me"
+
+
+def test_image_only_processing():
+    """Test processing observation with only images."""
+    processor = VanillaObservationProcessorStep()
+
+    image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8)
+    observation = {"pixels": image}
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    assert OBS_IMAGE in processed_obs
+    assert len(processed_obs) == 1
+
+
+def test_state_only_processing():
+    """Test processing observation with only states."""
+    processor = VanillaObservationProcessorStep()
+
+    agent_pos = np.array([1.0, 2.0], dtype=np.float32)
+    observation = {"agent_pos": agent_pos}
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    assert OBS_STATE in processed_obs
+    assert "agent_pos" not in processed_obs
+
+
+def test_empty_observation():
+    """Test processing empty observation."""
+    processor = VanillaObservationProcessorStep()
+
+    observation = {}
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    assert processed_obs == {}
+
+
+def test_equivalent_to_original_function():
+    """Test that ObservationProcessor produces equivalent results to preprocess_observation."""
+    # Import the original function for comparison
+    from lerobot.envs.utils import preprocess_observation
+
+    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)
+    env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32)
+    agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32)
+
+    observation = {"pixels": image, "environment_state": env_state, "agent_pos": agent_pos}
+
+    # Process with original function
+    original_result = preprocess_observation(observation)
+
+    # Process with new processor
+    transition = create_transition(observation=observation)
+    processor_result = processor(transition)[TransitionKey.OBSERVATION]
+
+    # Compare results
+    assert set(original_result.keys()) == set(processor_result.keys())
+
+    for key in original_result:
+        torch.testing.assert_close(original_result[key], processor_result[key])
+
+
+def test_equivalent_with_image_dict():
+    """Test equivalence with dictionary of images."""
+    from lerobot.envs.utils import preprocess_observation
+
+    processor = VanillaObservationProcessorStep()
+
+    # Create test data with multiple cameras
+    image1 = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8)
+    image2 = np.random.randint(0, 256, size=(48, 48, 3), dtype=np.uint8)
+    agent_pos = np.array([1.0, 2.0], dtype=np.float32)
+
+    observation = {"pixels": {"cam1": image1, "cam2": image2}, "agent_pos": agent_pos}
+
+    # Process with original function
+    original_result = preprocess_observation(observation)
+
+    # Process with new processor
+    transition = create_transition(observation=observation)
+    processor_result = processor(transition)[TransitionKey.OBSERVATION]
+
+    # Compare results
+    assert set(original_result.keys()) == set(processor_result.keys())
+
+    for key in original_result:
+        torch.testing.assert_close(original_result[key], processor_result[key])
+
+
+def test_image_processor_features_pixels_to_image(policy_feature_factory):
+    processor = VanillaObservationProcessorStep()
+    features = {
+        PipelineFeatureType.OBSERVATION: {
+            "pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
+            "keep": policy_feature_factory(FeatureType.ENV, (1,)),
+        },
+    }
+    out = processor.transform_features(features.copy())
+
+    assert (
+        OBS_IMAGE in out[PipelineFeatureType.OBSERVATION]
+        and out[PipelineFeatureType.OBSERVATION][OBS_IMAGE]
+        == features[PipelineFeatureType.OBSERVATION]["pixels"]
+    )
+    assert "pixels" not in out[PipelineFeatureType.OBSERVATION]
+    assert out[PipelineFeatureType.OBSERVATION]["keep"] == features[PipelineFeatureType.OBSERVATION]["keep"]
+    assert_contract_is_typed(out)
+
+
+def test_image_processor_features_observation_pixels_to_image(policy_feature_factory):
+    processor = VanillaObservationProcessorStep()
+    features = {
+        PipelineFeatureType.OBSERVATION: {
+            "observation.pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
+            "keep": policy_feature_factory(FeatureType.ENV, (1,)),
+        },
+    }
+    out = processor.transform_features(features.copy())
+
+    assert (
+        OBS_IMAGE in out[PipelineFeatureType.OBSERVATION]
+        and out[PipelineFeatureType.OBSERVATION][OBS_IMAGE]
+        == features[PipelineFeatureType.OBSERVATION]["observation.pixels"]
+    )
+    assert "observation.pixels" not in out[PipelineFeatureType.OBSERVATION]
+    assert out[PipelineFeatureType.OBSERVATION]["keep"] == features[PipelineFeatureType.OBSERVATION]["keep"]
+    assert_contract_is_typed(out)
+
+
+def test_image_processor_features_multi_camera_and_prefixed(policy_feature_factory):
+    processor = VanillaObservationProcessorStep()
+    features = {
+        PipelineFeatureType.OBSERVATION: {
+            "pixels.front": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
+            "pixels.wrist": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
+            "observation.pixels.rear": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
+            "keep": policy_feature_factory(FeatureType.ENV, (7,)),
+        },
+    }
+    out = processor.transform_features(features.copy())
+
+    assert (
+        f"{OBS_IMAGES}.front" in out[PipelineFeatureType.OBSERVATION]
+        and out[PipelineFeatureType.OBSERVATION][f"{OBS_IMAGES}.front"]
+        == features[PipelineFeatureType.OBSERVATION]["pixels.front"]
+    )
+    assert (
+        f"{OBS_IMAGES}.wrist" in out[PipelineFeatureType.OBSERVATION]
+        and out[PipelineFeatureType.OBSERVATION][f"{OBS_IMAGES}.wrist"]
+        == features[PipelineFeatureType.OBSERVATION]["pixels.wrist"]
+    )
+    assert (
+        f"{OBS_IMAGES}.rear" in out[PipelineFeatureType.OBSERVATION]
+        and out[PipelineFeatureType.OBSERVATION][f"{OBS_IMAGES}.rear"]
+        == features[PipelineFeatureType.OBSERVATION]["observation.pixels.rear"]
+    )
+    assert (
+        "pixels.front" not in out[PipelineFeatureType.OBSERVATION]
+        and "pixels.wrist" not in out[PipelineFeatureType.OBSERVATION]
+        and "observation.pixels.rear" not in out[PipelineFeatureType.OBSERVATION]
+    )
+    assert out[PipelineFeatureType.OBSERVATION]["keep"] == features[PipelineFeatureType.OBSERVATION]["keep"]
+    assert_contract_is_typed(out)
+
+
+def test_state_processor_features_environment_and_agent_pos(policy_feature_factory):
+    processor = VanillaObservationProcessorStep()
+    features = {
+        PipelineFeatureType.OBSERVATION: {
+            "environment_state": policy_feature_factory(FeatureType.STATE, (3,)),
+            "agent_pos": policy_feature_factory(FeatureType.STATE, (7,)),
+            "keep": policy_feature_factory(FeatureType.ENV, (1,)),
+        },
+    }
+    out = processor.transform_features(features.copy())
+
+    assert (
+        OBS_ENV_STATE in out[PipelineFeatureType.OBSERVATION]
+        and out[PipelineFeatureType.OBSERVATION][OBS_ENV_STATE]
+        == features[PipelineFeatureType.OBSERVATION]["environment_state"]
+    )
+    assert (
+        OBS_STATE in out[PipelineFeatureType.OBSERVATION]
+        and out[PipelineFeatureType.OBSERVATION][OBS_STATE]
+        == features[PipelineFeatureType.OBSERVATION]["agent_pos"]
+    )
+    assert (
+        "environment_state" not in out[PipelineFeatureType.OBSERVATION]
+        and "agent_pos" not in out[PipelineFeatureType.OBSERVATION]
+    )
+    assert out[PipelineFeatureType.OBSERVATION]["keep"] == features[PipelineFeatureType.OBSERVATION]["keep"]
+    assert_contract_is_typed(out)
+
+
+def test_state_processor_features_prefixed_inputs(policy_feature_factory):
+    proc = VanillaObservationProcessorStep()
+    features = {
+        PipelineFeatureType.OBSERVATION: {
+            OBS_ENV_STATE: policy_feature_factory(FeatureType.STATE, (2,)),
+            "observation.agent_pos": policy_feature_factory(FeatureType.STATE, (4,)),
+        },
+    }
+    out = proc.transform_features(features.copy())
+
+    assert (
+        OBS_ENV_STATE in out[PipelineFeatureType.OBSERVATION]
+        and out[PipelineFeatureType.OBSERVATION][OBS_ENV_STATE]
+        == features[PipelineFeatureType.OBSERVATION][OBS_ENV_STATE]
+    )
+    assert (
+        OBS_STATE in out[PipelineFeatureType.OBSERVATION]
+        and out[PipelineFeatureType.OBSERVATION][OBS_STATE]
+        == features[PipelineFeatureType.OBSERVATION]["observation.agent_pos"]
+    )
+    assert (
+        "environment_state" not in out[PipelineFeatureType.OBSERVATION]
+        and "agent_pos" not in out[PipelineFeatureType.OBSERVATION]
+    )
+    assert_contract_is_typed(out)

tests/processor/test_pi05_processor.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Compare the PI0.5 processor pipeline against the vendored OpenPI reference processors."""
+
+import os
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from lerobot.configs import FeatureType, PolicyFeature  # noqa: E402
+from lerobot.policies.pi05 import PI05Policy  # noqa: E402
+from lerobot.policies.pi05.configuration_pi05 import PI05Config  # noqa: E402
+from lerobot.policies.pi05.processor_pi05 import make_pi05_pre_post_processors  # noqa: E402
+from lerobot.utils.constants import ACTION, OBS_STATE  # noqa: E402
+from tests.policies.pi0_pi05.utils.openpi_parity import (  # noqa: E402
+    IMAGE_KEYS,
+    assert_processor_inputs_match_lerobot,
+    clone_batch,
+    make_openpi_observation_from_raw,
+    openpi_model_actions_from_raw,
+)
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
+    reason="OpenPI processor parity uses the PaliGemma tokenizer; run manually outside CI.",
+)
+
+DUMMY_ACTION_DIM = 32
+DUMMY_STATE_DIM = 32
+DUMMY_ACTION_HORIZON = 50
+DUMMY_MAX_TOKEN_LEN = 200
+DEVICE = torch.device("cpu")
+
+DUMMY_DATASET_STATS = {
+    OBS_STATE: {
+        "mean": torch.zeros(DUMMY_STATE_DIM),
+        "std": torch.ones(DUMMY_STATE_DIM),
+        "q01": torch.zeros(DUMMY_STATE_DIM),
+        "q99": torch.ones(DUMMY_STATE_DIM),
+    },
+    ACTION: {
+        "mean": torch.zeros(DUMMY_ACTION_DIM),
+        "std": torch.ones(DUMMY_ACTION_DIM),
+        "q01": torch.zeros(DUMMY_ACTION_DIM),
+        "q99": torch.ones(DUMMY_ACTION_DIM),
+    },
+    "images": {
+        key: {
+            "mean": torch.zeros(3, 224, 224),
+            "std": torch.ones(3, 224, 224),
+            "q01": torch.zeros(3, 224, 224),
+            "q99": torch.ones(3, 224, 224),
+        }
+        for key in IMAGE_KEYS
+    },
+}
+
+
+class PI05PolicyInputAdapter(torch.nn.Module):
+    """Minimal adapter exposing PI0.5 policy image preparation without loading model weights."""
+
+    _preprocess_images = PI05Policy._preprocess_images
+
+    def __init__(self, config: PI05Config) -> None:
+        super().__init__()
+        self.config = config
+        self._device_anchor = torch.nn.Parameter(torch.empty((), device=config.device), requires_grad=False)
+
+
+def create_pi05_config() -> PI05Config:
+    config = PI05Config(device=str(DEVICE))
+    config.max_state_dim = DUMMY_STATE_DIM
+    config.max_action_dim = DUMMY_ACTION_DIM
+    config.chunk_size = DUMMY_ACTION_HORIZON
+    config.n_action_steps = DUMMY_ACTION_HORIZON
+    config.tokenizer_max_length = DUMMY_MAX_TOKEN_LEN
+    config.input_features = {
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(DUMMY_STATE_DIM,)),
+        **{
+            f"observation.images.{key}": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224))
+            for key in IMAGE_KEYS
+        },
+    }
+    config.output_features = {
+        ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(DUMMY_ACTION_DIM,)),
+    }
+    return config
+
+
+def create_dummy_data() -> dict:
+    batch_size = 2
+    prompt = "Pick up the red block and place it in the bin"
+    return {
+        OBS_STATE: torch.randn(batch_size, DUMMY_STATE_DIM, dtype=torch.float32, device=DEVICE),
+        ACTION: torch.randn(
+            batch_size, DUMMY_ACTION_HORIZON, DUMMY_ACTION_DIM, dtype=torch.float32, device=DEVICE
+        ),
+        **{
+            f"observation.images.{key}": torch.rand(
+                batch_size, 3, 224, 224, dtype=torch.float32, device=DEVICE
+            )
+            for key in IMAGE_KEYS
+        },
+        "task": [prompt for _ in range(batch_size)],
+    }
+
+
+def test_pi05_processor_inputs_match_openpi_reference():
+    torch.manual_seed(0)
+    config = create_pi05_config()
+    preprocessor, _ = make_pi05_pre_post_processors(config=config, dataset_stats=DUMMY_DATASET_STATS)
+
+    raw_batch = create_dummy_data()
+    lerobot_batch = preprocessor(clone_batch(raw_batch))
+    openpi_observation = make_openpi_observation_from_raw(
+        raw_batch,
+        action_dim=DUMMY_ACTION_DIM,
+        max_token_len=DUMMY_MAX_TOKEN_LEN,
+        dataset_stats=DUMMY_DATASET_STATS,
+        pi05=True,
+    )
+
+    assert_processor_inputs_match_lerobot(
+        PI05PolicyInputAdapter(config),
+        lerobot_batch,
+        openpi_observation,
+        compare_state=False,
+    )
+    torch.testing.assert_close(
+        lerobot_batch[ACTION],
+        openpi_model_actions_from_raw(
+            raw_batch,
+            action_dim=DUMMY_ACTION_DIM,
+            dataset_stats=DUMMY_DATASET_STATS,
+            pi05=True,
+        ),
+        rtol=0,
+        atol=0,
+    )

tests/processor/test_pi0_processor.py

@@ -0,0 +1,156 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Compare the PI0 processor pipeline against the vendored OpenPI reference processors."""
+
+import os
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from lerobot.configs import FeatureType, PolicyFeature  # noqa: E402
+from lerobot.policies.pi0 import PI0Policy  # noqa: E402
+from lerobot.policies.pi0.configuration_pi0 import PI0Config  # noqa: E402
+from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors  # noqa: E402
+from lerobot.utils.constants import ACTION, OBS_STATE  # noqa: E402
+from tests.policies.pi0_pi05.utils.openpi_parity import (  # noqa: E402
+    IMAGE_KEYS,
+    assert_processor_inputs_match_lerobot,
+    clone_batch,
+    make_openpi_observation_from_raw,
+    openpi_model_actions_from_raw,
+)
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
+    reason="OpenPI processor parity uses the PaliGemma tokenizer; run manually outside CI.",
+)
+
+DUMMY_ACTION_DIM = 32
+DUMMY_STATE_DIM = 32
+DUMMY_ACTION_HORIZON = 50
+DUMMY_MAX_TOKEN_LEN = 48
+DEVICE = torch.device("cpu")
+
+DUMMY_DATASET_STATS = {
+    OBS_STATE: {
+        "mean": torch.zeros(DUMMY_STATE_DIM),
+        "std": torch.ones(DUMMY_STATE_DIM),
+        "q01": torch.zeros(DUMMY_STATE_DIM),
+        "q99": torch.ones(DUMMY_STATE_DIM),
+    },
+    ACTION: {
+        "mean": torch.zeros(DUMMY_ACTION_DIM),
+        "std": torch.ones(DUMMY_ACTION_DIM),
+        "q01": torch.zeros(DUMMY_ACTION_DIM),
+        "q99": torch.ones(DUMMY_ACTION_DIM),
+    },
+    "images": {
+        key: {
+            "mean": torch.zeros(3, 224, 224),
+            "std": torch.ones(3, 224, 224),
+            "q01": torch.zeros(3, 224, 224),
+            "q99": torch.ones(3, 224, 224),
+        }
+        for key in IMAGE_KEYS
+    },
+}
+
+
+class PI0PolicyInputAdapter(torch.nn.Module):
+    """Minimal adapter exposing PI0 policy input-preparation helpers without loading model weights."""
+
+    _preprocess_images = PI0Policy._preprocess_images
+    prepare_state = PI0Policy.prepare_state
+
+    def __init__(self, config: PI0Config) -> None:
+        super().__init__()
+        self.config = config
+        self._device_anchor = torch.nn.Parameter(torch.empty((), device=config.device), requires_grad=False)
+
+
+def create_pi0_config() -> PI0Config:
+    config = PI0Config(device=str(DEVICE))
+    config.max_state_dim = DUMMY_STATE_DIM
+    config.max_action_dim = DUMMY_ACTION_DIM
+    config.chunk_size = DUMMY_ACTION_HORIZON
+    config.n_action_steps = DUMMY_ACTION_HORIZON
+    config.tokenizer_max_length = DUMMY_MAX_TOKEN_LEN
+    config.input_features = {
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(DUMMY_STATE_DIM,)),
+        **{
+            f"observation.images.{key}": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224))
+            for key in IMAGE_KEYS
+        },
+    }
+    config.output_features = {
+        ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(DUMMY_ACTION_DIM,)),
+    }
+    return config
+
+
+def create_dummy_data() -> dict:
+    batch_size = 2
+    prompt = "Pick up the red block and place it in the bin"
+    return {
+        OBS_STATE: torch.randn(batch_size, DUMMY_STATE_DIM, dtype=torch.float32, device=DEVICE),
+        ACTION: torch.randn(
+            batch_size, DUMMY_ACTION_HORIZON, DUMMY_ACTION_DIM, dtype=torch.float32, device=DEVICE
+        ),
+        **{
+            f"observation.images.{key}": torch.rand(
+                batch_size, 3, 224, 224, dtype=torch.float32, device=DEVICE
+            )
+            for key in IMAGE_KEYS
+        },
+        "task": [prompt for _ in range(batch_size)],
+    }
+
+
+def test_pi0_processor_inputs_match_openpi_reference():
+    torch.manual_seed(0)
+    config = create_pi0_config()
+    preprocessor, _ = make_pi0_pre_post_processors(config=config, dataset_stats=DUMMY_DATASET_STATS)
+
+    raw_batch = create_dummy_data()
+    lerobot_batch = preprocessor(clone_batch(raw_batch))
+    openpi_observation = make_openpi_observation_from_raw(
+        raw_batch,
+        action_dim=DUMMY_ACTION_DIM,
+        max_token_len=DUMMY_MAX_TOKEN_LEN,
+        dataset_stats=DUMMY_DATASET_STATS,
+        pi05=False,
+    )
+
+    assert_processor_inputs_match_lerobot(
+        PI0PolicyInputAdapter(config),
+        lerobot_batch,
+        openpi_observation,
+        compare_state=True,
+    )
+    torch.testing.assert_close(
+        lerobot_batch[ACTION],
+        openpi_model_actions_from_raw(
+            raw_batch,
+            action_dim=DUMMY_ACTION_DIM,
+            dataset_stats=DUMMY_DATASET_STATS,
+            pi05=False,
+        ),
+        rtol=0,
+        atol=0,
+    )

tests/processor/test_pipeline_from_pretrained_helpers.py

@@ -0,0 +1,259 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Tests for DataProcessorPipeline.from_pretrained helper methods.
+
+These tests focus on the individual private methods that were extracted from
+the main from_pretrained method to improve modularity and testability.
+"""
+
+import json
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from lerobot.processor.pipeline import DataProcessorPipeline, ProcessorMigrationError
+
+# Simplified Config Loading Tests
+
+
+def test_load_config_directory():
+    """Test loading config from directory."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create a config file
+        config_file = tmp_path / "processor.json"
+        test_config = {"name": "TestProcessor", "steps": []}
+        config_file.write_text(json.dumps(test_config))
+
+        # Load from directory
+        loaded_config, base_path = DataProcessorPipeline._load_config(str(tmp_path), "processor.json", {})
+
+        assert loaded_config == test_config
+        assert base_path == tmp_path
+
+
+def test_load_config_single_file():
+    """Test loading config from a single file path."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create a config file
+        config_file = tmp_path / "processor.json"
+        test_config = {"name": "TestProcessor", "steps": []}
+        config_file.write_text(json.dumps(test_config))
+
+        # Load using file path directly
+        loaded_config, base_path = DataProcessorPipeline._load_config(
+            str(config_file), "any_filename_ignored", {}
+        )
+
+        assert loaded_config == test_config
+        assert base_path == tmp_path
+
+
+def test_load_config_directory_file_not_found():
+    """Test directory loading when config file doesn't exist."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Directory exists but no processor.json
+        with pytest.raises(FileNotFoundError, match="not found in directory"):
+            DataProcessorPipeline._load_config(str(tmp_path), "processor.json", {})
+
+
+def test_load_config_directory_with_migration_detection():
+    """Test that missing config triggers migration detection."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create old-style config to trigger migration
+        (tmp_path / "config.json").write_text(json.dumps({"type": "act"}))
+
+        # Try to load processor.json (doesn't exist), should trigger migration
+        with pytest.raises(ProcessorMigrationError):
+            DataProcessorPipeline._load_config(str(tmp_path), "processor.json", {})
+
+
+def test_load_config_nonexistent_path_tries_hub():
+    """Test that nonexistent paths try Hub (simplified logic)."""
+    # This path doesn't exist locally, should try Hub
+    with pytest.raises(FileNotFoundError, match="on the HuggingFace Hub"):
+        DataProcessorPipeline._load_config("nonexistent/path", "processor.json", {})
+
+
+# Config Validation Tests
+
+
+def test_validate_loaded_config_valid_config():
+    """Test validation with valid processor config."""
+    valid_config = {"name": "TestProcessor", "steps": []}
+
+    # Should not raise any exception
+    DataProcessorPipeline._validate_loaded_config("any-path", valid_config, "processor.json")
+
+
+def test_validate_loaded_config_invalid_config():
+    """Test validation with invalid processor config."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Create non-processor config to trigger migration
+        (tmp_path / "config.json").write_text(json.dumps({"type": "act"}))
+
+        invalid_config = {"type": "act", "hidden_dim": 256}
+
+        with pytest.raises(ProcessorMigrationError):
+            DataProcessorPipeline._validate_loaded_config(str(tmp_path), invalid_config, "config.json")
+
+
+def test_validate_loaded_config_invalid_config_no_migration():
+    """Test validation with invalid config when no migration is detected."""
+    # Non-directory path (Hub repo) - no migration detection
+    invalid_config = {"type": "act", "hidden_dim": 256}
+
+    with pytest.raises(ValueError, match="not a valid processor configuration"):
+        DataProcessorPipeline._validate_loaded_config("user/repo", invalid_config, "config.json")
+
+
+# Step Class Resolution Tests
+
+
+def test_resolve_step_class_registry_name():
+    """Test resolution using registry name."""
+    from lerobot.processor.pipeline import ProcessorStep, ProcessorStepRegistry
+
+    # Register a test step
+    @ProcessorStepRegistry.register("test_step")
+    class TestStep(ProcessorStep):
+        def __call__(self, transition):
+            return transition
+
+        def transform_features(self, features):
+            return features
+
+    try:
+        step_entry = {"registry_name": "test_step"}
+        step_class, step_key = DataProcessorPipeline._resolve_step_class(step_entry)
+
+        assert step_class is TestStep
+        assert step_key == "test_step"
+    finally:
+        ProcessorStepRegistry.unregister("test_step")
+
+
+def test_resolve_step_class_registry_name_not_found():
+    """Test resolution with non-existent registry name."""
+    step_entry = {"registry_name": "nonexistent_step"}
+
+    with pytest.raises(ImportError, match="Failed to load processor step from registry"):
+        DataProcessorPipeline._resolve_step_class(step_entry)
+
+
+def test_resolve_step_class_import_path():
+    """Test resolution using full import path."""
+    # Use a valid existing class (this should work)
+    step_entry = {"class": "lerobot.processor.pipeline.ProcessorStep"}
+
+    # This should succeed - ProcessorStep can be imported, just not instantiated
+    step_class, step_key = DataProcessorPipeline._resolve_step_class(step_entry)
+
+    from lerobot.processor.pipeline import ProcessorStep
+
+    assert step_class is ProcessorStep
+    assert step_key == "ProcessorStep"
+
+
+def test_resolve_step_class_invalid_import_path():
+    """Test resolution with invalid import path."""
+    step_entry = {"class": "nonexistent.module.ClassName"}
+
+    with pytest.raises(ImportError, match="Failed to load processor step"):
+        DataProcessorPipeline._resolve_step_class(step_entry)
+
+
+# Override Validation Tests
+
+
+def test_validate_overrides_used_all_used():
+    """Test validation when all overrides are used."""
+    # Empty set means all overrides were used
+    remaining_overrides = set()
+    config = {"steps": [{"class": "SomeStep"}]}
+
+    # Should not raise
+    DataProcessorPipeline._validate_overrides_used(remaining_overrides, config)
+
+
+def test_validate_overrides_used_some_unused():
+    """Test validation when some overrides are unused."""
+    remaining_overrides = {"NonExistentStep", "AnotherMissingStep"}
+    config = {
+        "steps": [
+            {"registry_name": "normalize_step"},
+            {"class": "some.module.TransformStep"},
+        ]
+    }
+
+    with pytest.raises(KeyError, match="Override keys.*do not match any step"):
+        DataProcessorPipeline._validate_overrides_used(remaining_overrides, config)
+
+
+def test_validate_overrides_used_helpful_error_message():
+    """Test that error message includes available step keys."""
+    remaining_overrides = {"WrongStep"}
+    config = {
+        "steps": [
+            {"registry_name": "correct_step"},
+            {"class": "module.path.CorrectClass"},
+        ]
+    }
+
+    with pytest.raises(KeyError) as exc_info:
+        DataProcessorPipeline._validate_overrides_used(remaining_overrides, config)
+
+    error_msg = str(exc_info.value)
+    assert "Available step keys" in error_msg
+    assert "correct_step" in error_msg
+    assert "CorrectClass" in error_msg
+
+
+# Integration Tests for Simplified Logic
+
+
+def test_simplified_three_way_loading():
+    """Test that the simplified 3-way loading logic works correctly."""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+
+        # Test 1: Directory loading
+        config_file = tmp_path / "processor.json"
+        test_config = {"name": "DirectoryTest", "steps": []}
+        config_file.write_text(json.dumps(test_config))
+
+        loaded_config, base_path = DataProcessorPipeline._load_config(str(tmp_path), "processor.json", {})
+        assert loaded_config["name"] == "DirectoryTest"
+        assert base_path == tmp_path
+
+        # Test 2: Single file loading
+        loaded_config, base_path = DataProcessorPipeline._load_config(
+            str(config_file), "ignored_filename", {}
+        )
+        assert loaded_config["name"] == "DirectoryTest"
+        assert base_path == tmp_path

tests/processor/test_policy_robot_bridge.py

@@ -0,0 +1,526 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tempfile
+from pathlib import Path
+
+import pytest
+import torch
+
+from lerobot.configs.types import FeatureType, PipelineFeatureType
+from lerobot.processor import (
+    DataProcessorPipeline,
+    PolicyActionToRobotActionProcessorStep,
+    ProcessorStepRegistry,
+    RobotActionToPolicyActionProcessorStep,
+)
+from lerobot.processor.converters import identity_transition
+from lerobot.utils.constants import ACTION
+from tests.conftest import assert_contract_is_typed
+
+
+def test_robot_to_policy_basic_action_conversion():
+    """Test basic robot action to policy action conversion."""
+    motor_names = ["joint1", "joint2", "joint3"]
+    processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
+
+    robot_action = {
+        "joint1.pos": 1.0,
+        "joint2.pos": 2.0,
+        "joint3.pos": 3.0,
+    }
+
+    policy_action = processor.action(robot_action)
+
+    assert isinstance(policy_action, torch.Tensor)
+    assert policy_action.shape == (3,)
+    torch.testing.assert_close(policy_action, torch.tensor([1.0, 2.0, 3.0]))
+
+
+def test_robot_to_policy_action_conversion_preserves_order():
+    """Test that motor names order is preserved in conversion."""
+    motor_names = ["gripper", "arm", "wrist"]
+    processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
+
+    robot_action = {
+        "arm.pos": 10.0,
+        "gripper.pos": 5.0,
+        "wrist.pos": 15.0,
+    }
+
+    policy_action = processor.action(robot_action)
+
+    expected = torch.tensor([5.0, 10.0, 15.0])
+    torch.testing.assert_close(policy_action, expected)
+
+
+def test_robot_to_policy_action_conversion_with_floats_and_tensors():
+    """Test conversion with mixed float and tensor values."""
+    motor_names = ["joint1", "joint2"]
+    processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
+
+    robot_action = {
+        "joint1.pos": torch.tensor(1.5),
+        "joint2.pos": 2.5,  # Regular float
+    }
+
+    policy_action = processor.action(robot_action)
+
+    assert isinstance(policy_action, torch.Tensor)
+    torch.testing.assert_close(policy_action, torch.tensor([1.5, 2.5]))
+
+
+def test_robot_to_policy_action_length_mismatch_error():
+    """Test error when robot action length doesn't match motor names."""
+    motor_names = ["joint1", "joint2", "joint3"]
+    processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
+
+    # Too few actions
+    robot_action = {"joint1.pos": 1.0, "joint2.pos": 2.0}
+
+    with pytest.raises(ValueError, match="Action must have 3 elements, got 2"):
+        processor.action(robot_action)
+
+    robot_action = {
+        "joint1.pos": 1.0,
+        "joint2.pos": 2.0,
+        "joint3.pos": 3.0,
+        "extra.pos": 4.0,
+    }
+
+    with pytest.raises(ValueError, match="Action must have 3 elements, got 4"):
+        processor.action(robot_action)
+
+
+def test_robot_to_policy_missing_motor_key_error():
+    """Test error when robot action is missing expected motor keys."""
+    motor_names = ["joint1", "joint2"]
+    processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
+
+    robot_action = {
+        "joint1.pos": 1.0,
+        "wrong_key.pos": 2.0,
+    }
+
+    with pytest.raises(KeyError):
+        processor.action(robot_action)
+
+
+def test_robot_to_policy_transform_features():
+    """Test feature transformation for robot to policy action processor."""
+    motor_names = ["joint1", "joint2", "joint3"]
+    processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
+
+    features = {
+        PipelineFeatureType.ACTION: {
+            "joint1.pos": {"type": FeatureType.ACTION, "shape": (1,)},
+            "joint2.pos": {"type": FeatureType.ACTION, "shape": (1,)},
+            "joint3.pos": {"type": FeatureType.ACTION, "shape": (1,)},
+            "other_data": {"type": FeatureType.ENV, "shape": (1,)},
+        }
+    }
+
+    transformed = processor.transform_features(features)
+
+    assert ACTION in transformed[PipelineFeatureType.ACTION]
+    action_feature = transformed[PipelineFeatureType.ACTION][ACTION]
+    assert action_feature.type == FeatureType.ACTION
+    assert action_feature.shape == (3,)
+
+    assert "joint1.pos" in transformed[PipelineFeatureType.ACTION]
+    assert "joint2.pos" in transformed[PipelineFeatureType.ACTION]
+    assert "joint3.pos" in transformed[PipelineFeatureType.ACTION]
+
+    assert "other_data" in transformed[PipelineFeatureType.ACTION]
+
+
+def test_robot_to_policy_get_config():
+    """Test configuration serialization."""
+    motor_names = ["motor1", "motor2"]
+    processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
+
+    config = processor.get_config()
+    assert config == {"motor_names": motor_names}
+
+
+def test_robot_to_policy_state_dict():
+    """Test state dict operations."""
+    processor = RobotActionToPolicyActionProcessorStep(motor_names=["joint1"])
+
+    state = processor.state_dict()
+    assert state == {}
+
+    processor.load_state_dict({})
+
+
+def test_robot_to_policy_single_motor():
+    """Test with single motor."""
+    processor = RobotActionToPolicyActionProcessorStep(motor_names=["single_joint"])
+
+    robot_action = {"single_joint.pos": 42.0}
+    policy_action = processor.action(robot_action)
+
+    assert policy_action.shape == (1,)
+    torch.testing.assert_close(policy_action, torch.tensor([42.0]))
+
+
+def test_policy_to_robot_basic_action_conversion():
+    """Test basic policy action to robot action conversion."""
+    motor_names = ["joint1", "joint2", "joint3"]
+    processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
+
+    policy_action = torch.tensor([1.0, 2.0, 3.0])
+    robot_action = processor.action(policy_action)
+
+    assert isinstance(robot_action, dict)
+    assert len(robot_action) == 3
+
+    expected = {
+        "joint1.pos": 1.0,
+        "joint2.pos": 2.0,
+        "joint3.pos": 3.0,
+    }
+
+    for key, expected_value in expected.items():
+        assert key in robot_action
+        actual_value = robot_action[key]
+        if isinstance(actual_value, torch.Tensor):
+            actual_value = actual_value.item()
+        assert actual_value == pytest.approx(expected_value)
+
+
+def test_policy_to_robot_action_conversion_preserves_order():
+    """Test that motor names order corresponds to tensor indices."""
+    motor_names = ["gripper", "arm", "wrist"]
+    processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
+
+    policy_action = torch.tensor([5.0, 10.0, 15.0])
+    robot_action = processor.action(policy_action)
+
+    assert robot_action["gripper.pos"] == pytest.approx(5.0)
+    assert robot_action["arm.pos"] == pytest.approx(10.0)
+    assert robot_action["wrist.pos"] == pytest.approx(15.0)
+
+
+def test_policy_to_robot_action_conversion_with_numpy_input():
+    """Test conversion with numpy array input."""
+    import numpy as np
+
+    motor_names = ["joint1", "joint2"]
+    processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
+
+    policy_action = np.array([1.5, 2.5])
+    robot_action = processor.action(policy_action)
+
+    assert robot_action["joint1.pos"] == pytest.approx(1.5)
+    assert robot_action["joint2.pos"] == pytest.approx(2.5)
+
+
+def test_policy_to_robot_action_length_mismatch_error():
+    """Test error when policy action length doesn't match motor names."""
+    motor_names = ["joint1", "joint2", "joint3"]
+    processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
+
+    policy_action = torch.tensor([1.0, 2.0])
+
+    with pytest.raises(ValueError, match="Action must have 3 elements, got 2"):
+        processor.action(policy_action)
+
+    policy_action = torch.tensor([1.0, 2.0, 3.0, 4.0])
+
+    with pytest.raises(ValueError, match="Action must have 3 elements, got 4"):
+        processor.action(policy_action)
+
+
+def test_policy_to_robot_transform_features():
+    """Test feature transformation for policy to robot action processor."""
+    motor_names = ["joint1", "joint2"]
+    processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
+
+    features = {
+        PipelineFeatureType.ACTION: {
+            ACTION: {"type": FeatureType.ACTION, "shape": (2,)},
+            "other_data": {"type": FeatureType.ENV, "shape": (1,)},
+        }
+    }
+
+    transformed = processor.transform_features(features)
+
+    assert "joint1.pos" in transformed[PipelineFeatureType.ACTION]
+    assert "joint2.pos" in transformed[PipelineFeatureType.ACTION]
+
+    for motor in motor_names:
+        motor_feature = transformed[PipelineFeatureType.ACTION][f"{motor}.pos"]
+        assert motor_feature.type == FeatureType.ACTION
+        assert motor_feature.shape == (1,)
+
+    assert ACTION in transformed[PipelineFeatureType.ACTION]
+
+    assert "other_data" in transformed[PipelineFeatureType.ACTION]
+
+
+def test_policy_to_robot_get_config():
+    """Test configuration serialization."""
+    motor_names = ["motor1", "motor2"]
+    processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
+
+    config = processor.get_config()
+    assert config == {"motor_names": motor_names}
+
+
+def test_policy_to_robot_state_dict():
+    """Test state dict operations."""
+    processor = PolicyActionToRobotActionProcessorStep(motor_names=["joint1"])
+
+    state = processor.state_dict()
+    assert state == {}
+
+    processor.load_state_dict({})
+
+
+def test_policy_to_robot_single_motor():
+    """Test with single motor."""
+    processor = PolicyActionToRobotActionProcessorStep(motor_names=["single_joint"])
+
+    policy_action = torch.tensor([42.0])
+    robot_action = processor.action(policy_action)
+
+    assert len(robot_action) == 1
+    assert robot_action["single_joint.pos"] == pytest.approx(42.0)
+
+
+def test_robot_to_policy_registry():
+    """Test RobotActionToPolicyActionProcessorStep registry."""
+    assert "robot_action_to_policy_action_processor" in ProcessorStepRegistry.list()
+
+    retrieved_class = ProcessorStepRegistry.get("robot_action_to_policy_action_processor")
+    assert retrieved_class is RobotActionToPolicyActionProcessorStep
+
+    instance = retrieved_class(motor_names=["test"])
+    assert isinstance(instance, RobotActionToPolicyActionProcessorStep)
+    assert instance.motor_names == ["test"]
+
+
+def test_policy_to_robot_registry():
+    """Test PolicyActionToRobotActionProcessorStep registry."""
+    assert "policy_action_to_robot_action_processor" in ProcessorStepRegistry.list()
+
+    retrieved_class = ProcessorStepRegistry.get("policy_action_to_robot_action_processor")
+    assert retrieved_class is PolicyActionToRobotActionProcessorStep
+
+    instance = retrieved_class(motor_names=["test"])
+    assert isinstance(instance, PolicyActionToRobotActionProcessorStep)
+    assert instance.motor_names == ["test"]
+
+
+def test_save_and_load_robot_to_policy():
+    """Test saving and loading RobotActionToPolicyActionProcessorStep."""
+    motor_names = ["joint1", "joint2", "joint3"]
+    processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
+    pipeline = DataProcessorPipeline([processor], name="TestRobotToPolicy")
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        # Save pipeline
+        pipeline.save_pretrained(tmp_dir)
+
+        # Check config file exists
+        config_path = Path(tmp_dir) / "testrobottopolicy.json"
+        assert config_path.exists()
+
+        # Load pipeline
+        loaded_pipeline = DataProcessorPipeline.from_pretrained(
+            tmp_dir,
+            "testrobottopolicy.json",
+            to_transition=identity_transition,
+            to_output=identity_transition,
+        )
+
+        assert loaded_pipeline.name == "TestRobotToPolicy"
+        assert len(loaded_pipeline) == 1
+
+        # Check loaded processor
+        loaded_processor = loaded_pipeline.steps[0]
+        assert isinstance(loaded_processor, RobotActionToPolicyActionProcessorStep)
+        assert loaded_processor.motor_names == motor_names
+
+        # Test functionality after loading
+        robot_action = {"joint1.pos": 1.0, "joint2.pos": 2.0, "joint3.pos": 3.0}
+        policy_action = loaded_processor.action(robot_action)
+        torch.testing.assert_close(policy_action, torch.tensor([1.0, 2.0, 3.0]))
+
+
+def test_save_and_load_policy_to_robot():
+    """Test saving and loading PolicyActionToRobotActionProcessorStep."""
+    motor_names = ["motor_a", "motor_b"]
+    processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
+    pipeline = DataProcessorPipeline([processor], name="TestPolicyToRobot")
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        # Save pipeline
+        pipeline.save_pretrained(tmp_dir)
+
+        # Load pipeline
+        loaded_pipeline = DataProcessorPipeline.from_pretrained(
+            tmp_dir,
+            "testpolicytorobot.json",
+            to_transition=identity_transition,
+            to_output=identity_transition,
+        )
+
+        loaded_processor = loaded_pipeline.steps[0]
+        assert isinstance(loaded_processor, PolicyActionToRobotActionProcessorStep)
+        assert loaded_processor.motor_names == motor_names
+
+        policy_action = torch.tensor([10.0, 20.0])
+        robot_action = loaded_processor.action(policy_action)
+        assert robot_action["motor_a.pos"] == pytest.approx(10.0)
+        assert robot_action["motor_b.pos"] == pytest.approx(20.0)
+
+
+# Integration and chaining tests
+
+
+def test_round_trip_conversion():
+    """Test that robot->policy->robot conversion preserves values."""
+    motor_names = ["joint1", "joint2", "joint3"]
+    robot_to_policy = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
+    policy_to_robot = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
+
+    original_robot_action = {
+        "joint1.pos": 1.5,
+        "joint2.pos": -2.3,
+        "joint3.pos": 0.7,
+    }
+
+    policy_action = robot_to_policy.action(original_robot_action)
+    final_robot_action = policy_to_robot.action(policy_action)
+
+    for key in original_robot_action:
+        original_val = original_robot_action[key]
+        final_val = final_robot_action[key]
+        if isinstance(final_val, torch.Tensor):
+            final_val = final_val.item()
+        assert final_val == pytest.approx(original_val, abs=1e-6)
+
+
+def test_chained_processors_in_pipeline():
+    """Test both processors chained in a pipeline."""
+    motor_names = ["joint1", "joint2"]
+    robot_to_policy = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
+    policy_to_robot = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
+
+    pipeline = DataProcessorPipeline(
+        [robot_to_policy, policy_to_robot],
+        to_transition=identity_transition,
+        to_output=identity_transition,
+    )
+
+    assert len(pipeline.steps) == 2
+    assert isinstance(pipeline.steps[0], RobotActionToPolicyActionProcessorStep)
+    assert isinstance(pipeline.steps[1], PolicyActionToRobotActionProcessorStep)
+
+
+def test_robot_to_policy_features_contract(policy_feature_factory):
+    """Test feature transformation maintains proper typing contract."""
+    processor = RobotActionToPolicyActionProcessorStep(motor_names=["j1", "j2"])
+    features = {
+        PipelineFeatureType.ACTION: {
+            "j1.pos": policy_feature_factory(FeatureType.ACTION, (1,)),
+            "j2.pos": policy_feature_factory(FeatureType.ACTION, (1,)),
+            "other": policy_feature_factory(FeatureType.ENV, (3,)),
+        }
+    }
+
+    out = processor.transform_features(features.copy())
+
+    assert_contract_is_typed(out)
+
+    assert ACTION in out[PipelineFeatureType.ACTION]
+    action_feature = out[PipelineFeatureType.ACTION][ACTION]
+    assert action_feature.type == FeatureType.ACTION
+    assert action_feature.shape == (2,)
+
+
+def test_policy_to_robot_features_contract(policy_feature_factory):
+    """Test feature transformation maintains proper typing contract."""
+    processor = PolicyActionToRobotActionProcessorStep(motor_names=["m1", "m2", "m3"])
+    features = {
+        PipelineFeatureType.ACTION: {
+            ACTION: policy_feature_factory(FeatureType.ACTION, (3,)),
+            "other": policy_feature_factory(FeatureType.ENV, (1,)),
+        }
+    }
+
+    out = processor.transform_features(features.copy())
+
+    assert_contract_is_typed(out)
+
+    for motor in ["m1", "m2", "m3"]:
+        key = f"{motor}.pos"
+        assert key in out[PipelineFeatureType.ACTION]
+        motor_feature = out[PipelineFeatureType.ACTION][key]
+        assert motor_feature.type == FeatureType.ACTION
+        assert motor_feature.shape == (1,)
+
+
+def test_empty_motor_names_list():
+    """Test behavior with empty motor names list."""
+    processor = RobotActionToPolicyActionProcessorStep(motor_names=[])
+
+    robot_action = {}
+    policy_action = processor.action(robot_action)
+
+    assert isinstance(policy_action, torch.Tensor)
+    assert policy_action.shape == (0,)
+
+
+def test_empty_motor_names_list_policy_to_robot():
+    """Test PolicyActionToRobotActionProcessorStep with empty motor names."""
+    processor = PolicyActionToRobotActionProcessorStep(motor_names=[])
+
+    policy_action = torch.tensor([])
+    robot_action = processor.action(policy_action)
+
+    assert isinstance(robot_action, dict)
+    assert len(robot_action) == 0
+
+
+def test_very_long_motor_names():
+    """Test with many motor names."""
+    motor_names = [f"joint_{i}" for i in range(100)]
+    processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
+
+    robot_action = {f"joint_{i}.pos": float(i) for i in range(100)}
+    policy_action = processor.action(robot_action)
+
+    assert policy_action.shape == (100,)
+    expected = torch.tensor([float(i) for i in range(100)])
+    torch.testing.assert_close(policy_action, expected)
+
+
+def test_special_characters_in_motor_names():
+    """Test with special characters in motor names."""
+    motor_names = ["motor-1", "motor_2", "motor.3"]
+    processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
+
+    robot_action = {
+        "motor-1.pos": 1.0,
+        "motor_2.pos": 2.0,
+        "motor.3.pos": 3.0,
+    }
+
+    policy_action = processor.action(robot_action)
+    torch.testing.assert_close(policy_action, torch.tensor([1.0, 2.0, 3.0]))

tests/processor/test_rename_processor.py

@@ -0,0 +1,498 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from lerobot.configs.types import FeatureType, PipelineFeatureType
+from lerobot.processor import (
+    DataProcessorPipeline,
+    ProcessorStepRegistry,
+    RenameObservationsProcessorStep,
+    TransitionKey,
+)
+from lerobot.processor.converters import create_transition, identity_transition
+from lerobot.processor.rename_processor import rename_stats
+from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_IMAGES, OBS_STATE
+from tests.conftest import assert_contract_is_typed
+
+
+def test_basic_renaming():
+    """Test basic key renaming functionality."""
+    rename_map = {
+        "old_key1": "new_key1",
+        "old_key2": "new_key2",
+    }
+    processor = RenameObservationsProcessorStep(rename_map=rename_map)
+
+    observation = {
+        "old_key1": torch.tensor([1.0, 2.0]),
+        "old_key2": np.array([3.0, 4.0]),
+        "unchanged_key": "keep_me",
+    }
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Check renamed keys
+    assert "new_key1" in processed_obs
+    assert "new_key2" in processed_obs
+    assert "old_key1" not in processed_obs
+    assert "old_key2" not in processed_obs
+
+    # Check values are preserved
+    torch.testing.assert_close(processed_obs["new_key1"], torch.tensor([1.0, 2.0]))
+    np.testing.assert_array_equal(processed_obs["new_key2"], np.array([3.0, 4.0]))
+
+    # Check unchanged key is preserved
+    assert processed_obs["unchanged_key"] == "keep_me"
+
+
+def test_empty_rename_map():
+    """Test processor with empty rename map (should pass through unchanged)."""
+    processor = RenameObservationsProcessorStep(rename_map={})
+
+    observation = {
+        "key1": torch.tensor([1.0]),
+        "key2": "value2",
+    }
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # All keys should be unchanged
+    assert processed_obs.keys() == observation.keys()
+    torch.testing.assert_close(processed_obs["key1"], observation["key1"])
+    assert processed_obs["key2"] == observation["key2"]
+
+
+def test_none_observation():
+    """Test processor with None observation."""
+    processor = RenameObservationsProcessorStep(rename_map={"old": "new"})
+
+    transition = create_transition(observation={})
+    result = processor(transition)
+
+    # Should return transition unchanged
+    assert result == transition
+
+
+def test_overlapping_rename():
+    """Test renaming when new names might conflict."""
+    rename_map = {
+        "a": "b",
+        "b": "c",  # This creates a potential conflict
+    }
+    processor = RenameObservationsProcessorStep(rename_map=rename_map)
+
+    observation = {
+        "a": 1,
+        "b": 2,
+        "x": 3,
+    }
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Check that renaming happens correctly
+    assert "a" not in processed_obs
+    assert processed_obs["b"] == 1  # 'a' renamed to 'b'
+    assert processed_obs["c"] == 2  # original 'b' renamed to 'c'
+    assert processed_obs["x"] == 3
+
+
+def test_partial_rename():
+    """Test renaming only some keys."""
+    rename_map = {
+        OBS_STATE: "observation.proprio_state",
+        "pixels": OBS_IMAGE,
+    }
+    processor = RenameObservationsProcessorStep(rename_map=rename_map)
+
+    observation = {
+        OBS_STATE: torch.randn(10),
+        "pixels": np.random.randint(0, 256, (64, 64, 3), dtype=np.uint8),
+        "reward": 1.0,
+        "info": {"episode": 1},
+    }
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Check renamed keys
+    assert "observation.proprio_state" in processed_obs
+    assert OBS_IMAGE in processed_obs
+    assert OBS_STATE not in processed_obs
+    assert "pixels" not in processed_obs
+
+    # Check unchanged keys
+    assert processed_obs["reward"] == 1.0
+    assert processed_obs["info"] == {"episode": 1}
+
+
+def test_get_config():
+    """Test configuration serialization."""
+    rename_map = {
+        "old1": "new1",
+        "old2": "new2",
+    }
+    processor = RenameObservationsProcessorStep(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 RenameProcessorStep)."""
+    processor = RenameObservationsProcessorStep(rename_map={"old": "new"})
+
+    state = processor.state_dict()
+    assert state == {}
+
+    # Load state dict should work even with empty dict
+    processor.load_state_dict({})
+
+
+def test_integration_with_robot_processor():
+    """Test integration with RobotProcessor pipeline."""
+    rename_map = {
+        "agent_pos": OBS_STATE,
+        "pixels": OBS_IMAGE,
+    }
+    rename_processor = RenameObservationsProcessorStep(rename_map=rename_map)
+
+    pipeline = DataProcessorPipeline(
+        [rename_processor], to_transition=identity_transition, to_output=identity_transition
+    )
+
+    observation = {
+        "agent_pos": np.array([1.0, 2.0, 3.0]),
+        "pixels": np.zeros((32, 32, 3), dtype=np.uint8),
+        "other_data": "preserve_me",
+    }
+    transition = create_transition(
+        observation=observation, reward=0.5, done=False, truncated=False, info={}, complementary_data={}
+    )
+
+    result = pipeline(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Check renaming worked through pipeline
+    assert OBS_STATE in processed_obs
+    assert OBS_IMAGE in processed_obs
+    assert "agent_pos" not in processed_obs
+    assert "pixels" not in processed_obs
+    assert processed_obs["other_data"] == "preserve_me"
+
+    # Check other transition elements unchanged
+    assert result[TransitionKey.REWARD] == 0.5
+    assert result[TransitionKey.DONE] is False
+
+
+def test_save_and_load_pretrained():
+    """Test saving and loading processor with RobotProcessor."""
+    rename_map = {
+        "old_state": OBS_STATE,
+        "old_image": OBS_IMAGE,
+    }
+    processor = RenameObservationsProcessorStep(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) / "testrenameprocessorstep.json"
+        )  # Based on name="TestRenameProcessorStep"
+        assert config_path.exists()
+
+        # No state files should be created for RenameProcessorStep
+        state_files = list(Path(tmp_dir).glob("*.safetensors"))
+        assert len(state_files) == 0
+
+        # Load pipeline
+        loaded_pipeline = DataProcessorPipeline.from_pretrained(
+            tmp_dir,
+            config_filename="testrenameprocessorstep.json",
+            to_transition=identity_transition,
+            to_output=identity_transition,
+        )
+
+        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, RenameObservationsProcessorStep)
+        assert loaded_processor.rename_map == rename_map
+
+        # Test functionality after loading
+        observation = {"old_state": [1, 2, 3], "old_image": "image_data"}
+        transition = create_transition(observation=observation)
+
+        result = loaded_pipeline(transition)
+        processed_obs = result[TransitionKey.OBSERVATION]
+
+        assert OBS_STATE in processed_obs
+        assert OBS_IMAGE in processed_obs
+        assert processed_obs[OBS_STATE] == [1, 2, 3]
+        assert processed_obs[OBS_IMAGE] == "image_data"
+
+
+def test_registry_functionality():
+    """Test that RenameProcessorStep is properly registered."""
+    # Check that it's registered
+    assert "rename_observations_processor" in ProcessorStepRegistry.list()
+
+    # Get from registry
+    retrieved_class = ProcessorStepRegistry.get("rename_observations_processor")
+    assert retrieved_class is RenameObservationsProcessorStep
+
+    # Create instance from registry
+    instance = retrieved_class(rename_map={"old": "new"})
+    assert isinstance(instance, RenameObservationsProcessorStep)
+    assert instance.rename_map == {"old": "new"}
+
+
+def test_registry_based_save_load():
+    """Test save/load using registry name instead of module path."""
+    processor = RenameObservationsProcessorStep(rename_map={"key1": "renamed_key1"})
+    pipeline = DataProcessorPipeline(
+        [processor], to_transition=identity_transition, to_output=identity_transition
+    )
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        # Save and load
+        pipeline.save_pretrained(tmp_dir)
+
+        # Verify config uses registry name
+        import json
+
+        with open(Path(tmp_dir) / "dataprocessorpipeline.json") as f:  # Default name is "RobotProcessor"
+            config = json.load(f)
+
+        assert "registry_name" in config["steps"][0]
+        assert config["steps"][0]["registry_name"] == "rename_observations_processor"
+        assert "class" not in config["steps"][0]  # Should use registry, not module path
+
+        # Load should work
+        loaded_pipeline = DataProcessorPipeline.from_pretrained(
+            tmp_dir, config_filename="dataprocessorpipeline.json"
+        )
+        loaded_processor = loaded_pipeline.steps[0]
+        assert isinstance(loaded_processor, RenameObservationsProcessorStep)
+        assert loaded_processor.rename_map == {"key1": "renamed_key1"}
+
+
+def test_chained_rename_processors():
+    """Test multiple RenameProcessorSteps in a pipeline."""
+    # First processor: rename raw keys to intermediate format
+    processor1 = RenameObservationsProcessorStep(
+        rename_map={
+            "pos": "agent_position",
+            "img": "camera_image",
+        }
+    )
+
+    # Second processor: rename to final format
+    processor2 = RenameObservationsProcessorStep(
+        rename_map={
+            "agent_position": OBS_STATE,
+            "camera_image": OBS_IMAGE,
+        }
+    )
+
+    pipeline = DataProcessorPipeline(
+        [processor1, processor2], to_transition=identity_transition, to_output=identity_transition
+    )
+
+    observation = {
+        "pos": np.array([1.0, 2.0]),
+        "img": "image_data",
+        "extra": "keep_me",
+    }
+    transition = create_transition(observation=observation)
+
+    # Step through to see intermediate results
+    results = list(pipeline.step_through(transition))
+
+    # After first processor
+    assert "agent_position" in results[1][TransitionKey.OBSERVATION]
+    assert "camera_image" in results[1][TransitionKey.OBSERVATION]
+
+    # After second processor
+    final_obs = results[2][TransitionKey.OBSERVATION]
+    assert OBS_STATE in final_obs
+    assert OBS_IMAGE in final_obs
+    assert final_obs["extra"] == "keep_me"
+
+    # Original keys should be gone
+    assert "pos" not in final_obs
+    assert "img" not in final_obs
+    assert "agent_position" not in final_obs
+    assert "camera_image" not in final_obs
+
+
+def test_nested_observation_rename():
+    """Test renaming with nested observation structures."""
+    rename_map = {
+        f"{OBS_IMAGES}.left": "observation.camera.left_view",
+        f"{OBS_IMAGES}.right": "observation.camera.right_view",
+        "observation.proprio": "observation.proprioception",
+    }
+    processor = RenameObservationsProcessorStep(rename_map=rename_map)
+
+    observation = {
+        f"{OBS_IMAGES}.left": torch.randn(3, 64, 64),
+        f"{OBS_IMAGES}.right": torch.randn(3, 64, 64),
+        "observation.proprio": torch.randn(7),
+        "observation.gripper": torch.tensor([0.0]),  # Not renamed
+    }
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Check renames
+    assert "observation.camera.left_view" in processed_obs
+    assert "observation.camera.right_view" in processed_obs
+    assert "observation.proprioception" in processed_obs
+
+    # Check unchanged key
+    assert "observation.gripper" in processed_obs
+
+    # Check old keys removed
+    assert f"{OBS_IMAGES}.left" not in processed_obs
+    assert f"{OBS_IMAGES}.right" not in processed_obs
+    assert "observation.proprio" not in processed_obs
+
+
+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 = RenameObservationsProcessorStep(rename_map=rename_map)
+
+    tensor_value = torch.randn(3, 3)
+    array_value = np.random.rand(2, 2)
+
+    observation = {
+        "old_tensor": tensor_value,
+        "old_array": array_value,
+        "old_scalar": 42,
+        "old_string": "hello",
+        "old_dict": {"nested": "value"},
+        "old_list": [1, 2, 3],
+    }
+    transition = create_transition(observation=observation)
+
+    result = processor(transition)
+    processed_obs = result[TransitionKey.OBSERVATION]
+
+    # Check that values and types are preserved
+    assert torch.equal(processed_obs["new_tensor"], tensor_value)
+    assert np.array_equal(processed_obs["new_array"], array_value)
+    assert processed_obs["new_scalar"] == 42
+    assert processed_obs["old_string"] == "hello"
+    assert processed_obs["old_dict"] == {"nested": "value"}
+    assert processed_obs["old_list"] == [1, 2, 3]
+
+
+def test_features_basic_renaming(policy_feature_factory):
+    processor = RenameObservationsProcessorStep(rename_map={"a": "x", "b": "y"})
+    features = {
+        PipelineFeatureType.OBSERVATION: {
+            "a": policy_feature_factory(FeatureType.VISUAL, (2,)),
+            "b": policy_feature_factory(FeatureType.VISUAL, (3,)),
+            "c": policy_feature_factory(FeatureType.VISUAL, (1,)),
+        },
+    }
+
+    out = processor.transform_features(features.copy())
+
+    # Values preserved and typed
+    assert out[PipelineFeatureType.OBSERVATION]["x"] == features[PipelineFeatureType.OBSERVATION]["a"]
+    assert out[PipelineFeatureType.OBSERVATION]["y"] == features[PipelineFeatureType.OBSERVATION]["b"]
+    assert out[PipelineFeatureType.OBSERVATION]["c"] == features[PipelineFeatureType.OBSERVATION]["c"]
+
+    assert_contract_is_typed(out)
+    # Input not mutated
+    assert set(features[PipelineFeatureType.OBSERVATION]) == {"a", "b", "c"}
+
+
+def test_features_overlapping_keys(policy_feature_factory):
+    # Overlapping renames: both 'a' and 'b' exist. 'a'->'b', 'b'->'c'
+    processor = RenameObservationsProcessorStep(rename_map={"a": "b", "b": "c"})
+    features = {
+        PipelineFeatureType.OBSERVATION: {
+            "a": policy_feature_factory(FeatureType.VISUAL, (1,)),
+            "b": policy_feature_factory(FeatureType.VISUAL, (2,)),
+        },
+    }
+    out = processor.transform_features(features)
+
+    assert set(out[PipelineFeatureType.OBSERVATION]) == {"b", "c"}
+    assert (
+        out[PipelineFeatureType.OBSERVATION]["b"] == features[PipelineFeatureType.OBSERVATION]["a"]
+    )  # 'a' renamed to'b'
+    assert (
+        out[PipelineFeatureType.OBSERVATION]["c"] == features[PipelineFeatureType.OBSERVATION]["b"]
+    )  # 'b' renamed to 'c'
+    assert_contract_is_typed(out)
+
+
+def test_features_chained_processors(policy_feature_factory):
+    # Chain two rename processors at the contract level
+    processor1 = RenameObservationsProcessorStep(rename_map={"pos": "agent_position", "img": "camera_image"})
+    processor2 = RenameObservationsProcessorStep(
+        rename_map={"agent_position": OBS_STATE, "camera_image": OBS_IMAGE}
+    )
+    pipeline = DataProcessorPipeline([processor1, processor2])
+
+    spec = {
+        PipelineFeatureType.OBSERVATION: {
+            "pos": policy_feature_factory(FeatureType.VISUAL, (7,)),
+            "img": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
+            "extra": policy_feature_factory(FeatureType.VISUAL, (1,)),
+        },
+    }
+    out = pipeline.transform_features(initial_features=spec)
+
+    assert set(out[PipelineFeatureType.OBSERVATION]) == {OBS_STATE, OBS_IMAGE, "extra"}
+    assert out[PipelineFeatureType.OBSERVATION][OBS_STATE] == spec[PipelineFeatureType.OBSERVATION]["pos"]
+    assert out[PipelineFeatureType.OBSERVATION][OBS_IMAGE] == spec[PipelineFeatureType.OBSERVATION]["img"]
+    assert out[PipelineFeatureType.OBSERVATION]["extra"] == spec[PipelineFeatureType.OBSERVATION]["extra"]
+    assert_contract_is_typed(out)
+
+
+def test_rename_stats_basic():
+    orig = {
+        OBS_STATE: {"mean": np.array([0.0]), "std": np.array([1.0])},
+        ACTION: {"mean": np.array([0.0])},
+    }
+    mapping = {OBS_STATE: "observation.robot_state"}
+    renamed = rename_stats(orig, mapping)
+    assert "observation.robot_state" in renamed and OBS_STATE not in renamed
+    # Ensure deep copy: mutate original and verify renamed unaffected
+    orig[OBS_STATE]["mean"][0] = 42.0
+    assert renamed["observation.robot_state"]["mean"][0] != 42.0

tests/processor/test_render_messages_processor.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python
+
+import pytest
+
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+
+import torch  # noqa: E402
+
+from lerobot.configs.recipe import MessageTurn, TrainingRecipe  # noqa: E402
+from lerobot.processor.converters import create_transition  # noqa: E402
+from lerobot.processor.render_messages_processor import RenderMessagesStep  # noqa: E402
+from lerobot.types import TransitionKey  # noqa: E402
+
+
+def test_render_messages_step_noops_without_language_columns():
+    recipe = TrainingRecipe(
+        messages=[
+            MessageTurn(role="user", content="${task}", stream="high_level"),
+            MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True),
+        ]
+    )
+    transition = create_transition(complementary_data={"task": "do it"})
+
+    assert RenderMessagesStep(recipe)(transition) == transition
+
+
+def test_render_messages_step_renders_and_drops_raw_language():
+    recipe = TrainingRecipe(
+        messages=[
+            MessageTurn(role="user", content="${task}", stream="high_level"),
+            MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True),
+        ]
+    )
+    transition = create_transition(
+        complementary_data={
+            "task": "do it",
+            "timestamp": torch.tensor(0.0),
+            "index": torch.tensor(7),
+            "language_persistent": [
+                {
+                    "role": "assistant",
+                    "content": "reach carefully",
+                    "style": "subtask",
+                    "timestamp": 0.0,
+                    "camera": None,
+                    "tool_calls": None,
+                }
+            ],
+            "language_events": [],
+        }
+    )
+
+    out = RenderMessagesStep(recipe)(transition)
+    data = out[TransitionKey.COMPLEMENTARY_DATA]
+
+    assert "language_persistent" not in data
+    assert "language_events" not in data
+    assert data["messages"][-1]["content"] == "reach carefully"
+    assert data["message_streams"] == ["high_level", "low_level"]
+    assert data["target_message_indices"] == [1]

tests/processor/test_smolvla_processor.py

@@ -0,0 +1,457 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for SmolVLA policy processor."""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from lerobot.configs.types import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature
+from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
+from lerobot.policies.smolvla.processor_smolvla import make_smolvla_pre_post_processors
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    EnvTransition,
+    NewLineTaskProcessorStep,
+    NormalizerProcessorStep,
+    ProcessorStep,
+    RenameObservationsProcessorStep,
+    TransitionKey,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import create_transition, transition_to_batch
+from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE
+
+
+class MockTokenizerProcessorStep(ProcessorStep):
+    """Mock tokenizer processor step for testing."""
+
+    def __init__(self, *args, **kwargs):
+        # Accept any arguments to mimic the real TokenizerProcessorStep interface
+        pass
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        # Pass through transition unchanged
+        return transition
+
+    def transform_features(self, features):
+        # Pass through features unchanged
+        return features
+
+
+def create_default_config():
+    """Create a default SmolVLA configuration for testing."""
+    config = SmolVLAConfig()
+    config.input_features = {
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,)),
+        OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
+    }
+    config.output_features = {
+        ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
+    }
+    config.normalization_mapping = {
+        FeatureType.STATE: NormalizationMode.MEAN_STD,
+        FeatureType.VISUAL: NormalizationMode.IDENTITY,
+        FeatureType.ACTION: NormalizationMode.MIN_MAX,
+    }
+    config.device = "cpu"
+    config.vlm_model_name = "HuggingFaceTB/SmolVLM-Instruct"
+    config.pad_language_to = "max_length"
+    config.tokenizer_max_length = 100
+    return config
+
+
+def create_default_stats():
+    """Create default dataset statistics for testing."""
+    return {
+        OBS_STATE: {"mean": torch.zeros(8), "std": torch.ones(8)},
+        OBS_IMAGE: {},  # No normalization for images
+        ACTION: {"min": torch.full((7,), -1.0), "max": torch.ones(7)},
+    }
+
+
+def test_make_smolvla_processor_basic():
+    """Test basic creation of SmolVLA processor."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    with patch(
+        "lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessorStep
+    ):
+        preprocessor, postprocessor = make_smolvla_pre_post_processors(
+            config,
+            stats,
+        )
+
+    # Check processor names
+    assert preprocessor.name == "policy_preprocessor"
+    assert postprocessor.name == "policy_postprocessor"
+
+    # Check steps in preprocessor
+    assert len(preprocessor.steps) == 6
+    assert isinstance(preprocessor.steps[0], RenameObservationsProcessorStep)
+    assert isinstance(preprocessor.steps[1], AddBatchDimensionProcessorStep)
+    assert isinstance(preprocessor.steps[2], NewLineTaskProcessorStep)
+    # Step 3 would be TokenizerProcessorStep but it's mocked
+    assert isinstance(preprocessor.steps[4], DeviceProcessorStep)
+    assert isinstance(preprocessor.steps[5], NormalizerProcessorStep)
+
+    # Check steps in postprocessor
+    assert len(postprocessor.steps) == 2
+    assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
+    assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
+
+
+def test_smolvla_newline_processor_single_task():
+    """Test NewLineTaskProcessorStep with single task string."""
+    processor = NewLineTaskProcessorStep()
+
+    # Test with task that doesn't have newline
+    transition = create_transition(complementary_data={"task": "test task"})
+    result = processor(transition)
+    assert result[TransitionKey.COMPLEMENTARY_DATA]["task"] == "test task\n"
+
+    # Test with task that already has newline
+    transition = create_transition(complementary_data={"task": "test task\n"})
+    result = processor(transition)
+    assert result[TransitionKey.COMPLEMENTARY_DATA]["task"] == "test task\n"
+
+
+def test_smolvla_newline_processor_list_of_tasks():
+    """Test NewLineTaskProcessorStep with list of task strings."""
+    processor = NewLineTaskProcessorStep()
+
+    # Test with list of tasks
+    tasks = ["task1", "task2\n", "task3"]
+    transition = create_transition(complementary_data={"task": tasks})
+    result = processor(transition)
+    expected = ["task1\n", "task2\n", "task3\n"]
+    assert result[TransitionKey.COMPLEMENTARY_DATA]["task"] == expected
+
+
+def test_smolvla_newline_processor_empty_transition():
+    """Test NewLineTaskProcessorStep with empty transition."""
+    processor = NewLineTaskProcessorStep()
+
+    # Test with no complementary_data
+    transition = create_transition()
+    result = processor(transition)
+    assert result == transition
+
+    # Test with complementary_data but no task
+    transition = create_transition(complementary_data={"other": "data"})
+    result = processor(transition)
+    assert result == transition
+
+    # Test with None task
+    transition = create_transition(complementary_data={"task": None})
+    result = processor(transition)
+    assert result == transition
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_smolvla_processor_cuda():
+    """Test SmolVLA processor with CUDA device."""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    # Mock the tokenizer processor to act as pass-through
+    class MockTokenizerProcessorStep(ProcessorStep):
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def __call__(self, transition):
+            return transition
+
+        def state_dict(self):
+            return {}
+
+        def load_state_dict(self, state):
+            pass
+
+        def reset(self):
+            pass
+
+        def get_config(self):
+            return {"tokenizer_name": "HuggingFaceTB/SmolVLM-Instruct"}
+
+        def transform_features(self, features):
+            return features
+
+    with patch(
+        "lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessorStep
+    ):
+        preprocessor, postprocessor = make_smolvla_pre_post_processors(
+            config,
+            stats,
+        )
+
+    # Create CPU data
+    observation = {
+        OBS_STATE: torch.randn(8),
+        OBS_IMAGE: torch.randn(3, 224, 224),
+    }
+    action = torch.randn(7)
+    transition = create_transition(observation, action, complementary_data={"task": "test task"})
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data is on CUDA
+    assert processed[OBS_STATE].device.type == "cuda"
+    assert processed[OBS_IMAGE].device.type == "cuda"
+    assert processed[TransitionKey.ACTION.value].device.type == "cuda"
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_smolvla_processor_accelerate_scenario():
+    """Test SmolVLA processor in simulated Accelerate scenario."""
+    config = create_default_config()
+    config.device = "cuda:0"
+    stats = create_default_stats()
+
+    # Mock the tokenizer processor to act as pass-through
+    class MockTokenizerProcessorStep(ProcessorStep):
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def __call__(self, transition):
+            return transition
+
+        def state_dict(self):
+            return {}
+
+        def load_state_dict(self, state):
+            pass
+
+        def reset(self):
+            pass
+
+        def get_config(self):
+            return {"tokenizer_name": "HuggingFaceTB/SmolVLM-Instruct"}
+
+        def transform_features(self, features):
+            return features
+
+    with patch(
+        "lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessorStep
+    ):
+        preprocessor, postprocessor = make_smolvla_pre_post_processors(
+            config,
+            stats,
+        )
+
+    # Simulate Accelerate: data already on GPU and batched
+    device = torch.device("cuda:0")
+    observation = {
+        OBS_STATE: torch.randn(1, 8).to(device),
+        OBS_IMAGE: torch.randn(1, 3, 224, 224).to(device),
+    }
+    action = torch.randn(1, 7).to(device)
+    transition = create_transition(observation, action, complementary_data={"task": ["test task"]})
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data stays on same GPU
+    assert processed[OBS_STATE].device == device
+    assert processed[OBS_IMAGE].device == device
+    assert processed[TransitionKey.ACTION.value].device == device
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs")
+def test_smolvla_processor_multi_gpu():
+    """Test SmolVLA processor with multi-GPU setup."""
+    config = create_default_config()
+    config.device = "cuda:0"
+    stats = create_default_stats()
+
+    # Mock the tokenizer processor to act as pass-through
+    class MockTokenizerProcessorStep(ProcessorStep):
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def __call__(self, transition):
+            return transition
+
+        def state_dict(self):
+            return {}
+
+        def load_state_dict(self, state):
+            pass
+
+        def reset(self):
+            pass
+
+        def get_config(self):
+            return {"tokenizer_name": "HuggingFaceTB/SmolVLM-Instruct"}
+
+        def transform_features(self, features):
+            return features
+
+    with patch(
+        "lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessorStep
+    ):
+        preprocessor, postprocessor = make_smolvla_pre_post_processors(
+            config,
+            stats,
+        )
+
+    # Simulate data on different GPU
+    device = torch.device("cuda:1")
+    observation = {
+        OBS_STATE: torch.randn(1, 8).to(device),
+        OBS_IMAGE: torch.randn(1, 3, 224, 224).to(device),
+    }
+    action = torch.randn(1, 7).to(device)
+    transition = create_transition(observation, action, complementary_data={"task": ["test task"]})
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data stays on cuda:1
+    assert processed[OBS_STATE].device == device
+    assert processed[OBS_IMAGE].device == device
+    assert processed[TransitionKey.ACTION.value].device == device
+
+
+def test_smolvla_processor_without_stats():
+    """Test SmolVLA processor creation without dataset statistics."""
+    config = create_default_config()
+
+    # Mock the tokenizer processor
+    with patch(
+        "lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessorStep
+    ):
+        preprocessor, postprocessor = make_smolvla_pre_post_processors(
+            config,
+            dataset_stats=None,
+        )
+
+    # Should still create processors
+    assert preprocessor is not None
+    assert postprocessor is not None
+
+
+def test_smolvla_newline_processor_state_dict():
+    """Test NewLineTaskProcessorStep state dict methods."""
+    processor = NewLineTaskProcessorStep()
+
+    # Test state_dict (should be empty)
+    state = processor.state_dict()
+    assert state == {}
+
+    # Test load_state_dict (should do nothing)
+    processor.load_state_dict({})
+
+    # Test reset (should do nothing)
+    processor.reset()
+
+    # Test get_config
+    config = processor.get_config()
+    assert config == {}
+
+
+def test_smolvla_newline_processor_transform_features():
+    """Test NewLineTaskProcessorStep transform_features method."""
+    processor = NewLineTaskProcessorStep()
+
+    # Test transform_features
+    features = {
+        PipelineFeatureType.OBSERVATION: {OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(10,))},
+    }
+    result = processor.transform_features(features)
+    assert result == features  # Should return unchanged
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_smolvla_processor_bfloat16_device_float32_normalizer():
+    """Test: DeviceProcessor(bfloat16) + NormalizerProcessor(float32) → output bfloat16 via automatic adaptation"""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    with patch(
+        "lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessorStep
+    ):
+        preprocessor, _ = make_smolvla_pre_post_processors(
+            config,
+            stats,
+        )
+
+    # Modify the pipeline to use bfloat16 device processor with float32 normalizer
+    modified_steps = []
+    for step in preprocessor.steps:
+        if isinstance(step, DeviceProcessorStep):
+            # Device processor converts to bfloat16
+            modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="bfloat16"))
+        elif isinstance(step, NormalizerProcessorStep):
+            # Normalizer stays configured as float32 (will auto-adapt to bfloat16)
+            modified_steps.append(
+                NormalizerProcessorStep(
+                    features=step.features,
+                    norm_map=step.norm_map,
+                    stats=step.stats,
+                    device=config.device,
+                    dtype=torch.float32,  # Deliberately configured as float32
+                )
+            )
+        else:
+            modified_steps.append(step)
+    preprocessor.steps = modified_steps
+
+    # Verify initial normalizer configuration (SmolVLA has NormalizerProcessorStep at index 5)
+    normalizer_step = preprocessor.steps[5]  # NormalizerProcessorStep
+    assert normalizer_step.dtype == torch.float32
+
+    # Create test data with both state and visual observations
+    observation = {
+        OBS_STATE: torch.randn(8, dtype=torch.float32),
+        OBS_IMAGE: torch.randn(3, 224, 224, dtype=torch.float32),
+    }
+    action = torch.randn(7, dtype=torch.float32)
+    transition = create_transition(
+        observation, action, complementary_data={"task": "test bfloat16 adaptation"}
+    )
+
+    batch = transition_to_batch(transition)
+
+    # Process through full pipeline
+    processed = preprocessor(batch)
+
+    # Verify: DeviceProcessor → bfloat16, NormalizerProcessor adapts → final output is bfloat16
+    assert processed[OBS_STATE].dtype == torch.bfloat16
+    assert processed[OBS_IMAGE].dtype == torch.bfloat16  # IDENTITY normalization still gets dtype conversion
+    assert processed[TransitionKey.ACTION.value].dtype == torch.bfloat16
+
+    # Verify normalizer automatically adapted its internal state
+    assert normalizer_step.dtype == torch.bfloat16
+    # Check state stats (has normalization)
+    for stat_tensor in normalizer_step._tensor_stats[OBS_STATE].values():
+        assert stat_tensor.dtype == torch.bfloat16
+    # OBS_IMAGE uses IDENTITY normalization, so no stats to check

tests/processor/test_tdmpc_processor.py

@@ -0,0 +1,467 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for TDMPC policy processor."""
+
+import tempfile
+
+import pytest
+import torch
+
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+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,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    RenameObservationsProcessorStep,
+    TransitionKey,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import create_transition, transition_to_batch
+from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE
+
+
+def create_default_config():
+    """Create a default TDMPC configuration for testing."""
+    config = TDMPCConfig()
+    config.input_features = {
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(12,)),
+        OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
+    }
+    config.output_features = {
+        ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(6,)),
+    }
+    config.normalization_mapping = {
+        FeatureType.STATE: NormalizationMode.MEAN_STD,
+        FeatureType.VISUAL: NormalizationMode.IDENTITY,
+        FeatureType.ACTION: NormalizationMode.MIN_MAX,
+    }
+    config.device = "cpu"
+    return config
+
+
+def create_default_stats():
+    """Create default dataset statistics for testing."""
+    return {
+        OBS_STATE: {"mean": torch.zeros(12), "std": torch.ones(12)},
+        OBS_IMAGE: {},  # No normalization for images
+        ACTION: {"min": torch.full((6,), -1.0), "max": torch.ones(6)},
+    }
+
+
+def test_make_tdmpc_processor_basic():
+    """Test basic creation of TDMPC processor."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Check processor names
+    assert preprocessor.name == "policy_preprocessor"
+    assert postprocessor.name == "policy_postprocessor"
+
+    # Check steps in preprocessor
+    assert len(preprocessor.steps) == 4
+    assert isinstance(preprocessor.steps[0], RenameObservationsProcessorStep)
+    assert isinstance(preprocessor.steps[1], AddBatchDimensionProcessorStep)
+    assert isinstance(preprocessor.steps[2], DeviceProcessorStep)
+    assert isinstance(preprocessor.steps[3], NormalizerProcessorStep)
+
+    # Check steps in postprocessor
+    assert len(postprocessor.steps) == 2
+    assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
+    assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
+
+
+def test_tdmpc_processor_normalization():
+    """Test that TDMPC processor correctly normalizes and unnormalizes data."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Create test data
+    observation = {
+        OBS_STATE: torch.randn(12),
+        OBS_IMAGE: torch.randn(3, 224, 224),
+    }
+    action = torch.randn(6)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data is processed and batched
+    assert processed[OBS_STATE].shape == (1, 12)
+    assert processed[OBS_IMAGE].shape == (1, 3, 224, 224)
+    assert processed[TransitionKey.ACTION.value].shape == (1, 6)
+
+    # Process action through postprocessor
+    postprocessed = postprocessor(processed[TransitionKey.ACTION.value])
+
+    # Check that action is unnormalized (but still batched)
+    assert postprocessed.shape == (1, 6)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_tdmpc_processor_cuda():
+    """Test TDMPC processor with CUDA device."""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Create CPU data
+    observation = {
+        OBS_STATE: torch.randn(12),
+        OBS_IMAGE: torch.randn(3, 224, 224),
+    }
+    action = torch.randn(6)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data is on CUDA
+    assert processed[OBS_STATE].device.type == "cuda"
+    assert processed[OBS_IMAGE].device.type == "cuda"
+    assert processed[TransitionKey.ACTION.value].device.type == "cuda"
+
+    # Process through postprocessor
+    postprocessed = postprocessor(processed[TransitionKey.ACTION.value])
+
+    # Check that action is back on CPU
+    assert postprocessed.device.type == "cpu"
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_tdmpc_processor_accelerate_scenario():
+    """Test TDMPC processor in simulated Accelerate scenario."""
+    config = create_default_config()
+    config.device = "cuda:0"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Simulate Accelerate: data already on GPU
+    device = torch.device("cuda:0")
+    observation = {
+        OBS_STATE: torch.randn(12).to(device),
+        OBS_IMAGE: torch.randn(3, 224, 224).to(device),
+    }
+    action = torch.randn(6).to(device)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data stays on same GPU
+    assert processed[OBS_STATE].device == device
+    assert processed[OBS_IMAGE].device == device
+    assert processed[TransitionKey.ACTION.value].device == device
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs")
+def test_tdmpc_processor_multi_gpu():
+    """Test TDMPC processor with multi-GPU setup."""
+    config = create_default_config()
+    config.device = "cuda:0"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Simulate data on different GPU
+    device = torch.device("cuda:1")
+    observation = {
+        OBS_STATE: torch.randn(12).to(device),
+        OBS_IMAGE: torch.randn(3, 224, 224).to(device),
+    }
+    action = torch.randn(6).to(device)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data stays on cuda:1
+    assert processed[OBS_STATE].device == device
+    assert processed[OBS_IMAGE].device == device
+    assert processed[TransitionKey.ACTION.value].device == device
+
+
+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)
+
+    # Should still create processors
+    assert preprocessor is not None
+    assert postprocessor is not None
+
+    # Process should still work
+    observation = {
+        OBS_STATE: torch.randn(12),
+        OBS_IMAGE: torch.randn(3, 224, 224),
+    }
+    action = torch.randn(6)
+    transition = create_transition(observation, action)
+    batch = transition_to_batch(transition)
+
+    processed = preprocessor(batch)
+    assert processed is not None
+
+
+def test_tdmpc_processor_save_and_load():
+    """Test saving and loading TDMPC processor."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Save preprocessor
+        preprocessor.save_pretrained(tmpdir)
+
+        # Load preprocessor
+        loaded_preprocessor = DataProcessorPipeline.from_pretrained(
+            tmpdir, config_filename="policy_preprocessor.json"
+        )
+
+        # Test that loaded processor works
+        observation = {
+            OBS_STATE: torch.randn(12),
+            OBS_IMAGE: torch.randn(3, 224, 224),
+        }
+        action = torch.randn(6)
+        transition = create_transition(observation, action)
+
+        batch = transition_to_batch(transition)
+        processed = loaded_preprocessor(batch)
+        assert processed[OBS_STATE].shape == (1, 12)
+        assert processed[OBS_IMAGE].shape == (1, 3, 224, 224)
+        assert processed[TransitionKey.ACTION.value].shape == (1, 6)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_tdmpc_processor_mixed_precision():
+    """Test TDMPC processor with mixed precision."""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    # Create processor
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Replace DeviceProcessorStep with one that uses float16
+    modified_steps = []
+    for step in preprocessor.steps:
+        if isinstance(step, DeviceProcessorStep):
+            modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16"))
+        elif isinstance(step, NormalizerProcessorStep):
+            # Update normalizer to use the same device as the device processor
+            modified_steps.append(
+                NormalizerProcessorStep(
+                    features=step.features,
+                    norm_map=step.norm_map,
+                    stats=step.stats,
+                    device=config.device,
+                    dtype=torch.float16,  # Match the float16 dtype
+                )
+            )
+        else:
+            modified_steps.append(step)
+    preprocessor.steps = modified_steps
+
+    # Create test data
+    observation = {
+        OBS_STATE: torch.randn(12, dtype=torch.float32),
+        OBS_IMAGE: torch.randn(3, 224, 224, dtype=torch.float32),
+    }
+    action = torch.randn(6, dtype=torch.float32)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data is converted to float16
+    assert processed[OBS_STATE].dtype == torch.float16
+    assert processed[OBS_IMAGE].dtype == torch.float16
+    assert processed[TransitionKey.ACTION.value].dtype == torch.float16
+
+
+def test_tdmpc_processor_batch_data():
+    """Test TDMPC processor with batched data."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Test with batched data
+    batch_size = 64
+    observation = {
+        OBS_STATE: torch.randn(batch_size, 12),
+        OBS_IMAGE: torch.randn(batch_size, 3, 224, 224),
+    }
+    action = torch.randn(batch_size, 6)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that batch dimension is preserved
+    assert processed[OBS_STATE].shape == (batch_size, 12)
+    assert processed[OBS_IMAGE].shape == (batch_size, 3, 224, 224)
+    assert processed[TransitionKey.ACTION.value].shape == (batch_size, 6)
+
+
+def test_tdmpc_processor_edge_cases():
+    """Test TDMPC processor with edge cases."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Test with only state observation (no image)
+    observation = {OBS_STATE: torch.randn(12)}
+    action = torch.randn(6)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    processed = preprocessor(batch)
+    assert processed[OBS_STATE].shape == (1, 12)
+    assert OBS_IMAGE not in processed
+
+    # Test with only image observation (no state)
+    observation = {OBS_IMAGE: torch.randn(3, 224, 224)}
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    processed = preprocessor(batch)
+    assert processed[OBS_IMAGE].shape == (1, 3, 224, 224)
+    assert OBS_STATE not in processed
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_tdmpc_processor_bfloat16_device_float32_normalizer():
+    """Test: DeviceProcessor(bfloat16) + NormalizerProcessor(float32) → output bfloat16 via automatic adaptation"""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    preprocessor, _ = make_tdmpc_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Modify the pipeline to use bfloat16 device processor with float32 normalizer
+    modified_steps = []
+    for step in preprocessor.steps:
+        if isinstance(step, DeviceProcessorStep):
+            # Device processor converts to bfloat16
+            modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="bfloat16"))
+        elif isinstance(step, NormalizerProcessorStep):
+            # Normalizer stays configured as float32 (will auto-adapt to bfloat16)
+            modified_steps.append(
+                NormalizerProcessorStep(
+                    features=step.features,
+                    norm_map=step.norm_map,
+                    stats=step.stats,
+                    device=config.device,
+                    dtype=torch.float32,  # Deliberately configured as float32
+                )
+            )
+        else:
+            modified_steps.append(step)
+    preprocessor.steps = modified_steps
+
+    # Verify initial normalizer configuration
+    normalizer_step = preprocessor.steps[3]  # NormalizerProcessorStep
+    assert normalizer_step.dtype == torch.float32
+
+    # Create test data with both state and visual observations
+    observation = {
+        OBS_STATE: torch.randn(12, dtype=torch.float32),
+        OBS_IMAGE: torch.randn(3, 224, 224, dtype=torch.float32),
+    }
+    action = torch.randn(6, dtype=torch.float32)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through full pipeline
+    processed = preprocessor(batch)
+
+    # Verify: DeviceProcessor → bfloat16, NormalizerProcessor adapts → final output is bfloat16
+    assert processed[OBS_STATE].dtype == torch.bfloat16
+    assert processed[OBS_IMAGE].dtype == torch.bfloat16  # IDENTITY normalization still gets dtype conversion
+    assert processed[TransitionKey.ACTION.value].dtype == torch.bfloat16
+
+    # Verify normalizer automatically adapted its internal state
+    assert normalizer_step.dtype == torch.bfloat16
+    # Check state stats (has normalization)
+    for stat_tensor in normalizer_step._tensor_stats[OBS_STATE].values():
+        assert stat_tensor.dtype == torch.bfloat16
+    # OBS_IMAGE uses IDENTITY normalization, so no stats to check

tests/processor/test_vqbet_processor.py

@@ -0,0 +1,462 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for VQBeT policy processor."""
+
+import tempfile
+
+import pytest
+import torch
+
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+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,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    RenameObservationsProcessorStep,
+    TransitionKey,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import create_transition, transition_to_batch
+from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE
+
+
+def create_default_config():
+    """Create a default VQBeT configuration for testing."""
+    config = VQBeTConfig()
+    config.input_features = {
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,)),
+        OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
+    }
+    config.output_features = {
+        ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
+    }
+    config.normalization_mapping = {
+        FeatureType.STATE: NormalizationMode.MEAN_STD,
+        FeatureType.VISUAL: NormalizationMode.IDENTITY,
+        FeatureType.ACTION: NormalizationMode.MIN_MAX,
+    }
+    config.device = "cpu"
+    return config
+
+
+def create_default_stats():
+    """Create default dataset statistics for testing."""
+    return {
+        OBS_STATE: {"mean": torch.zeros(8), "std": torch.ones(8)},
+        OBS_IMAGE: {},  # No normalization for images
+        ACTION: {"min": torch.full((7,), -1.0), "max": torch.ones(7)},
+    }
+
+
+def test_make_vqbet_processor_basic():
+    """Test basic creation of VQBeT processor."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Check processor names
+    assert preprocessor.name == "policy_preprocessor"
+    assert postprocessor.name == "policy_postprocessor"
+
+    # Check steps in preprocessor
+    assert len(preprocessor.steps) == 4
+    assert isinstance(preprocessor.steps[0], RenameObservationsProcessorStep)
+    assert isinstance(preprocessor.steps[1], AddBatchDimensionProcessorStep)
+    assert isinstance(preprocessor.steps[2], DeviceProcessorStep)
+    assert isinstance(preprocessor.steps[3], NormalizerProcessorStep)
+
+    # Check steps in postprocessor
+    assert len(postprocessor.steps) == 2
+    assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
+    assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
+
+
+def test_vqbet_processor_with_images():
+    """Test VQBeT processor with image and state observations."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Create test data with images and states
+    observation = {
+        OBS_STATE: torch.randn(8),
+        OBS_IMAGE: torch.randn(3, 224, 224),
+    }
+    action = torch.randn(7)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data is batched
+    assert processed[OBS_STATE].shape == (1, 8)
+    assert processed[OBS_IMAGE].shape == (1, 3, 224, 224)
+    assert processed[TransitionKey.ACTION.value].shape == (1, 7)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_vqbet_processor_cuda():
+    """Test VQBeT processor with CUDA device."""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Create CPU data
+    observation = {
+        OBS_STATE: torch.randn(8),
+        OBS_IMAGE: torch.randn(3, 224, 224),
+    }
+    action = torch.randn(7)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data is on CUDA
+    assert processed[OBS_STATE].device.type == "cuda"
+    assert processed[OBS_IMAGE].device.type == "cuda"
+    assert processed[TransitionKey.ACTION.value].device.type == "cuda"
+
+    # Process through postprocessor
+    postprocessed = postprocessor(processed[TransitionKey.ACTION.value])
+
+    # Check that action is back on CPU
+    assert postprocessed.device.type == "cpu"
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_vqbet_processor_accelerate_scenario():
+    """Test VQBeT processor in simulated Accelerate scenario."""
+    config = create_default_config()
+    config.device = "cuda:0"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Simulate Accelerate: data already on GPU and batched
+    device = torch.device("cuda:0")
+    observation = {
+        OBS_STATE: torch.randn(1, 8).to(device),
+        OBS_IMAGE: torch.randn(1, 3, 224, 224).to(device),
+    }
+    action = torch.randn(1, 7).to(device)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data stays on same GPU
+    assert processed[OBS_STATE].device == device
+    assert processed[OBS_IMAGE].device == device
+    assert processed[TransitionKey.ACTION.value].device == device
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs")
+def test_vqbet_processor_multi_gpu():
+    """Test VQBeT processor with multi-GPU setup."""
+    config = create_default_config()
+    config.device = "cuda:0"
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Simulate data on different GPU
+    device = torch.device("cuda:1")
+    observation = {
+        OBS_STATE: torch.randn(1, 8).to(device),
+        OBS_IMAGE: torch.randn(1, 3, 224, 224).to(device),
+    }
+    action = torch.randn(1, 7).to(device)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data stays on cuda:1
+    assert processed[OBS_STATE].device == device
+    assert processed[OBS_IMAGE].device == device
+    assert processed[TransitionKey.ACTION.value].device == device
+
+
+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)
+
+    # Should still create processors
+    assert preprocessor is not None
+    assert postprocessor is not None
+
+    # Process should still work
+    observation = {
+        OBS_STATE: torch.randn(8),
+        OBS_IMAGE: torch.randn(3, 224, 224),
+    }
+    action = torch.randn(7)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    processed = preprocessor(batch)
+    assert processed is not None
+
+
+def test_vqbet_processor_save_and_load():
+    """Test saving and loading VQBeT processor."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Save preprocessor
+        preprocessor.save_pretrained(tmpdir)
+
+        # Load preprocessor
+        loaded_preprocessor = DataProcessorPipeline.from_pretrained(
+            tmpdir, config_filename="policy_preprocessor.json"
+        )
+
+        # Test that loaded processor works
+        observation = {
+            OBS_STATE: torch.randn(8),
+            OBS_IMAGE: torch.randn(3, 224, 224),
+        }
+        action = torch.randn(7)
+        transition = create_transition(observation, action)
+
+        batch = transition_to_batch(transition)
+        processed = loaded_preprocessor(batch)
+        assert processed[OBS_STATE].shape == (1, 8)
+        assert processed[OBS_IMAGE].shape == (1, 3, 224, 224)
+        assert processed[TransitionKey.ACTION.value].shape == (1, 7)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_vqbet_processor_mixed_precision():
+    """Test VQBeT processor with mixed precision."""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    # Create processor
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Replace DeviceProcessorStep with one that uses float16
+    modified_steps = []
+    for step in preprocessor.steps:
+        if isinstance(step, DeviceProcessorStep):
+            modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16"))
+        elif isinstance(step, NormalizerProcessorStep):
+            # Update normalizer to use the same device as the device processor
+            modified_steps.append(
+                NormalizerProcessorStep(
+                    features=step.features,
+                    norm_map=step.norm_map,
+                    stats=step.stats,
+                    device=config.device,
+                    dtype=torch.float16,  # Match the float16 dtype
+                )
+            )
+        else:
+            modified_steps.append(step)
+    preprocessor.steps = modified_steps
+
+    # Create test data
+    observation = {
+        OBS_STATE: torch.randn(8, dtype=torch.float32),
+        OBS_IMAGE: torch.randn(3, 224, 224, dtype=torch.float32),
+    }
+    action = torch.randn(7, dtype=torch.float32)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that data is converted to float16
+    assert processed[OBS_STATE].dtype == torch.float16
+    assert processed[OBS_IMAGE].dtype == torch.float16
+    assert processed[TransitionKey.ACTION.value].dtype == torch.float16
+
+
+def test_vqbet_processor_large_batch():
+    """Test VQBeT processor with large batch sizes."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Test with large batch
+    batch_size = 128
+    observation = {
+        OBS_STATE: torch.randn(batch_size, 8),
+        OBS_IMAGE: torch.randn(batch_size, 3, 224, 224),
+    }
+    action = torch.randn(batch_size, 7)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through preprocessor
+
+    processed = preprocessor(batch)
+
+    # Check that batch dimension is preserved
+    assert processed[OBS_STATE].shape == (batch_size, 8)
+    assert processed[OBS_IMAGE].shape == (batch_size, 3, 224, 224)
+    assert processed[TransitionKey.ACTION.value].shape == (batch_size, 7)
+
+
+def test_vqbet_processor_sequential_processing():
+    """Test VQBeT processor with sequential data processing."""
+    config = create_default_config()
+    stats = create_default_stats()
+
+    preprocessor, postprocessor = make_vqbet_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Process multiple samples sequentially
+    results = []
+    for _ in range(5):
+        observation = {
+            OBS_STATE: torch.randn(8),
+            OBS_IMAGE: torch.randn(3, 224, 224),
+        }
+        action = torch.randn(7)
+        transition = create_transition(observation, action)
+
+        batch = transition_to_batch(transition)
+
+        processed = preprocessor(batch)
+        results.append(processed)
+
+    # Check that all results are consistent
+    for result in results:
+        assert result[OBS_STATE].shape == (1, 8)
+        assert result[OBS_IMAGE].shape == (1, 3, 224, 224)
+        assert result[TransitionKey.ACTION.value].shape == (1, 7)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_vqbet_processor_bfloat16_device_float32_normalizer():
+    """Test: DeviceProcessor(bfloat16) + NormalizerProcessor(float32) → output bfloat16 via automatic adaptation"""
+    config = create_default_config()
+    config.device = "cuda"
+    stats = create_default_stats()
+
+    preprocessor, _ = make_vqbet_pre_post_processors(
+        config,
+        stats,
+    )
+
+    # Modify the pipeline to use bfloat16 device processor with float32 normalizer
+    modified_steps = []
+    for step in preprocessor.steps:
+        if isinstance(step, DeviceProcessorStep):
+            # Device processor converts to bfloat16
+            modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="bfloat16"))
+        elif isinstance(step, NormalizerProcessorStep):
+            # Normalizer stays configured as float32 (will auto-adapt to bfloat16)
+            modified_steps.append(
+                NormalizerProcessorStep(
+                    features=step.features,
+                    norm_map=step.norm_map,
+                    stats=step.stats,
+                    device=config.device,
+                    dtype=torch.float32,  # Deliberately configured as float32
+                )
+            )
+        else:
+            modified_steps.append(step)
+    preprocessor.steps = modified_steps
+
+    # Verify initial normalizer configuration
+    normalizer_step = preprocessor.steps[3]  # NormalizerProcessorStep
+    assert normalizer_step.dtype == torch.float32
+
+    # Create test data with both state and visual observations
+    observation = {
+        OBS_STATE: torch.randn(8, dtype=torch.float32),
+        OBS_IMAGE: torch.randn(3, 224, 224, dtype=torch.float32),
+    }
+    action = torch.randn(7, dtype=torch.float32)
+    transition = create_transition(observation, action)
+
+    batch = transition_to_batch(transition)
+
+    # Process through full pipeline
+    processed = preprocessor(batch)
+
+    # Verify: DeviceProcessor → bfloat16, NormalizerProcessor adapts → final output is bfloat16
+    assert processed[OBS_STATE].dtype == torch.bfloat16
+    assert processed[OBS_IMAGE].dtype == torch.bfloat16  # IDENTITY normalization still gets dtype conversion
+    assert processed[TransitionKey.ACTION.value].dtype == torch.bfloat16
+
+    # Verify normalizer automatically adapted its internal state
+    assert normalizer_step.dtype == torch.bfloat16
+    # Check state stats (has normalization)
+    for stat_tensor in normalizer_step._tensor_stats[OBS_STATE].values():
+        assert stat_tensor.dtype == torch.bfloat16
+    # OBS_IMAGE uses IDENTITY normalization, so no stats to check