refactor(pipeline): feature contract now categorizes between OBS or Action (#1867)

* refactor(processor): signature of transform_features

* refactor(processor): remove prefixes + processor respect new transform_features signature + update test accordingly

* refactor(processor): rename now is only for visual

* refactor(processor): update normalize processor

* refactor(processor): update vanilla processor features

* refactor(processor): feature contract now uses its own enum

* chore(processor): rename renameprocessor

* chore(processor): minor changes

* refactor(processor): add create & change aggregate

* refactor(processor): update aggregate

* refactor(processor): simplify to functions, fix features contracts and rename function

* test(processor): remove to converter tests as now they are very simple

* chore(docs): recover docs joint observations processor

* fix(processor): update RKP

* fix(tests): recv diff test_pipeline

* chore(tests): add docs to test

* chore(processor): leave obs language constant untouched

* fix(processor): correct new shape of feature in crop image processor

tests/conftest.py

@@ -19,7 +19,7 @@ import traceback
 import pytest
 from serial import SerialException
 
-from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
 from tests.utils import DEVICE
 
 # Import fixture modules as plugins
@@ -83,7 +83,9 @@ def policy_feature_factory():
     return _pf
 
 
-def assert_contract_is_typed(features: dict[str, PolicyFeature]) -> None:
+def assert_contract_is_typed(features: dict[PipelineFeatureType, dict[str, PolicyFeature]]) -> None:
     assert isinstance(features, dict)
-    assert all(isinstance(k, str) for k in features.keys())
-    assert all(isinstance(v, PolicyFeature) for v in features.values())
+    assert all(isinstance(k, PipelineFeatureType) for k in features.keys())
+    assert all(isinstance(v, dict) for v in features.values())
+    assert all(all(isinstance(nk, str) for nk in v.keys()) for v in features.values())
+    assert all(all(isinstance(nv, PolicyFeature) for nv in v.values()) for v in features.values())

tests/processor/test_act_processor.py

@@ -29,7 +29,7 @@ from lerobot.processor import (
     DataProcessorPipeline,
     DeviceProcessorStep,
     NormalizerProcessorStep,
-    RenameProcessorStep,
+    RenameObservationsProcessorStep,
     TransitionKey,
     UnnormalizerProcessorStep,
 )
@@ -86,7 +86,7 @@ def test_make_act_processor_basic():
 
     # Check steps in preprocessor
     assert len(preprocessor.steps) == 4
-    assert isinstance(preprocessor.steps[0], RenameProcessorStep)
+    assert isinstance(preprocessor.steps[0], RenameObservationsProcessorStep)
     assert isinstance(preprocessor.steps[1], AddBatchDimensionProcessorStep)
     assert isinstance(preprocessor.steps[2], DeviceProcessorStep)
     assert isinstance(preprocessor.steps[3], NormalizerProcessorStep)

tests/processor/test_converters.py

@@ -4,111 +4,13 @@ import torch
 
 from lerobot.processor import TransitionKey
 from lerobot.processor.converters import (
-    action_to_transition,
     batch_to_transition,
-    observation_to_transition,
     to_tensor,
     transition_to_batch,
     transition_to_dataset_frame,
-    transition_to_robot_action,
 )
 
 
-def test_to_transition_teleop_action_prefix_and_tensor_conversion():
-    # Scalars, arrays, and uint8 arrays are all converted to tensors
-    img = np.zeros((8, 12, 3), dtype=np.uint8)
-    act = {
-        "ee.x": 0.5,  # scalar to torch tensor
-        "delta": np.array([1.0, 2.0]),  # ndarray to torch tensor
-        "raw_img": img,  # uint8 HWC to torch tensor
-    }
-
-    tr = action_to_transition(act)
-
-    # Should be an EnvTransition-like dict with ACTION populated
-    assert isinstance(tr, dict)
-    assert TransitionKey.ACTION in tr
-    assert "action.ee.x" in tr[TransitionKey.ACTION]
-    assert "action.delta" in tr[TransitionKey.ACTION]
-    assert "action.raw_img" in tr[TransitionKey.ACTION]
-
-    # Types: all values -> torch tensor
-    assert isinstance(tr[TransitionKey.ACTION]["action.ee.x"], torch.Tensor)
-    assert tr[TransitionKey.ACTION]["action.ee.x"].item() == pytest.approx(0.5)
-
-    assert isinstance(tr[TransitionKey.ACTION]["action.delta"], torch.Tensor)
-    assert tr[TransitionKey.ACTION]["action.delta"].shape == (2,)
-    assert torch.allclose(tr[TransitionKey.ACTION]["action.delta"], torch.tensor([1.0, 2.0]))
-
-    assert isinstance(tr[TransitionKey.ACTION]["action.raw_img"], torch.Tensor)
-    assert tr[TransitionKey.ACTION]["action.raw_img"].dtype == torch.float32  # converted from uint8
-    assert tr[TransitionKey.ACTION]["action.raw_img"].shape == (8, 12, 3)
-
-    # Observation is created as empty dict by make_transition
-    assert TransitionKey.OBSERVATION in tr
-    assert isinstance(tr[TransitionKey.OBSERVATION], dict)
-    assert tr[TransitionKey.OBSERVATION] == {}
-
-
-def test_to_transition_robot_observation_state_vs_images_split():
-    # Create an observation with mixed content
-    img = np.full((10, 20, 3), 255, dtype=np.uint8)  # image (uint8 HWC)
-    obs = {
-        "j1.pos": 10.0,  # scalar to state to torch tensor
-        "j2.pos": np.float32(20.0),  # scalar np to state to torch tensor
-        "image_front": img,  # to images passthrough
-        "flag": np.int32(7),  # scalar to state to torch tensor
-        "arr": np.array([1.5, 2.5]),  # vector to state to torch tensor
-    }
-
-    tr = observation_to_transition(obs)
-    assert isinstance(tr, dict)
-    assert TransitionKey.OBSERVATION in tr
-
-    out = tr[TransitionKey.OBSERVATION]
-    # Check state keys are present and converted to tensors
-    for k in ("j1.pos", "j2.pos", "flag", "arr"):
-        key = f"observation.state.{k}"
-        assert key in out
-        v = out[key]
-        if k != "arr":
-            assert isinstance(v, torch.Tensor) and v.ndim == 0
-        else:
-            assert isinstance(v, torch.Tensor) and v.ndim == 1 and v.shape == (2,)
-
-    # Check image present as is
-    assert "observation.images.image_front" in out
-    assert isinstance(out["observation.images.image_front"], np.ndarray)
-    assert out["observation.images.image_front"].dtype == np.uint8
-    assert out["observation.images.image_front"].shape == (10, 20, 3)
-
-    # ACTION should be empty dict by make_transition
-    assert TransitionKey.ACTION in tr
-    assert isinstance(tr[TransitionKey.ACTION], dict)
-    assert tr[TransitionKey.ACTION] == {}
-
-
-def test_to_output_robot_action_strips_prefix_and_filters_pos_keys_only():
-    # Build a transition with mixed action keys
-    tr = {
-        TransitionKey.ACTION: {
-            "action.j1.pos": 11.0,  # keep "j1.pos"
-            "action.gripper.pos": torch.tensor(33.0),  # keep: tensor accepted
-            "action.ee.x": 0.5,  # ignore (doesn't end with .pos)
-            "misc": "ignore_me",  # ignore (no 'action.' prefix)
-        }
-    }
-
-    out = transition_to_robot_action(tr)
-    # Only ".pos" keys with "action." prefix are retained and stripped to base names
-    assert set(out.keys()) == {"j1.pos", "gripper.pos"}
-    # Values converted to float
-    assert isinstance(out["j1.pos"], float)
-    assert isinstance(out["gripper.pos"], float)
-    assert out["j1.pos"] == pytest.approx(11.0)
-    assert out["gripper.pos"] == pytest.approx(33.0)
-
-
 def test_transition_to_dataset_frame_merge_and_pack_vectors_and_metadata():
     # Fabricate dataset features (as stored in dataset.meta["features"])
     features = {

tests/processor/test_device_processor.py

@@ -18,7 +18,7 @@ import tempfile
 import pytest
 import torch
 
-from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
 from lerobot.processor import DataProcessorPipeline, DeviceProcessorStep, TransitionKey
 
 
@@ -292,8 +292,10 @@ def test_features():
     processor = DeviceProcessorStep(device="cpu")
 
     features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(10,)),
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(5,)),
+        PipelineFeatureType.OBSERVATION: {
+            "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(10,))
+        },
+        PipelineFeatureType.ACTION: {"action": PolicyFeature(type=FeatureType.ACTION, shape=(5,))},
     }
 
     result = processor.transform_features(features)

