mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-15 08:39:49 +00:00
Adil Zouitine
453e0a995f
- Added `RenameProcessor` to facilitate key renaming in observations, improving data handling flexibility. - Updated `__init__.py` to include `RenameProcessor` in module exports. - Refactored `NormalizationProcessor` and `ObservationNormalizer` to use `rsplit` for better key handling. - Introduced comprehensive tests for `NormalizationProcessor` and `RenameProcessor` to ensure functionality and robustness.
394 lines
13 KiB
Python
394 lines
13 KiB
Python
#!/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.processor.pipeline import ProcessorStepRegistry, RobotProcessor, TransitionIndex
|
|
from lerobot.processor.rename_processor import RenameProcessor
|
|
|
|
|
|
def test_basic_renaming():
|
|
"""Test basic key renaming functionality."""
|
|
rename_map = {
|
|
"old_key1": "new_key1",
|
|
"old_key2": "new_key2",
|
|
}
|
|
processor = RenameProcessor(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 = (observation, None, None, None, None, None, None)
|
|
|
|
result = processor(transition)
|
|
processed_obs = result[TransitionIndex.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 = RenameProcessor(rename_map={})
|
|
|
|
observation = {
|
|
"key1": torch.tensor([1.0]),
|
|
"key2": "value2",
|
|
}
|
|
transition = (observation, None, None, None, None, None, None)
|
|
|
|
result = processor(transition)
|
|
processed_obs = result[TransitionIndex.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 = RenameProcessor(rename_map={"old": "new"})
|
|
|
|
transition = (None, None, None, None, None, None, None)
|
|
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 = RenameProcessor(rename_map=rename_map)
|
|
|
|
observation = {
|
|
"a": 1,
|
|
"b": 2,
|
|
"x": 3,
|
|
}
|
|
transition = (observation, None, None, None, None, None, None)
|
|
|
|
result = processor(transition)
|
|
processed_obs = result[TransitionIndex.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 = {
|
|
"observation.state": "observation.proprio_state",
|
|
"pixels": "observation.image",
|
|
}
|
|
processor = RenameProcessor(rename_map=rename_map)
|
|
|
|
observation = {
|
|
"observation.state": torch.randn(10),
|
|
"pixels": np.random.randint(0, 256, (64, 64, 3), dtype=np.uint8),
|
|
"reward": 1.0,
|
|
"info": {"episode": 1},
|
|
}
|
|
transition = (observation, None, None, None, None, None, None)
|
|
|
|
result = processor(transition)
|
|
processed_obs = result[TransitionIndex.OBSERVATION]
|
|
|
|
# Check renamed keys
|
|
assert "observation.proprio_state" in processed_obs
|
|
assert "observation.image" in processed_obs
|
|
assert "observation.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 = RenameProcessor(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 RenameProcessor)."""
|
|
processor = RenameProcessor(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": "observation.state",
|
|
"pixels": "observation.image",
|
|
}
|
|
rename_processor = RenameProcessor(rename_map=rename_map)
|
|
|
|
pipeline = RobotProcessor([rename_processor])
|
|
|
|
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 = (observation, None, 0.5, False, False, {}, {})
|
|
|
|
result = pipeline(transition)
|
|
processed_obs = result[TransitionIndex.OBSERVATION]
|
|
|
|
# Check renaming worked through pipeline
|
|
assert "observation.state" in processed_obs
|
|
assert "observation.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[TransitionIndex.REWARD] == 0.5
|
|
assert result[TransitionIndex.DONE] is False
|
|
|
|
|
|
def test_save_and_load_pretrained():
|
|
"""Test saving and loading processor with RobotProcessor."""
|
|
rename_map = {
|
|
"old_state": "observation.state",
|
|
"old_image": "observation.image",
|
|
}
|
|
processor = RenameProcessor(rename_map=rename_map)
|
|
pipeline = RobotProcessor([processor], name="TestRenameProcessor")
|
|
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
# Save pipeline
|
|
pipeline.save_pretrained(tmp_dir)
|
|
|
|
# Check files were created
|
|
config_path = Path(tmp_dir) / "processor.json"
|
|
assert config_path.exists()
|
|
|
|
# No state files should be created for RenameProcessor
|
|
state_files = list(Path(tmp_dir).glob("*.safetensors"))
|
|
assert len(state_files) == 0
|
|
|
|
# Load pipeline
|
|
loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir)
|
|
|
|
assert loaded_pipeline.name == "TestRenameProcessor"
|
|
assert len(loaded_pipeline) == 1
|
|
|
|
# Check that loaded processor works correctly
|
|
loaded_processor = loaded_pipeline.steps[0]
|
|
assert isinstance(loaded_processor, RenameProcessor)
|
|
assert loaded_processor.rename_map == rename_map
|
|
|
|
# Test functionality after loading
|
|
observation = {"old_state": [1, 2, 3], "old_image": "image_data"}
|
|
transition = (observation, None, None, None, None, None, None)
|
|
|
|
result = loaded_pipeline(transition)
|
|
processed_obs = result[TransitionIndex.OBSERVATION]
|
|
|
|
assert "observation.state" in processed_obs
|
|
assert "observation.image" in processed_obs
|
|
assert processed_obs["observation.state"] == [1, 2, 3]
|
|
assert processed_obs["observation.image"] == "image_data"
|
|
|
|
|
|
def test_registry_functionality():
|
|
"""Test that RenameProcessor is properly registered."""
|
|
# Check that it's registered
|
|
assert "rename_processor" in ProcessorStepRegistry.list()
|
|
|
|
# Get from registry
|
|
retrieved_class = ProcessorStepRegistry.get("rename_processor")
|
|
assert retrieved_class is RenameProcessor
|
|
|
|
# Create instance from registry
|
|
instance = retrieved_class(rename_map={"old": "new"})
|
|
assert isinstance(instance, RenameProcessor)
|
|
assert instance.rename_map == {"old": "new"}
|
|
|
|
|
|
def test_registry_based_save_load():
|
|
"""Test save/load using registry name instead of module path."""
|
|
processor = RenameProcessor(rename_map={"key1": "renamed_key1"})
|
|
pipeline = RobotProcessor([processor])
|
|
|
|
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) / "processor.json") as f:
|
|
config = json.load(f)
|
|
|
|
assert "registry_name" in config["steps"][0]
|
|
assert config["steps"][0]["registry_name"] == "rename_processor"
|
|
assert "class" not in config["steps"][0] # Should use registry, not module path
|
|
|
|
# Load should work
|
|
loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir)
|
|
loaded_processor = loaded_pipeline.steps[0]
|
|
assert isinstance(loaded_processor, RenameProcessor)
|
|
assert loaded_processor.rename_map == {"key1": "renamed_key1"}
|
|
|
|
|
|
def test_chained_rename_processors():
|
|
"""Test multiple RenameProcessors in a pipeline."""
|
|
# First processor: rename raw keys to intermediate format
|
|
processor1 = RenameProcessor(
|
|
rename_map={
|
|
"pos": "agent_position",
|
|
"img": "camera_image",
|
|
}
|
|
)
|
|
|
|
# Second processor: rename to final format
|
|
processor2 = RenameProcessor(
|
|
rename_map={
|
|
"agent_position": "observation.state",
|
|
"camera_image": "observation.image",
|
|
}
|
|
)
|
|
|
|
pipeline = RobotProcessor([processor1, processor2])
|
|
|
|
observation = {
|
|
"pos": np.array([1.0, 2.0]),
|
|
"img": "image_data",
|
|
"extra": "keep_me",
|
|
}
|
|
transition = (observation, None, None, None, None, None, None)
|
|
|
|
# Step through to see intermediate results
|
|
results = list(pipeline.step_through(transition))
|
|
|
|
# After first processor
|
|
assert "agent_position" in results[1][TransitionIndex.OBSERVATION]
|
|
assert "camera_image" in results[1][TransitionIndex.OBSERVATION]
|
|
|
|
# After second processor
|
|
final_obs = results[2][TransitionIndex.OBSERVATION]
|
|
assert "observation.state" in final_obs
|
|
assert "observation.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 = {
|
|
"observation.images.left": "observation.camera.left_view",
|
|
"observation.images.right": "observation.camera.right_view",
|
|
"observation.proprio": "observation.proprioception",
|
|
}
|
|
processor = RenameProcessor(rename_map=rename_map)
|
|
|
|
observation = {
|
|
"observation.images.left": torch.randn(3, 64, 64),
|
|
"observation.images.right": torch.randn(3, 64, 64),
|
|
"observation.proprio": torch.randn(7),
|
|
"observation.gripper": torch.tensor([0.0]), # Not renamed
|
|
}
|
|
transition = (observation, None, None, None, None, None, None)
|
|
|
|
result = processor(transition)
|
|
processed_obs = result[TransitionIndex.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 "observation.images.left" not in processed_obs
|
|
assert "observation.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 = RenameProcessor(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 = (observation, None, None, None, None, None, None)
|
|
|
|
result = processor(transition)
|
|
processed_obs = result[TransitionIndex.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]
|