mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-16 00:59:46 +00:00
Adil Zouitine
f6c7287ae7
- Introduced `RobotPipeline` and `ObservationProcessor` for handling observation transformations. - Updated `preprocess_observation` to maintain backward compatibility while leveraging the new pipeline. - Added tests for the new processing components and ensured they match the original functionality. - Removed hardcoded logic in favor of a more flexible, composable architecture.
408 lines
13 KiB
Python
408 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 json
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from dataclasses import dataclass
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from lerobot.processor.pipeline import RobotPipeline, EnvTransition, PipelineStep
|
|
|
|
|
|
@dataclass
|
|
class MockStep:
|
|
"""Mock pipeline step for testing - demonstrates best practices.
|
|
|
|
This example shows the proper separation:
|
|
- JSON-serializable attributes (name, counter) go in get_config()
|
|
- Only torch tensors go in state_dict()
|
|
|
|
Note: The counter is part of the configuration, so it will be restored
|
|
when the step is recreated from config during loading.
|
|
"""
|
|
|
|
name: str = "mock_step"
|
|
counter: int = 0
|
|
|
|
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
|
"""Add a counter to the complementary_data."""
|
|
obs, action, reward, done, truncated, info, comp_data = transition
|
|
|
|
if comp_data is None:
|
|
comp_data = {}
|
|
else:
|
|
comp_data = dict(comp_data) # Make a copy
|
|
|
|
comp_data[f"{self.name}_counter"] = self.counter
|
|
self.counter += 1
|
|
|
|
return (obs, action, reward, done, truncated, info, comp_data)
|
|
|
|
def get_config(self) -> dict[str, Any]:
|
|
# Return all JSON-serializable attributes that should be persisted
|
|
# These will be passed to __init__ when loading
|
|
return {"name": self.name, "counter": self.counter}
|
|
|
|
def state_dict(self) -> dict[str, torch.Tensor]:
|
|
# Only return torch tensors (empty in this case since we have no tensor state)
|
|
return {}
|
|
|
|
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
|
# No tensor state to load
|
|
pass
|
|
|
|
def reset(self) -> None:
|
|
self.counter = 0
|
|
|
|
|
|
@dataclass
|
|
class MockStepWithoutOptionalMethods:
|
|
"""Mock step that only implements the required __call__ method."""
|
|
|
|
multiplier: float = 2.0
|
|
|
|
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
|
"""Multiply reward by multiplier."""
|
|
obs, action, reward, done, truncated, info, comp_data = transition
|
|
|
|
if reward is not None:
|
|
reward = reward * self.multiplier
|
|
|
|
return (obs, action, reward, done, truncated, info, comp_data)
|
|
|
|
|
|
@dataclass
|
|
class MockStepWithTensorState:
|
|
"""Mock step demonstrating mixed JSON attributes and tensor state."""
|
|
|
|
name: str = "tensor_step"
|
|
learning_rate: float = 0.01
|
|
window_size: int = 10
|
|
|
|
def __init__(self, name: str = "tensor_step", learning_rate: float = 0.01, window_size: int = 10):
|
|
self.name = name
|
|
self.learning_rate = learning_rate
|
|
self.window_size = window_size
|
|
# Tensor state
|
|
self.running_mean = torch.zeros(window_size)
|
|
self.running_count = torch.tensor(0)
|
|
|
|
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
|
"""Update running statistics."""
|
|
obs, action, reward, done, truncated, info, comp_data = transition
|
|
|
|
if reward is not None:
|
|
# Update running mean
|
|
idx = self.running_count % self.window_size
|
|
self.running_mean[idx] = reward
|
|
self.running_count += 1
|
|
|
|
return transition
|
|
|
|
def get_config(self) -> dict[str, Any]:
|
|
# Only JSON-serializable attributes
|
|
return {
|
|
"name": self.name,
|
|
"learning_rate": self.learning_rate,
|
|
"window_size": self.window_size,
|
|
}
|
|
|
|
def state_dict(self) -> dict[str, torch.Tensor]:
|
|
# Only tensor state
|
|
return {
|
|
"running_mean": self.running_mean,
|
|
"running_count": self.running_count,
|
|
}
|
|
|
|
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
|
self.running_mean = state["running_mean"]
|
|
self.running_count = state["running_count"]
|
|
|
|
def reset(self) -> None:
|
|
self.running_mean.zero_()
|
|
self.running_count.zero_()
|
|
|
|
|
|
def test_empty_pipeline():
|
|
"""Test pipeline with no steps."""
|
|
pipeline = RobotPipeline()
|
|
|
|
transition = (None, None, 0.0, False, False, {}, {})
|
|
result = pipeline(transition)
|
|
|
|
assert result == transition
|
|
assert len(pipeline) == 0
|
|
|
|
def test_single_step_pipeline():
|
|
"""Test pipeline with a single step."""
|
|
step = MockStep("test_step")
|
|
pipeline = RobotPipeline([step])
|
|
|
|
transition = (None, None, 0.0, False, False, {}, {})
|
|
result = pipeline(transition)
|
|
|
|
assert len(pipeline) == 1
|
|
assert result[6]["test_step_counter"] == 0 # complementary_data
|
|
|
|
# Call again to test counter increment
|
|
result = pipeline(transition)
|
|
assert result[6]["test_step_counter"] == 1
|
|
|
|
def test_multiple_steps_pipeline():
|
|
"""Test pipeline with multiple steps."""
|
|
step1 = MockStep("step1")
|
|
step2 = MockStep("step2")
|
|
pipeline = RobotPipeline([step1, step2])
|
|
|
|
transition = (None, None, 0.0, False, False, {}, {})
|
|
result = pipeline(transition)
|
|
|
|
assert len(pipeline) == 2
|
|
assert result[6]["step1_counter"] == 0
|
|
assert result[6]["step2_counter"] == 0
|
|
|
|
def test_invalid_transition_format():
|
|
"""Test pipeline with invalid transition format."""
|
|
pipeline = RobotPipeline([MockStep()])
|
|
|
|
# Test with wrong number of elements
|
|
with pytest.raises(ValueError, match="EnvTransition must be a 7-tuple"):
|
|
pipeline((None, None, 0.0)) # Only 3 elements
|
|
|
|
# Test with wrong type
|
|
with pytest.raises(ValueError, match="EnvTransition must be a 7-tuple"):
|
|
pipeline("not a tuple")
|
|
|
|
def test_step_through():
|
|
"""Test step_through method."""
|
|
step1 = MockStep("step1")
|
|
step2 = MockStep("step2")
|
|
pipeline = RobotPipeline([step1, step2])
|
|
|
|
transition = (None, None, 0.0, False, False, {}, {})
|
|
|
|
results = list(pipeline.step_through(transition))
|
|
|
|
assert len(results) == 3 # Original + 2 steps
|
|
assert results[0] == transition # Original
|
|
assert "step1_counter" in results[1][6] # After step1
|
|
assert "step2_counter" in results[2][6] # After step2
|
|
|
|
def test_indexing():
|
|
"""Test pipeline indexing."""
|
|
step1 = MockStep("step1")
|
|
step2 = MockStep("step2")
|
|
pipeline = RobotPipeline([step1, step2])
|
|
|
|
# Test integer indexing
|
|
assert pipeline[0] is step1
|
|
assert pipeline[1] is step2
|
|
|
|
# Test slice indexing
|
|
sub_pipeline = pipeline[0:1]
|
|
assert isinstance(sub_pipeline, RobotPipeline)
|
|
assert len(sub_pipeline) == 1
|
|
assert sub_pipeline[0] is step1
|
|
|
|
def test_hooks():
|
|
"""Test before/after step hooks."""
|
|
step = MockStep("test_step")
|
|
pipeline = RobotPipeline([step])
|
|
|
|
before_calls = []
|
|
after_calls = []
|
|
|
|
def before_hook(idx: int, transition: EnvTransition):
|
|
before_calls.append(idx)
|
|
return transition
|
|
|
|
def after_hook(idx: int, transition: EnvTransition):
|
|
after_calls.append(idx)
|
|
return transition
|
|
|
|
pipeline.register_before_step_hook(before_hook)
|
|
pipeline.register_after_step_hook(after_hook)
|
|
|
|
transition = (None, None, 0.0, False, False, {}, {})
|
|
pipeline(transition)
|
|
|
|
assert before_calls == [0]
|
|
assert after_calls == [0]
|
|
|
|
def test_hook_modification():
|
|
"""Test that hooks can modify transitions."""
|
|
step = MockStep("test_step")
|
|
pipeline = RobotPipeline([step])
|
|
|
|
def modify_reward_hook(idx: int, transition: EnvTransition):
|
|
obs, action, reward, done, truncated, info, comp_data = transition
|
|
return (obs, action, 42.0, done, truncated, info, comp_data)
|
|
|
|
pipeline.register_before_step_hook(modify_reward_hook)
|
|
|
|
transition = (None, None, 0.0, False, False, {}, {})
|
|
result = pipeline(transition)
|
|
|
|
assert result[2] == 42.0 # reward modified by hook
|
|
|
|
def test_reset():
|
|
"""Test pipeline reset functionality."""
|
|
step = MockStep("test_step")
|
|
pipeline = RobotPipeline([step])
|
|
|
|
reset_called = []
|
|
|
|
def reset_hook():
|
|
reset_called.append(True)
|
|
|
|
pipeline.register_reset_hook(reset_hook)
|
|
|
|
# Make some calls to increment counter
|
|
transition = (None, None, 0.0, False, False, {}, {})
|
|
pipeline(transition)
|
|
pipeline(transition)
|
|
|
|
assert step.counter == 2
|
|
|
|
# Reset should reset step and call hook
|
|
pipeline.reset()
|
|
|
|
assert step.counter == 0
|
|
assert len(reset_called) == 1
|
|
|
|
def test_profile_steps():
|
|
"""Test step profiling functionality."""
|
|
step1 = MockStep("step1")
|
|
step2 = MockStep("step2")
|
|
pipeline = RobotPipeline([step1, step2])
|
|
|
|
transition = (None, None, 0.0, False, False, {}, {})
|
|
|
|
profile_results = pipeline.profile_steps(transition, num_runs=10)
|
|
|
|
assert len(profile_results) == 2
|
|
assert "step_0_MockStep" in profile_results
|
|
assert "step_1_MockStep" in profile_results
|
|
assert all(isinstance(time, float) and time >= 0 for time in profile_results.values())
|
|
|
|
def test_save_and_load_pretrained():
|
|
"""Test saving and loading pipeline.
|
|
|
|
This test demonstrates that JSON-serializable attributes (like counter)
|
|
are saved in the config and restored when the step is recreated.
|
|
"""
|
|
step1 = MockStep("step1")
|
|
step2 = MockStep("step2")
|
|
|
|
# Increment counters to have some state
|
|
step1.counter = 5
|
|
step2.counter = 10
|
|
|
|
pipeline = RobotPipeline([step1, step2], name="TestPipeline", seed=42)
|
|
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
# Save pipeline
|
|
pipeline.save_pretrained(tmp_dir)
|
|
|
|
# Check files were created
|
|
config_path = Path(tmp_dir) / "pipeline.json"
|
|
assert config_path.exists()
|
|
|
|
# Check config content
|
|
with open(config_path) as f:
|
|
config = json.load(f)
|
|
|
|
assert config["name"] == "TestPipeline"
|
|
assert config["seed"] == 42
|
|
assert len(config["steps"]) == 2
|
|
|
|
# Verify counters are saved in config, not in separate state files
|
|
assert config["steps"][0]["config"]["counter"] == 5
|
|
assert config["steps"][1]["config"]["counter"] == 10
|
|
|
|
# Load pipeline
|
|
loaded_pipeline = RobotPipeline.from_pretrained(tmp_dir)
|
|
|
|
assert loaded_pipeline.name == "TestPipeline"
|
|
assert loaded_pipeline.seed == 42
|
|
assert len(loaded_pipeline) == 2
|
|
|
|
# Check that counter was restored from config
|
|
assert loaded_pipeline.steps[0].counter == 5
|
|
assert loaded_pipeline.steps[1].counter == 10
|
|
|
|
def test_step_without_optional_methods():
|
|
"""Test pipeline with steps that don't implement optional methods."""
|
|
step = MockStepWithoutOptionalMethods(multiplier=3.0)
|
|
pipeline = RobotPipeline([step])
|
|
|
|
transition = (None, None, 2.0, False, False, {}, {})
|
|
result = pipeline(transition)
|
|
|
|
assert result[2] == 6.0 # 2.0 * 3.0
|
|
|
|
# Reset should work even if step doesn't implement reset
|
|
pipeline.reset()
|
|
|
|
# Save/load should work even without optional methods
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
pipeline.save_pretrained(tmp_dir)
|
|
loaded_pipeline = RobotPipeline.from_pretrained(tmp_dir)
|
|
assert len(loaded_pipeline) == 1
|
|
|
|
def test_mixed_json_and_tensor_state():
|
|
"""Test step with both JSON attributes and tensor state."""
|
|
step = MockStepWithTensorState(name="stats", learning_rate=0.05, window_size=5)
|
|
pipeline = RobotPipeline([step])
|
|
|
|
# Process some transitions with rewards
|
|
for i in range(10):
|
|
transition = (None, None, float(i), False, False, {}, {})
|
|
pipeline(transition)
|
|
|
|
# Check state
|
|
assert step.running_count.item() == 10
|
|
assert step.learning_rate == 0.05
|
|
|
|
# Save and load
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
pipeline.save_pretrained(tmp_dir)
|
|
|
|
# Check that both config and state files were created
|
|
config_path = Path(tmp_dir) / "pipeline.json"
|
|
state_path = Path(tmp_dir) / "step_0.safetensors"
|
|
assert config_path.exists()
|
|
assert state_path.exists()
|
|
|
|
# Load and verify
|
|
loaded_pipeline = RobotPipeline.from_pretrained(tmp_dir)
|
|
loaded_step = loaded_pipeline.steps[0]
|
|
|
|
# Check JSON attributes were restored
|
|
assert loaded_step.name == "stats"
|
|
assert loaded_step.learning_rate == 0.05
|
|
assert loaded_step.window_size == 5
|
|
|
|
# Check tensor state was restored
|
|
assert loaded_step.running_count.item() == 10
|
|
assert torch.allclose(loaded_step.running_mean, step.running_mean)
|
|
|
|
|