tests/processor/test_diffusion_processor.py

@@ -29,7 +29,7 @@ from lerobot.processor import (
     DataProcessorPipeline,
     DeviceProcessorStep,
     NormalizerProcessorStep,
-    RenameProcessorStep,
+    RenameObservationsProcessorStep,
     TransitionKey,
     UnnormalizerProcessorStep,
 )
@@ -89,7 +89,7 @@ def test_make_diffusion_processor_basic():
 
     # Check steps in preprocessor
     assert len(preprocessor.steps) == 4
-    assert isinstance(preprocessor.steps[0], RenameProcessorStep)
+    assert isinstance(preprocessor.steps[0], RenameObservationsProcessorStep)
     assert isinstance(preprocessor.steps[1], AddBatchDimensionProcessorStep)
     assert isinstance(preprocessor.steps[2], DeviceProcessorStep)
     assert isinstance(preprocessor.steps[3], NormalizerProcessorStep)

tests/processor/test_observation_processor.py

@@ -18,7 +18,7 @@ import numpy as np
 import pytest
 import torch
 
-from lerobot.configs.types import FeatureType
+from lerobot.configs.types import FeatureType, PipelineFeatureType
 from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 from lerobot.processor import TransitionKey, VanillaObservationProcessorStep
 from tests.conftest import assert_contract_is_typed
