mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-16 00:59:46 +00:00
Refactor observation preprocessing to use a modular pipeline system
- 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.
This commit is contained in:
Adil Zouitine
@@ -0,0 +1,466 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 202 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.observation_processor import (
|
||||
ImageProcessor,
|
||||
StateProcessor,
|
||||
ObservationProcessor,
|
||||
)
|
||||
from lerobot.processor.pipeline import EnvTransition
|
||||
|
||||
|
||||
def test_process_single_image():
|
||||
"""Test processing a single image."""
|
||||
processor = ImageProcessor()
|
||||
|
||||
# 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 = (observation, None, None, None, None, None, None)
|
||||
|
||||
result = processor(transition)
|
||||
processed_obs = result[0]
|
||||
|
||||
# Check that the image was processed correctly
|
||||
assert "observation.image" in processed_obs
|
||||
processed_img = processed_obs["observation.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 = ImageProcessor()
|
||||
|
||||
# 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 = (observation, None, None, None, None, None, None)
|
||||
|
||||
result = processor(transition)
|
||||
processed_obs = result[0]
|
||||
|
||||
# Check that both images were processed
|
||||
assert "observation.images.camera1" in processed_obs
|
||||
assert "observation.images.camera2" in processed_obs
|
||||
|
||||
# Check shapes
|
||||
assert processed_obs["observation.images.camera1"].shape == (1, 3, 32, 32)
|
||||
assert processed_obs["observation.images.camera2"].shape == (1, 3, 48, 48)
|
||||
|
||||
def test_process_batched_image():
|
||||
"""Test processing already batched images."""
|
||||
processor = ImageProcessor()
|
||||
|
||||
# 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 = (observation, None, None, None, None, None, None)
|
||||
|
||||
result = processor(transition)
|
||||
processed_obs = result[0]
|
||||
|
||||
# Check that batch dimension is preserved
|
||||
assert processed_obs["observation.image"].shape == (2, 3, 64, 64)
|
||||
|
||||
def test_invalid_image_format():
|
||||
"""Test error handling for invalid image formats."""
|
||||
processor = ImageProcessor()
|
||||
|
||||
# Test wrong channel order (channels first)
|
||||
image = np.random.randint(0, 256, size=(3, 64, 64), dtype=np.uint8)
|
||||
observation = {"pixels": image}
|
||||
transition = (observation, None, None, None, None, None, None)
|
||||
|
||||
with pytest.raises(ValueError, match="Expected channel-last images"):
|
||||
processor(transition)
|
||||
|
||||
def test_invalid_image_dtype():
|
||||
"""Test error handling for invalid image dtype."""
|
||||
processor = ImageProcessor()
|
||||
|
||||
# Test wrong dtype
|
||||
image = np.random.rand(64, 64, 3).astype(np.float32)
|
||||
observation = {"pixels": image}
|
||||
transition = (observation, None, None, None, None, None, None)
|
||||
|
||||
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 = ImageProcessor()
|
||||
|
||||
observation = {"other_data": np.array([1, 2, 3])}
|
||||
transition = (observation, None, None, None, None, None, None)
|
||||
|
||||
result = processor(transition)
|
||||
processed_obs = result[0]
|
||||
|
||||
# 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 = ImageProcessor()
|
||||
|
||||
transition = (None, None, None, None, None, None, None)
|
||||
result = processor(transition)
|
||||
|
||||
assert result == transition
|
||||
|
||||
def test_serialization_methods():
|
||||
"""Test serialization methods."""
|
||||
processor = ImageProcessor()
|
||||
|
||||
# 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 = StateProcessor()
|
||||
|
||||
env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32)
|
||||
observation = {"environment_state": env_state}
|
||||
transition = (observation, None, None, None, None, None, None)
|
||||
|
||||
result = processor(transition)
|
||||
processed_obs = result[0]
|
||||
|
||||
# Check that environment_state was renamed and processed
|
||||
assert "observation.environment_state" in processed_obs
|
||||
assert "environment_state" not in processed_obs
|
||||
|
||||
processed_state = processed_obs["observation.environment_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 = StateProcessor()
|
||||
|
||||
agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32)
|
||||
observation = {"agent_pos": agent_pos}
|
||||
transition = (observation, None, None, None, None, None, None)
|
||||
|
||||
result = processor(transition)
|
||||
processed_obs = result[0]
|
||||
|
||||
# Check that agent_pos was renamed and processed
|
||||
assert "observation.state" in processed_obs
|
||||
assert "agent_pos" not in processed_obs
|
||||
|
||||
processed_state = processed_obs["observation.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 = StateProcessor()
|
||||
|
||||
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 = (observation, None, None, None, None, None, None)
|
||||
|
||||
result = processor(transition)
|
||||
processed_obs = result[0]
|
||||
|
||||
# Check that batch dimensions are preserved
|
||||
assert processed_obs["observation.environment_state"].shape == (2, 2)
|
||||
assert processed_obs["observation.state"].shape == (2, 2)
|
||||
|
||||
def test_process_both_states():
|
||||
"""Test processing both environment_state and agent_pos."""
|
||||
processor = StateProcessor()
|
||||
|
||||
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 = (observation, None, None, None, None, None, None)
|
||||
|
||||
result = processor(transition)
|
||||
processed_obs = result[0]
|
||||
|
||||
# Check that both states were processed
|
||||
assert "observation.environment_state" in processed_obs
|
||||
assert "observation.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 = StateProcessor()
|
||||
|
||||
observation = {"other_data": np.array([1, 2, 3])}
|
||||
transition = (observation, None, None, None, None, None, None)
|
||||
|
||||
result = processor(transition)
|
||||
processed_obs = result[0]
|
||||
|
||||
# Should preserve data unchanged
|
||||
assert processed_obs == observation
|
||||
|
||||
def test_none_observation():
|
||||
"""Test processor with None observation."""
|
||||
processor = StateProcessor()
|
||||
|
||||
transition = (None, None, None, None, None, None, None)
|
||||
result = processor(transition)
|
||||
|
||||
assert result == transition
|
||||
|
||||
def test_serialization_methods():
|
||||
"""Test serialization methods."""
|
||||
processor = StateProcessor()
|
||||
|
||||
# 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_complete_observation_processing():
|
||||
"""Test processing a complete observation with both images and states."""
|
||||
processor = ObservationProcessor()
|
||||
|
||||
# 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 = (observation, None, None, None, None, None, None)
|
||||
|
||||
result = processor(transition)
|
||||
processed_obs = result[0]
|
||||
|
||||
# Check that image was processed
|
||||
assert "observation.image" in processed_obs
|
||||
assert processed_obs["observation.image"].shape == (1, 3, 32, 32)
|
||||
|
||||
# Check that states were processed
|
||||
assert "observation.environment_state" in processed_obs
|
||||
assert "observation.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 = ObservationProcessor()
|
||||
|
||||
image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8)
|
||||
observation = {"pixels": image}
|
||||
transition = (observation, None, None, None, None, None, None)
|
||||
|
||||
result = processor(transition)
|
||||
processed_obs = result[0]
|
||||
|
||||
assert "observation.image" in processed_obs
|
||||
assert len(processed_obs) == 1
|
||||
|
||||
def test_state_only_processing():
|
||||
"""Test processing observation with only states."""
|
||||
processor = ObservationProcessor()
|
||||
|
||||
agent_pos = np.array([1.0, 2.0], dtype=np.float32)
|
||||
observation = {"agent_pos": agent_pos}
|
||||
transition = (observation, None, None, None, None, None, None)
|
||||
|
||||
result = processor(transition)
|
||||
processed_obs = result[0]
|
||||
|
||||
assert "observation.state" in processed_obs
|
||||
assert "agent_pos" not in processed_obs
|
||||
|
||||
def test_empty_observation():
|
||||
"""Test processing empty observation."""
|
||||
processor = ObservationProcessor()
|
||||
|
||||
observation = {}
|
||||
transition = (observation, None, None, None, None, None, None)
|
||||
|
||||
result = processor(transition)
|
||||
processed_obs = result[0]
|
||||
|
||||
assert processed_obs == {}
|
||||
|
||||
def test_none_observation():
|
||||
"""Test processing None observation."""
|
||||
processor = ObservationProcessor()
|
||||
|
||||
transition = (None, None, None, None, None, None, None)
|
||||
result = processor(transition)
|
||||
|
||||
assert result == transition
|
||||
|
||||
def test_serialization_methods():
|
||||
"""Test serialization methods."""
|
||||
processor = ObservationProcessor()
|
||||
|
||||
# Test get_config
|
||||
config = processor.get_config()
|
||||
assert isinstance(config, dict)
|
||||
assert "image_processor" in config
|
||||
assert "state_processor" in config
|
||||
|
||||
# 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_custom_sub_processors():
|
||||
"""Test ObservationProcessor with custom sub-processors."""
|
||||
image_proc = ImageProcessor()
|
||||
state_proc = StateProcessor()
|
||||
processor = ObservationProcessor(image_processor=image_proc, state_processor=state_proc)
|
||||
|
||||
# Should use the provided processors
|
||||
assert processor.image_processor is image_proc
|
||||
assert processor.state_processor is state_proc
|
||||
|
||||
|
||||
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 = ObservationProcessor()
|
||||
|
||||
# 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 = (observation, None, None, None, None, None, None)
|
||||
processor_result = processor(transition)[0]
|
||||
|
||||
# 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 = ObservationProcessor()
|
||||
|
||||
# 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 = (observation, None, None, None, None, None, None)
|
||||
processor_result = processor(transition)[0]
|
||||
|
||||
# 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])
|
||||
@@ -0,0 +1,408 @@
|
||||
#!/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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user