mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-15 16:49:55 +00:00
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot]
committed by
Adil Zouitine
Adil Zouitine
parent
f6c7287ae7
commit
769f531603
@@ -20,447 +20,447 @@ import torch
|
||||
|
||||
from lerobot.processor.observation_processor import (
|
||||
ImageProcessor,
|
||||
StateProcessor,
|
||||
ObservationProcessor,
|
||||
StateProcessor,
|
||||
)
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
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]]))
|
||||
|
||||
"""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]]))
|
||||
|
||||
"""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)
|
||||
|
||||
"""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"
|
||||
}
|
||||
|
||||
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"
|
||||
|
||||
"""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
|
||||
|
||||
"""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
|
||||
|
||||
"""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 == {}
|
||||
|
||||
"""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
|
||||
|
||||
"""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()
|
||||
|
||||
"""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
|
||||
"""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])
|
||||
|
||||
"""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])
|
||||
"""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])
|
||||
|
||||
Reference in New Issue
Block a user