@@ -412,74 +412,130 @@ def test_equivalent_with_image_dict():
 def test_image_processor_features_pixels_to_image(policy_feature_factory):
     processor = VanillaObservationProcessorStep()
     features = {
-        "pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
-        "keep": policy_feature_factory(FeatureType.ENV, (1,)),
+        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 and out[OBS_IMAGE] == features["pixels"]
-    assert "pixels" not in out
-    assert out["keep"] == features["keep"]
+    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 = {
-        "observation.pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
-        "keep": policy_feature_factory(FeatureType.ENV, (1,)),
+        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 and out[OBS_IMAGE] == features["observation.pixels"]
-    assert "observation.pixels" not in out
-    assert out["keep"] == features["keep"]
+    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 = {
-        "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,)),
+        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 and out[f"{OBS_IMAGES}.front"] == features["pixels.front"]
-    assert f"{OBS_IMAGES}.wrist" in out and out[f"{OBS_IMAGES}.wrist"] == features["pixels.wrist"]
-    assert f"{OBS_IMAGES}.rear" in out and out[f"{OBS_IMAGES}.rear"] == features["observation.pixels.rear"]
-    assert "pixels.front" not in out and "pixels.wrist" not in out and "observation.pixels.rear" not in out
-    assert out["keep"] == features["keep"]
+    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 = {
-        "environment_state": policy_feature_factory(FeatureType.STATE, (3,)),
-        "agent_pos": policy_feature_factory(FeatureType.STATE, (7,)),
-        "keep": policy_feature_factory(FeatureType.ENV, (1,)),
+        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 and out[OBS_ENV_STATE] == features["environment_state"]
-    assert OBS_STATE in out and out[OBS_STATE] == features["agent_pos"]
-    assert "environment_state" not in out and "agent_pos" not in out
-    assert out["keep"] == features["keep"]
+    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 = {
-        "observation.environment_state": policy_feature_factory(FeatureType.STATE, (2,)),
-        "observation.agent_pos": policy_feature_factory(FeatureType.STATE, (4,)),
+        PipelineFeatureType.OBSERVATION: {
+            "observation.environment_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 and out[OBS_ENV_STATE] == features["observation.environment_state"]
-    assert OBS_STATE in out and out[OBS_STATE] == features["observation.agent_pos"]
-    assert "environment_state" not in out and "agent_pos" not in out
+    assert (
+        OBS_ENV_STATE in out[PipelineFeatureType.OBSERVATION]
+        and out[PipelineFeatureType.OBSERVATION][OBS_ENV_STATE]
+        == features[PipelineFeatureType.OBSERVATION]["observation.environment_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_pi0_processor.py

@@ -30,7 +30,7 @@ from lerobot.processor import (
     EnvTransition,
     NormalizerProcessorStep,
     ProcessorStep,
-    RenameProcessorStep,
+    RenameObservationsProcessorStep,
     TransitionKey,
     UnnormalizerProcessorStep,
 )
@@ -115,7 +115,7 @@ def test_make_pi0_processor_basic():
 
     # Check steps in preprocessor
     assert len(preprocessor.steps) == 6
-    assert isinstance(preprocessor.steps[0], RenameProcessorStep)
+    assert isinstance(preprocessor.steps[0], RenameObservationsProcessorStep)
     assert isinstance(preprocessor.steps[1], AddBatchDimensionProcessorStep)
     assert isinstance(preprocessor.steps[2], Pi0NewLineProcessor)
     # Step 3 would be TokenizerProcessorStep but it's mocked

tests/processor/test_pipeline.py

@@ -25,7 +25,7 @@ import pytest
 import torch
 import torch.nn as nn
 
-from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
 from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features
 from lerobot.processor import (
     DataProcessorPipeline,
@@ -96,7 +96,9 @@ class MockStep(ProcessorStep):
     def reset(self) -> None:
         self.counter = 0
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         # We do not test features here
         return features
 
@@ -118,7 +120,9 @@ class MockStepWithoutOptionalMethods(ProcessorStep):
 
         return transition
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         # We do not test features here
         return features
 
@@ -174,7 +178,9 @@ class MockStepWithTensorState(ProcessorStep):
         self.running_mean.zero_()
         self.running_count.zero_()
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         # We do not test features here
         return features
 
@@ -670,7 +676,9 @@ class MockModuleStep(ProcessorStep, nn.Module):
         self.running_mean.zero_()
         self.counter = 0
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         # We do not test features here
         return features
 
@@ -752,7 +760,9 @@ class MockNonModuleStepWithState(ProcessorStep):
         self.step_count.zero_()
         self.history.clear()
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         # We do not test features here
         return features
 
@@ -807,7 +817,9 @@ class MockStepWithNonSerializableParam(ProcessorStep):
     def reset(self) -> None:
         pass
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         # We do not test features here
         return features
 
@@ -846,7 +858,9 @@ class RegisteredMockStep(ProcessorStep):
     def reset(self) -> None:
         pass
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         # We do not test features here
         return features
 
@@ -1406,7 +1420,9 @@ def test_state_file_naming_with_registry():
         def load_state_dict(self, state):
             self.state_tensor = state["state_tensor"]
 
-        def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        def transform_features(
+            self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+        ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
             # We do not test features here
             return features
 
@@ -1463,7 +1479,9 @@ def test_override_with_nested_config():
         def get_config(self):
             return {"name": self.name, "simple_param": self.simple_param, "nested_config": self.nested_config}
 
-        def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        def transform_features(
+            self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+        ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
             # We do not test features here
             return features
 
@@ -1557,7 +1575,9 @@ def test_override_with_callables():
         def get_config(self):
             return {"name": self.name}
 
-        def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        def transform_features(
+            self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+        ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
             # We do not test features here
             return features
 
@@ -1692,7 +1712,9 @@ def test_override_with_device_strings():
         def load_state_dict(self, state):
             self.buffer = state["buffer"]
 
-        def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        def transform_features(
+            self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+        ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
             # We do not test features here
             return features
 
@@ -1805,16 +1827,20 @@ class NonCompliantStep:
         return transition
 
 
-class NonCallableStep:
+class NonCallableStep(ProcessorStep):
     """Intentionally non-compliant: missing __call__."""
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         return features
 
 
-def test_construction_rejects_step_without_processorstep():
+def test_construction_rejects_step_without_call():
     """Test that DataProcessorPipeline rejects steps that don't inherit from ProcessorStep."""
-    with pytest.raises(TypeError, match=r"must inherit from ProcessorStep"):
+    with pytest.raises(
+        TypeError, match=r"Can't instantiate abstract class NonCallableStep with abstract method __call_"
+    ):
         DataProcessorPipeline([NonCallableStep()])
 
     with pytest.raises(TypeError, match=r"must inherit from ProcessorStep"):
@@ -1831,8 +1857,10 @@ class FeatureContractAddStep(ProcessorStep):
     def __call__(self, transition: EnvTransition) -> EnvTransition:
         return transition
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        features[self.key] = self.value
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        features[PipelineFeatureType.OBSERVATION][self.key] = self.value
         return features
 
 
@@ -1846,8 +1874,12 @@ class FeatureContractMutateStep(ProcessorStep):
     def __call__(self, transition: EnvTransition) -> EnvTransition:
         return transition
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        features[self.key] = self.fn(features.get(self.key))
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        features[PipelineFeatureType.OBSERVATION][self.key] = self.fn(
+            features[PipelineFeatureType.OBSERVATION].get(self.key)
+        )
         return features
 
 
@@ -1858,7 +1890,9 @@ class FeatureContractBadReturnStep(ProcessorStep):
     def __call__(self, transition: EnvTransition) -> EnvTransition:
         return transition
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         return ["not-a-dict"]
 
 
@@ -1871,8 +1905,10 @@ class FeatureContractRemoveStep(ProcessorStep):
     def __call__(self, transition: EnvTransition) -> EnvTransition:
         return transition
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        features.pop(self.key, None)
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        features[PipelineFeatureType.OBSERVATION].pop(self.key, None)
         return features
 
 
@@ -1884,17 +1920,22 @@ def test_features_orders_and_merges(policy_feature_factory):
             FeatureContractAddStep("b", policy_feature_factory(FeatureType.ENV, (2,))),
         ]
     )
-    out = p.transform_features({})
-
-    assert out["a"].type == FeatureType.STATE and out["a"].shape == (3,)
-    assert out["b"].type == FeatureType.ENV and out["b"].shape == (2,)
+    out = p.transform_features({PipelineFeatureType.OBSERVATION: {}})
+    assert out[PipelineFeatureType.OBSERVATION]["a"].type == FeatureType.STATE and out[
+        PipelineFeatureType.OBSERVATION
+    ]["a"].shape == (3,)
+    assert out[PipelineFeatureType.OBSERVATION]["b"].type == FeatureType.ENV and out[
+        PipelineFeatureType.OBSERVATION
+    ]["b"].shape == (2,)
     assert_contract_is_typed(out)
 
 
 def test_features_respects_initial_without_mutation(policy_feature_factory):
     initial = {
-        "seed": policy_feature_factory(FeatureType.STATE, (7,)),
-        "nested": policy_feature_factory(FeatureType.ENV, (0,)),
+        PipelineFeatureType.OBSERVATION: {
+            "seed": policy_feature_factory(FeatureType.STATE, (7,)),
+            "nested": policy_feature_factory(FeatureType.ENV, (0,)),
+        }
     }
     p = DataProcessorPipeline(
         [
@@ -1906,11 +1947,11 @@ def test_features_respects_initial_without_mutation(policy_feature_factory):
     )
     out = p.transform_features(initial_features=initial)
 
-    assert out["seed"].shape == (8,)
-    assert out["nested"].shape == (5,)
+    assert out[PipelineFeatureType.OBSERVATION]["seed"].shape == (8,)
+    assert out[PipelineFeatureType.OBSERVATION]["nested"].shape == (5,)
     # Initial dict must be preserved
-    assert initial["seed"].shape == (7,)
-    assert initial["nested"].shape == (0,)
+    assert initial[PipelineFeatureType.OBSERVATION]["seed"].shape == (7,)
+    assert initial[PipelineFeatureType.OBSERVATION]["nested"].shape == (0,)
 
     assert_contract_is_typed(out)
 
@@ -1923,14 +1964,22 @@ def test_features_execution_order_tracking():
         def __call__(self, transition: EnvTransition) -> EnvTransition:
             return transition
 
-        def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        def transform_features(
+            self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+        ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
             code = {"A": 1, "B": 2, "C": 3}[self.label]
-            pf = features.get("order", PolicyFeature(type=FeatureType.ENV, shape=()))
-            features["order"] = PolicyFeature(type=pf.type, shape=pf.shape + (code,))
+            pf = features[PipelineFeatureType.OBSERVATION].get(
+                "order", PolicyFeature(type=FeatureType.ENV, shape=())
+            )
+            features[PipelineFeatureType.OBSERVATION]["order"] = PolicyFeature(
+                type=pf.type, shape=pf.shape + (code,)
+            )
             return features
 
-    out = DataProcessorPipeline([Track("A"), Track("B"), Track("C")]).transform_features({})
-    assert out["order"].shape == (1, 2, 3)
+    out = DataProcessorPipeline([Track("A"), Track("B"), Track("C")]).transform_features(
+        initial_features={PipelineFeatureType.OBSERVATION: {}}
+    )
+    assert out[PipelineFeatureType.OBSERVATION]["order"].shape == (1, 2, 3)
 
 
 def test_features_remove_key(policy_feature_factory):
@@ -1940,18 +1989,23 @@ def test_features_remove_key(policy_feature_factory):
             FeatureContractRemoveStep("a"),
         ]
     )
-    out = p.transform_features({})
-    assert "a" not in out
+    out = p.transform_features({PipelineFeatureType.OBSERVATION: {}})
+    assert "a" not in out[PipelineFeatureType.OBSERVATION]
 
 
 def test_features_remove_from_initial(policy_feature_factory):
     initial = {
-        "keep": policy_feature_factory(FeatureType.STATE, (1,)),
-        "drop": policy_feature_factory(FeatureType.STATE, (1,)),
+        PipelineFeatureType.OBSERVATION: {
+            "keep": policy_feature_factory(FeatureType.STATE, (1,)),
+            "drop": policy_feature_factory(FeatureType.STATE, (1,)),
+        },
     }
     p = DataProcessorPipeline([FeatureContractRemoveStep("drop")])
     out = p.transform_features(initial_features=initial)
-    assert "drop" not in out and out["keep"] == initial["keep"]
+    assert (
+        "drop" not in out[PipelineFeatureType.OBSERVATION]
+        and out[PipelineFeatureType.OBSERVATION]["keep"] == initial[PipelineFeatureType.OBSERVATION]["keep"]
+    )
 
 
 @dataclass
@@ -1961,13 +2015,15 @@ class AddActionEEAndJointFeatures(ProcessorStep):
     def __call__(self, tr):
         return tr
 
-    def transform_features(self, features: dict) -> dict:
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         # EE features
-        features["action.ee.x"] = float
-        features["action.ee.y"] = float
+        features[PipelineFeatureType.ACTION]["action.ee.x"] = float
+        features[PipelineFeatureType.ACTION]["action.ee.y"] = float
         # JOINT features
-        features["action.j1.pos"] = float
-        features["action.j2.pos"] = float
+        features[PipelineFeatureType.ACTION]["action.j1.pos"] = float
+        features[PipelineFeatureType.ACTION]["action.j2.pos"] = float
         return features
 
 
@@ -1981,18 +2037,20 @@ class AddObservationStateFeatures(ProcessorStep):
     def __call__(self, tr):
         return tr
 
-    def transform_features(self, features: dict) -> dict:
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         # State features (mix EE and a joint state)
-        features["observation.state.ee.x"] = float
-        features["observation.state.j1.pos"] = float
+        features[PipelineFeatureType.OBSERVATION]["observation.state.ee.x"] = float
+        features[PipelineFeatureType.OBSERVATION]["observation.state.j1.pos"] = float
         if self.add_front_image:
-            features["observation.images.front"] = self.front_image_shape
+            features[PipelineFeatureType.OBSERVATION]["observation.images.front"] = self.front_image_shape
         return features
 
 
 def test_aggregate_joint_action_only():
     rp = DataProcessorPipeline([AddActionEEAndJointFeatures()])
-    initial = {"front": (480, 640, 3)}
+    initial = {PipelineFeatureType.OBSERVATION: {"front": (480, 640, 3)}, PipelineFeatureType.ACTION: {}}
 
     out = aggregate_pipeline_dataset_features(
         pipeline=rp,
@@ -2014,7 +2072,7 @@ def test_aggregate_ee_action_and_observation_with_videos():
 
     out = aggregate_pipeline_dataset_features(
         pipeline=rp,
-        initial_features=initial,
+        initial_features={PipelineFeatureType.OBSERVATION: initial, PipelineFeatureType.ACTION: {}},
         use_videos=True,
         patterns=["action.ee", "observation.state"],
     )
@@ -2042,7 +2100,7 @@ def test_aggregate_both_action_types():
     rp = DataProcessorPipeline([AddActionEEAndJointFeatures()])
     out = aggregate_pipeline_dataset_features(
         pipeline=rp,
-        initial_features={},
+        initial_features={PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: {}},
         use_videos=True,
         patterns=["action.ee", "action.j1", "action.j2.pos"],
     )
@@ -2059,7 +2117,7 @@ def test_aggregate_images_when_use_videos_false():
 
     out = aggregate_pipeline_dataset_features(
         pipeline=rp,
-        initial_features=initial,
+        initial_features={PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: initial},
         use_videos=False,  # expect "image" dtype
         patterns=None,
     )
@@ -2076,7 +2134,7 @@ def test_aggregate_images_when_use_videos_true():
 
     out = aggregate_pipeline_dataset_features(
         pipeline=rp,
-        initial_features=initial,
+        initial_features={PipelineFeatureType.OBSERVATION: initial, PipelineFeatureType.ACTION: {}},
         use_videos=True,
         patterns=None,
     )
@@ -2100,7 +2158,7 @@ def test_initial_camera_not_overridden_by_step_image():
 
     out = aggregate_pipeline_dataset_features(
         pipeline=rp,
-        initial_features=initial,
+        initial_features={PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: initial},
         use_videos=True,
         patterns=["observation.images.front"],
     )

tests/processor/test_rename_processor.py

@@ -19,11 +19,11 @@ from pathlib import Path
 import numpy as np
 import torch
 
-from lerobot.configs.types import FeatureType
+from lerobot.configs.types import FeatureType, PipelineFeatureType
 from lerobot.processor import (
     DataProcessorPipeline,
     ProcessorStepRegistry,
-    RenameProcessorStep,
+    RenameObservationsProcessorStep,
     TransitionKey,
 )
 from lerobot.processor.rename_processor import rename_stats
@@ -51,7 +51,7 @@ def test_basic_renaming():
         "old_key1": "new_key1",
         "old_key2": "new_key2",
     }
-    processor = RenameProcessorStep(rename_map=rename_map)
+    processor = RenameObservationsProcessorStep(rename_map=rename_map)
 
     observation = {
         "old_key1": torch.tensor([1.0, 2.0]),
@@ -79,7 +79,7 @@ def test_basic_renaming():
 
 def test_empty_rename_map():
     """Test processor with empty rename map (should pass through unchanged)."""
-    processor = RenameProcessorStep(rename_map={})
+    processor = RenameObservationsProcessorStep(rename_map={})
 
     observation = {
         "key1": torch.tensor([1.0]),
@@ -98,7 +98,7 @@ def test_empty_rename_map():
 
 def test_none_observation():
     """Test processor with None observation."""
-    processor = RenameProcessorStep(rename_map={"old": "new"})
+    processor = RenameObservationsProcessorStep(rename_map={"old": "new"})
 
     transition = create_transition()
     result = processor(transition)
@@ -113,7 +113,7 @@ def test_overlapping_rename():
         "a": "b",
         "b": "c",  # This creates a potential conflict
     }
-    processor = RenameProcessorStep(rename_map=rename_map)
+    processor = RenameObservationsProcessorStep(rename_map=rename_map)
 
     observation = {
         "a": 1,
@@ -138,7 +138,7 @@ def test_partial_rename():
         "observation.state": "observation.proprio_state",
         "pixels": "observation.image",
     }
-    processor = RenameProcessorStep(rename_map=rename_map)
+    processor = RenameObservationsProcessorStep(rename_map=rename_map)
 
     observation = {
         "observation.state": torch.randn(10),
@@ -168,7 +168,7 @@ def test_get_config():
         "old1": "new1",
         "old2": "new2",
     }
-    processor = RenameProcessorStep(rename_map=rename_map)
+    processor = RenameObservationsProcessorStep(rename_map=rename_map)
 
     config = processor.get_config()
     assert config == {"rename_map": rename_map}
@@ -176,7 +176,7 @@ def test_get_config():
 
 def test_state_dict():
     """Test state dict (should be empty for RenameProcessorStep)."""
-    processor = RenameProcessorStep(rename_map={"old": "new"})
+    processor = RenameObservationsProcessorStep(rename_map={"old": "new"})
 
     state = processor.state_dict()
     assert state == {}
@@ -191,7 +191,7 @@ def test_integration_with_robot_processor():
         "agent_pos": "observation.state",
         "pixels": "observation.image",
     }
-    rename_processor = RenameProcessorStep(rename_map=rename_map)
+    rename_processor = RenameObservationsProcessorStep(rename_map=rename_map)
 
     pipeline = DataProcessorPipeline([rename_processor], to_transition=lambda x: x, to_output=lambda x: x)
 
@@ -225,7 +225,7 @@ def test_save_and_load_pretrained():
         "old_state": "observation.state",
         "old_image": "observation.image",
     }
-    processor = RenameProcessorStep(rename_map=rename_map)
+    processor = RenameObservationsProcessorStep(rename_map=rename_map)
     pipeline = DataProcessorPipeline([processor], name="TestRenameProcessorStep")
 
     with tempfile.TemporaryDirectory() as tmp_dir:
@@ -252,7 +252,7 @@ def test_save_and_load_pretrained():
 
         # Check that loaded processor works correctly
         loaded_processor = loaded_pipeline.steps[0]
-        assert isinstance(loaded_processor, RenameProcessorStep)
+        assert isinstance(loaded_processor, RenameObservationsProcessorStep)
         assert loaded_processor.rename_map == rename_map
 
         # Test functionality after loading
@@ -271,21 +271,21 @@ def test_save_and_load_pretrained():
 def test_registry_functionality():
     """Test that RenameProcessorStep is properly registered."""
     # Check that it's registered
-    assert "rename_processor" in ProcessorStepRegistry.list()
+    assert "rename_observations_processor" in ProcessorStepRegistry.list()
 
     # Get from registry
-    retrieved_class = ProcessorStepRegistry.get("rename_processor")
-    assert retrieved_class is RenameProcessorStep
+    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, RenameProcessorStep)
+    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 = RenameProcessorStep(rename_map={"key1": "renamed_key1"})
+    processor = RenameObservationsProcessorStep(rename_map={"key1": "renamed_key1"})
     pipeline = DataProcessorPipeline([processor], to_transition=lambda x: x, to_output=lambda x: x)
 
     with tempfile.TemporaryDirectory() as tmp_dir:
@@ -299,20 +299,20 @@ def test_registry_based_save_load():
             config = json.load(f)
 
         assert "registry_name" in config["steps"][0]
-        assert config["steps"][0]["registry_name"] == "rename_processor"
+        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)
         loaded_processor = loaded_pipeline.steps[0]
-        assert isinstance(loaded_processor, RenameProcessorStep)
+        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 = RenameProcessorStep(
+    processor1 = RenameObservationsProcessorStep(
         rename_map={
             "pos": "agent_position",
             "img": "camera_image",
@@ -320,7 +320,7 @@ def test_chained_rename_processors():
     )
 
     # Second processor: rename to final format
-    processor2 = RenameProcessorStep(
+    processor2 = RenameObservationsProcessorStep(
         rename_map={
             "agent_position": "observation.state",
             "camera_image": "observation.image",
@@ -365,7 +365,7 @@ def test_nested_observation_rename():
         "observation.images.right": "observation.camera.right_view",
         "observation.proprio": "observation.proprioception",
     }
-    processor = RenameProcessorStep(rename_map=rename_map)
+    processor = RenameObservationsProcessorStep(rename_map=rename_map)
 
     observation = {
         "observation.images.left": torch.randn(3, 64, 64),
@@ -395,7 +395,7 @@ def test_nested_observation_rename():
 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 = RenameProcessorStep(rename_map=rename_map)
+    processor = RenameObservationsProcessorStep(rename_map=rename_map)
 
     tensor_value = torch.randn(3, 3)
     array_value = np.random.rand(2, 2)
@@ -423,59 +423,75 @@ def test_value_types_preserved():
 
 
 def test_features_basic_renaming(policy_feature_factory):
-    processor = RenameProcessorStep(rename_map={"a": "x", "b": "y"})
+    processor = RenameObservationsProcessorStep(rename_map={"a": "x", "b": "y"})
     features = {
-        "a": policy_feature_factory(FeatureType.STATE, (2,)),
-        "b": policy_feature_factory(FeatureType.ACTION, (3,)),
-        "c": policy_feature_factory(FeatureType.ENV, (1,)),
+        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["x"] == features["a"]
-    assert out["y"] == features["b"]
-    assert out["c"] == features["c"]
+    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) == {"a", "b", "c"}
+    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 = RenameProcessorStep(rename_map={"a": "b", "b": "c"})
+    processor = RenameObservationsProcessorStep(rename_map={"a": "b", "b": "c"})
     features = {
-        "a": policy_feature_factory(FeatureType.STATE, (1,)),
-        "b": policy_feature_factory(FeatureType.STATE, (2,)),
+        PipelineFeatureType.OBSERVATION: {
+            "a": policy_feature_factory(FeatureType.VISUAL, (1,)),
+            "b": policy_feature_factory(FeatureType.VISUAL, (2,)),
+        },
     }
     out = processor.transform_features(features)
 
-    assert set(out) == {"b", "c"}
-    assert out["b"] == features["a"]  # 'a' renamed to'b'
-    assert out["c"] == features["b"]  # 'b' renamed to 'c'
+    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 = RenameProcessorStep(rename_map={"pos": "agent_position", "img": "camera_image"})
-    processor2 = RenameProcessorStep(
+    processor1 = RenameObservationsProcessorStep(rename_map={"pos": "agent_position", "img": "camera_image"})
+    processor2 = RenameObservationsProcessorStep(
         rename_map={"agent_position": "observation.state", "camera_image": "observation.image"}
     )
     pipeline = DataProcessorPipeline([processor1, processor2])
 
     spec = {
-        "pos": policy_feature_factory(FeatureType.STATE, (7,)),
-        "img": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
-        "extra": policy_feature_factory(FeatureType.ENV, (1,)),
+        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) == {"observation.state", "observation.image", "extra"}
-    assert out["observation.state"] == spec["pos"]
-    assert out["observation.image"] == spec["img"]
-    assert out["extra"] == spec["extra"]
+    assert set(out[PipelineFeatureType.OBSERVATION]) == {"observation.state", "observation.image", "extra"}
+    assert (
+        out[PipelineFeatureType.OBSERVATION]["observation.state"]
+        == spec[PipelineFeatureType.OBSERVATION]["pos"]
+    )
+    assert (
+        out[PipelineFeatureType.OBSERVATION]["observation.image"]
+        == spec[PipelineFeatureType.OBSERVATION]["img"]
+    )
+    assert out[PipelineFeatureType.OBSERVATION]["extra"] == spec[PipelineFeatureType.OBSERVATION]["extra"]
     assert_contract_is_typed(out)
 
 

tests/processor/test_sac_processor.py

@@ -29,7 +29,7 @@ from lerobot.processor import (
     DataProcessorPipeline,
     DeviceProcessorStep,
     NormalizerProcessorStep,
-    RenameProcessorStep,
+    RenameObservationsProcessorStep,
     TransitionKey,
     UnnormalizerProcessorStep,
 )
@@ -91,7 +91,7 @@ def test_make_sac_processor_basic():
 
     # Check steps in preprocessor
     assert len(preprocessor.steps) == 4
-    assert isinstance(preprocessor.steps[0], RenameProcessorStep)
+    assert isinstance(preprocessor.steps[0], RenameObservationsProcessorStep)
     assert isinstance(preprocessor.steps[1], AddBatchDimensionProcessorStep)
     assert isinstance(preprocessor.steps[2], DeviceProcessorStep)
     assert isinstance(preprocessor.steps[3], NormalizerProcessorStep)

tests/processor/test_smolvla_processor.py

@@ -20,7 +20,7 @@ from unittest.mock import patch
 import pytest
 import torch
 
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.configs.types import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature
 from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.smolvla.processor_smolvla import (
@@ -33,7 +33,7 @@ from lerobot.processor import (
     EnvTransition,
     NormalizerProcessorStep,
     ProcessorStep,
-    RenameProcessorStep,
+    RenameObservationsProcessorStep,
     TransitionKey,
     UnnormalizerProcessorStep,
 )
@@ -122,7 +122,7 @@ def test_make_smolvla_processor_basic():
 
     # Check steps in preprocessor
     assert len(preprocessor.steps) == 6
-    assert isinstance(preprocessor.steps[0], RenameProcessorStep)
+    assert isinstance(preprocessor.steps[0], RenameObservationsProcessorStep)
     assert isinstance(preprocessor.steps[1], AddBatchDimensionProcessorStep)
     assert isinstance(preprocessor.steps[2], SmolVLANewLineProcessor)
     # Step 3 would be TokenizerProcessorStep but it's mocked
@@ -400,7 +400,7 @@ def test_smolvla_newline_processor_transform_features():
 
     # Test transform_features
     features = {
-        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(10,)),
+        PipelineFeatureType.OBSERVATION: {OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(10,))},
     }
     result = processor.transform_features(features)
     assert result == features  # Should return unchanged

tests/processor/test_tdmpc_processor.py

@@ -29,7 +29,7 @@ from lerobot.processor import (
     DataProcessorPipeline,
     DeviceProcessorStep,
     NormalizerProcessorStep,
-    RenameProcessorStep,
+    RenameObservationsProcessorStep,
     TransitionKey,
     UnnormalizerProcessorStep,
 )
@@ -94,7 +94,7 @@ def test_make_tdmpc_processor_basic():
 
     # Check steps in preprocessor
     assert len(preprocessor.steps) == 4
-    assert isinstance(preprocessor.steps[0], RenameProcessorStep)
+    assert isinstance(preprocessor.steps[0], RenameObservationsProcessorStep)
     assert isinstance(preprocessor.steps[1], AddBatchDimensionProcessorStep)
     assert isinstance(preprocessor.steps[2], DeviceProcessorStep)
     assert isinstance(preprocessor.steps[3], NormalizerProcessorStep)

tests/processor/test_tokenizer_processor.py

@@ -8,7 +8,7 @@ from unittest.mock import patch
 import pytest
 import torch
 
-from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
 from lerobot.constants import OBS_LANGUAGE
 from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep, TransitionKey
 from tests.utils import require_package
@@ -512,23 +512,27 @@ def test_features_basic():
     processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=128)
 
     input_features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(10,)),
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(5,)),
+        PipelineFeatureType.OBSERVATION: {
+            "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(10,))
+        },
+        PipelineFeatureType.ACTION: {"action": PolicyFeature(type=FeatureType.ACTION, shape=(5,))},
     }
 
     output_features = processor.transform_features(input_features)
 
     # Check that original features are preserved
-    assert "observation.state" in output_features
-    assert "action" in output_features
+    assert "observation.state" in output_features[PipelineFeatureType.OBSERVATION]
+    assert "action" in output_features[PipelineFeatureType.ACTION]
 
     # Check that tokenized features are added
-    assert f"{OBS_LANGUAGE}.tokens" in output_features
-    assert f"{OBS_LANGUAGE}.attention_mask" in output_features
+    assert f"{OBS_LANGUAGE}.tokens" in output_features[PipelineFeatureType.OBSERVATION]
+    assert f"{OBS_LANGUAGE}.attention_mask" in output_features[PipelineFeatureType.OBSERVATION]
 
     # Check feature properties
-    tokens_feature = output_features[f"{OBS_LANGUAGE}.tokens"]
-    attention_mask_feature = output_features[f"{OBS_LANGUAGE}.attention_mask"]
+    tokens_feature = output_features[PipelineFeatureType.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
+    attention_mask_feature = output_features[PipelineFeatureType.OBSERVATION][
+        f"{OBS_LANGUAGE}.attention_mask"
+    ]
 
     assert tokens_feature.type == FeatureType.LANGUAGE
     assert tokens_feature.shape == (128,)
@@ -542,15 +546,17 @@ def test_features_with_custom_max_length():
     mock_tokenizer = MockTokenizer(vocab_size=100)
     processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=64)
 
-    input_features = {}
+    input_features = {PipelineFeatureType.OBSERVATION: {}}
     output_features = processor.transform_features(input_features)
 
     # Check that features use correct max_length
-    assert f"{OBS_LANGUAGE}.tokens" in output_features
-    assert f"{OBS_LANGUAGE}.attention_mask" in output_features
+    assert f"{OBS_LANGUAGE}.tokens" in output_features[PipelineFeatureType.OBSERVATION]
+    assert f"{OBS_LANGUAGE}.attention_mask" in output_features[PipelineFeatureType.OBSERVATION]
 
-    tokens_feature = output_features[f"{OBS_LANGUAGE}.tokens"]
-    attention_mask_feature = output_features[f"{OBS_LANGUAGE}.attention_mask"]
+    tokens_feature = output_features[PipelineFeatureType.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
+    attention_mask_feature = output_features[PipelineFeatureType.OBSERVATION][
+        f"{OBS_LANGUAGE}.attention_mask"
+    ]
 
     assert tokens_feature.shape == (64,)
     assert attention_mask_feature.shape == (64,)
@@ -563,15 +569,19 @@ def test_features_existing_features():
     processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=256)
 
     input_features = {
-        f"{OBS_LANGUAGE}.tokens": PolicyFeature(type=FeatureType.LANGUAGE, shape=(100,)),
-        f"{OBS_LANGUAGE}.attention_mask": PolicyFeature(type=FeatureType.LANGUAGE, shape=(100,)),
+        PipelineFeatureType.OBSERVATION: {
+            f"{OBS_LANGUAGE}.tokens": PolicyFeature(type=FeatureType.LANGUAGE, shape=(100,)),
+            f"{OBS_LANGUAGE}.attention_mask": PolicyFeature(type=FeatureType.LANGUAGE, shape=(100,)),
+        }
     }
 
     output_features = processor.transform_features(input_features)
 
     # Should not overwrite existing features
-    assert output_features[f"{OBS_LANGUAGE}.tokens"].shape == (100,)  # Original shape preserved
-    assert output_features[f"{OBS_LANGUAGE}.attention_mask"].shape == (100,)
+    assert output_features[PipelineFeatureType.OBSERVATION][f"{OBS_LANGUAGE}.tokens"].shape == (
+        100,
+    )  # Original shape preserved
+    assert output_features[PipelineFeatureType.OBSERVATION][f"{OBS_LANGUAGE}.attention_mask"].shape == (100,)
 
 
 @require_package("transformers")

tests/processor/test_vqbet_processor.py

@@ -29,7 +29,7 @@ from lerobot.processor import (
     DataProcessorPipeline,
     DeviceProcessorStep,
     NormalizerProcessorStep,
-    RenameProcessorStep,
+    RenameObservationsProcessorStep,
     TransitionKey,
     UnnormalizerProcessorStep,
 )
@@ -94,7 +94,7 @@ def test_make_vqbet_processor_basic():
 
     # Check steps in preprocessor
     assert len(preprocessor.steps) == 4
-    assert isinstance(preprocessor.steps[0], RenameProcessorStep)
+    assert isinstance(preprocessor.steps[0], RenameObservationsProcessorStep)
     assert isinstance(preprocessor.steps[1], AddBatchDimensionProcessorStep)
     assert isinstance(preprocessor.steps[2], DeviceProcessorStep)
     assert isinstance(preprocessor.steps[3], NormalizerProcessorStep)