mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-15 16:49:55 +00:00
feat(processor): enhance type safety with generic DataProcessorPipeline for policy and robot pipelines (#1915)
* refactor(processor): enhance type annotations for processors in record, replay, teleoperate, and control utils - Updated type annotations for preprocessor and postprocessor parameters in record_loop and predict_action functions to specify the expected dictionary types. - Adjusted robot_action_processor type in ReplayConfig and TeleoperateConfig to improve clarity and maintainability. - Ensured consistency in type definitions across multiple files, enhancing overall code readability. * refactor(processor): enhance type annotations for RobotProcessorPipeline in various files - Updated type annotations for RobotProcessorPipeline instances in evaluate.py, record.py, replay.py, teleoperate.py, and other related files to specify input and output types more clearly. - Introduced new type conversions for PolicyAction and EnvTransition to improve type safety and maintainability across the processing pipelines. - Ensured consistency in type definitions, enhancing overall code readability and reducing potential runtime errors. * refactor(processor): update transition handling in processors to use transition_to_batch - Replaced direct transition handling with transition_to_batch in various processor tests and implementations to ensure consistent batching of input data. - Updated assertions in tests to reflect changes in data structure, enhancing clarity and maintainability. - Improved overall code readability by standardizing the way transitions are processed across different processor types. * refactor(tests): standardize transition key usage in processor tests - Updated assertions in processor test files to utilize the TransitionKey for action references, enhancing consistency across tests. - Replaced direct string references with TransitionKey constants for improved readability and maintainability. - Ensured that all relevant tests reflect these changes, contributing to a more uniform approach in handling transitions.
This commit is contained in:
Adil Zouitine
@@ -33,7 +33,7 @@ from lerobot.processor import (
|
||||
TransitionKey,
|
||||
UnnormalizerProcessorStep,
|
||||
)
|
||||
from lerobot.processor.converters import create_transition, identity_transition
|
||||
from lerobot.processor.converters import create_transition, transition_to_batch
|
||||
|
||||
|
||||
def create_default_config():
|
||||
@@ -69,8 +69,6 @@ def test_make_sac_processor_basic():
|
||||
preprocessor, postprocessor = make_sac_pre_post_processors(
|
||||
config,
|
||||
stats,
|
||||
preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
)
|
||||
|
||||
# Check processor names
|
||||
@@ -98,30 +96,28 @@ def test_sac_processor_normalization_modes():
|
||||
preprocessor, postprocessor = make_sac_pre_post_processors(
|
||||
config,
|
||||
stats,
|
||||
preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
)
|
||||
|
||||
# Create test data
|
||||
observation = {OBS_STATE: torch.randn(10) * 2} # Larger values to test normalization
|
||||
action = torch.rand(5) * 2 - 1 # Range [-1, 1]
|
||||
transition = create_transition(observation, action)
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
# Process through preprocessor
|
||||
processed = preprocessor(transition)
|
||||
processed = preprocessor(batch)
|
||||
|
||||
# Check that data is normalized and batched
|
||||
# State should be mean-std normalized
|
||||
# Action should be min-max normalized to [-1, 1]
|
||||
assert processed[TransitionKey.OBSERVATION][OBS_STATE].shape == (1, 10)
|
||||
assert processed[TransitionKey.ACTION].shape == (1, 5)
|
||||
assert processed[OBS_STATE].shape == (1, 10)
|
||||
assert processed[TransitionKey.ACTION.value].shape == (1, 5)
|
||||
|
||||
# Process action through postprocessor
|
||||
action_transition = create_transition(action=processed[TransitionKey.ACTION])
|
||||
postprocessed = postprocessor(action_transition)
|
||||
postprocessed = postprocessor(processed[TransitionKey.ACTION.value])
|
||||
|
||||
# Check that action is unnormalized (but still batched)
|
||||
assert postprocessed[TransitionKey.ACTION].shape == (1, 5)
|
||||
assert postprocessed.shape == (1, 5)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
@@ -134,28 +130,26 @@ def test_sac_processor_cuda():
|
||||
preprocessor, postprocessor = make_sac_pre_post_processors(
|
||||
config,
|
||||
stats,
|
||||
preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
)
|
||||
|
||||
# Create CPU data
|
||||
observation = {OBS_STATE: torch.randn(10)}
|
||||
action = torch.randn(5)
|
||||
transition = create_transition(observation, action)
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
# Process through preprocessor
|
||||
processed = preprocessor(transition)
|
||||
processed = preprocessor(batch)
|
||||
|
||||
# Check that data is on CUDA
|
||||
assert processed[TransitionKey.OBSERVATION][OBS_STATE].device.type == "cuda"
|
||||
assert processed[TransitionKey.ACTION].device.type == "cuda"
|
||||
assert processed[OBS_STATE].device.type == "cuda"
|
||||
assert processed[TransitionKey.ACTION.value].device.type == "cuda"
|
||||
|
||||
# Process through postprocessor
|
||||
action_transition = create_transition(action=processed[TransitionKey.ACTION])
|
||||
postprocessed = postprocessor(action_transition)
|
||||
postprocessed = postprocessor(processed[TransitionKey.ACTION.value])
|
||||
|
||||
# Check that action is back on CPU
|
||||
assert postprocessed[TransitionKey.ACTION].device.type == "cpu"
|
||||
assert postprocessed.device.type == "cpu"
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
@@ -168,8 +162,6 @@ def test_sac_processor_accelerate_scenario():
|
||||
preprocessor, postprocessor = make_sac_pre_post_processors(
|
||||
config,
|
||||
stats,
|
||||
preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
)
|
||||
|
||||
# Simulate Accelerate: data already on GPU
|
||||
@@ -177,13 +169,14 @@ def test_sac_processor_accelerate_scenario():
|
||||
observation = {OBS_STATE: torch.randn(10).to(device)}
|
||||
action = torch.randn(5).to(device)
|
||||
transition = create_transition(observation, action)
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
# Process through preprocessor
|
||||
processed = preprocessor(transition)
|
||||
processed = preprocessor(batch)
|
||||
|
||||
# Check that data stays on same GPU
|
||||
assert processed[TransitionKey.OBSERVATION][OBS_STATE].device == device
|
||||
assert processed[TransitionKey.ACTION].device == device
|
||||
assert processed[OBS_STATE].device == device
|
||||
assert processed[TransitionKey.ACTION.value].device == device
|
||||
|
||||
|
||||
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs")
|
||||
@@ -196,8 +189,6 @@ def test_sac_processor_multi_gpu():
|
||||
preprocessor, postprocessor = make_sac_pre_post_processors(
|
||||
config,
|
||||
stats,
|
||||
preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
)
|
||||
|
||||
# Simulate data on different GPU
|
||||
@@ -205,35 +196,21 @@ def test_sac_processor_multi_gpu():
|
||||
observation = {OBS_STATE: torch.randn(10).to(device)}
|
||||
action = torch.randn(5).to(device)
|
||||
transition = create_transition(observation, action)
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
# Process through preprocessor
|
||||
processed = preprocessor(transition)
|
||||
processed = preprocessor(batch)
|
||||
|
||||
# Check that data stays on cuda:1
|
||||
assert processed[TransitionKey.OBSERVATION][OBS_STATE].device == device
|
||||
assert processed[TransitionKey.ACTION].device == device
|
||||
assert processed[OBS_STATE].device == device
|
||||
assert processed[TransitionKey.ACTION.value].device == device
|
||||
|
||||
|
||||
def test_sac_processor_without_stats():
|
||||
"""Test SAC processor creation without dataset statistics."""
|
||||
config = create_default_config()
|
||||
|
||||
# Get the steps from the factory function
|
||||
factory_preprocessor, factory_postprocessor = make_sac_pre_post_processors(config, dataset_stats=None)
|
||||
|
||||
# Create new processors with EnvTransition input/output
|
||||
preprocessor = DataProcessorPipeline(
|
||||
factory_preprocessor.steps,
|
||||
name=factory_preprocessor.name,
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
)
|
||||
postprocessor = DataProcessorPipeline(
|
||||
factory_postprocessor.steps,
|
||||
name=factory_postprocessor.name,
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
)
|
||||
preprocessor, postprocessor = make_sac_pre_post_processors(config, dataset_stats=None)
|
||||
|
||||
# Should still create processors
|
||||
assert preprocessor is not None
|
||||
@@ -243,8 +220,9 @@ def test_sac_processor_without_stats():
|
||||
observation = {OBS_STATE: torch.randn(10)}
|
||||
action = torch.randn(5)
|
||||
transition = create_transition(observation, action)
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
processed = preprocessor(transition)
|
||||
processed = preprocessor(batch)
|
||||
assert processed is not None
|
||||
|
||||
|
||||
@@ -256,8 +234,6 @@ def test_sac_processor_save_and_load():
|
||||
preprocessor, postprocessor = make_sac_pre_post_processors(
|
||||
config,
|
||||
stats,
|
||||
preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
@@ -265,18 +241,17 @@ def test_sac_processor_save_and_load():
|
||||
preprocessor.save_pretrained(tmpdir)
|
||||
|
||||
# Load preprocessor
|
||||
loaded_preprocessor = DataProcessorPipeline.from_pretrained(
|
||||
tmpdir, to_transition=identity_transition, to_output=identity_transition
|
||||
)
|
||||
loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir)
|
||||
|
||||
# Test that loaded processor works
|
||||
observation = {OBS_STATE: torch.randn(10)}
|
||||
action = torch.randn(5)
|
||||
transition = create_transition(observation, action)
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
processed = loaded_preprocessor(transition)
|
||||
assert processed[TransitionKey.OBSERVATION][OBS_STATE].shape == (1, 10)
|
||||
assert processed[TransitionKey.ACTION].shape == (1, 5)
|
||||
processed = loaded_preprocessor(batch)
|
||||
assert processed[OBS_STATE].shape == (1, 10)
|
||||
assert processed[TransitionKey.ACTION.value].shape == (1, 5)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
@@ -290,8 +265,6 @@ def test_sac_processor_mixed_precision():
|
||||
preprocessor, postprocessor = make_sac_pre_post_processors(
|
||||
config,
|
||||
stats,
|
||||
preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
)
|
||||
|
||||
# Replace DeviceProcessorStep with one that uses float16
|
||||
@@ -301,11 +274,12 @@ def test_sac_processor_mixed_precision():
|
||||
modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16"))
|
||||
elif isinstance(step, NormalizerProcessorStep):
|
||||
# Update normalizer to use the same device as the device processor
|
||||
norm_step = step # Now type checker knows this is NormalizerProcessorStep
|
||||
modified_steps.append(
|
||||
NormalizerProcessorStep(
|
||||
features=step.features,
|
||||
norm_map=step.norm_map,
|
||||
stats=step.stats,
|
||||
features=norm_step.features,
|
||||
norm_map=norm_step.norm_map,
|
||||
stats=norm_step.stats,
|
||||
device=config.device,
|
||||
dtype=torch.float16, # Match the float16 dtype
|
||||
)
|
||||
@@ -318,13 +292,14 @@ def test_sac_processor_mixed_precision():
|
||||
observation = {OBS_STATE: torch.randn(10, dtype=torch.float32)}
|
||||
action = torch.randn(5, dtype=torch.float32)
|
||||
transition = create_transition(observation, action)
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
# Process through preprocessor
|
||||
processed = preprocessor(transition)
|
||||
processed = preprocessor(batch)
|
||||
|
||||
# Check that data is converted to float16
|
||||
assert processed[TransitionKey.OBSERVATION][OBS_STATE].dtype == torch.float16
|
||||
assert processed[TransitionKey.ACTION].dtype == torch.float16
|
||||
assert processed[OBS_STATE].dtype == torch.float16
|
||||
assert processed[TransitionKey.ACTION.value].dtype == torch.float16
|
||||
|
||||
|
||||
def test_sac_processor_batch_data():
|
||||
@@ -335,8 +310,6 @@ def test_sac_processor_batch_data():
|
||||
preprocessor, postprocessor = make_sac_pre_post_processors(
|
||||
config,
|
||||
stats,
|
||||
preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
)
|
||||
|
||||
# Test with batched data
|
||||
@@ -344,13 +317,14 @@ def test_sac_processor_batch_data():
|
||||
observation = {OBS_STATE: torch.randn(batch_size, 10)}
|
||||
action = torch.randn(batch_size, 5)
|
||||
transition = create_transition(observation, action)
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
# Process through preprocessor
|
||||
processed = preprocessor(transition)
|
||||
processed = preprocessor(batch)
|
||||
|
||||
# Check that batch dimension is preserved
|
||||
assert processed[TransitionKey.OBSERVATION][OBS_STATE].shape == (batch_size, 10)
|
||||
assert processed[TransitionKey.ACTION].shape == (batch_size, 5)
|
||||
assert processed[OBS_STATE].shape == (batch_size, 10)
|
||||
assert processed[TransitionKey.ACTION.value].shape == (batch_size, 5)
|
||||
|
||||
|
||||
def test_sac_processor_edge_cases():
|
||||
@@ -361,22 +335,24 @@ def test_sac_processor_edge_cases():
|
||||
preprocessor, postprocessor = make_sac_pre_post_processors(
|
||||
config,
|
||||
stats,
|
||||
preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
)
|
||||
|
||||
# Test with empty observation
|
||||
transition = create_transition(observation={}, action=torch.randn(5))
|
||||
processed = preprocessor(transition)
|
||||
assert processed[TransitionKey.OBSERVATION] == {}
|
||||
assert processed[TransitionKey.ACTION].shape == (1, 5)
|
||||
# Test with observation that has no state key but still exists
|
||||
observation = {"observation.dummy": torch.randn(1)} # Some dummy observation to pass validation
|
||||
action = torch.randn(5)
|
||||
batch = {TransitionKey.ACTION.value: action, **observation}
|
||||
processed = preprocessor(batch)
|
||||
# observation.state wasn't in original, so it won't be in processed
|
||||
assert OBS_STATE not in processed
|
||||
assert processed[TransitionKey.ACTION.value].shape == (1, 5)
|
||||
|
||||
# Test with zero action (representing "null" action)
|
||||
transition = create_transition(observation={OBS_STATE: torch.randn(10)}, action=torch.zeros(5))
|
||||
processed = preprocessor(transition)
|
||||
assert processed[TransitionKey.OBSERVATION][OBS_STATE].shape == (1, 10)
|
||||
batch = transition_to_batch(transition)
|
||||
processed = preprocessor(batch)
|
||||
assert processed[OBS_STATE].shape == (1, 10)
|
||||
# Action should be present and batched, even if it's zeros
|
||||
assert processed[TransitionKey.ACTION].shape == (1, 5)
|
||||
assert processed[TransitionKey.ACTION.value].shape == (1, 5)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
@@ -389,8 +365,6 @@ def test_sac_processor_bfloat16_device_float32_normalizer():
|
||||
preprocessor, _ = make_sac_pre_post_processors(
|
||||
config,
|
||||
stats,
|
||||
preprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
postprocessor_kwargs={"to_transition": identity_transition, "to_output": identity_transition},
|
||||
)
|
||||
|
||||
# Modify the pipeline to use bfloat16 device processor with float32 normalizer
|
||||
@@ -401,11 +375,12 @@ def test_sac_processor_bfloat16_device_float32_normalizer():
|
||||
modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="bfloat16"))
|
||||
elif isinstance(step, NormalizerProcessorStep):
|
||||
# Normalizer stays configured as float32 (will auto-adapt to bfloat16)
|
||||
norm_step = step # Now type checker knows this is NormalizerProcessorStep
|
||||
modified_steps.append(
|
||||
NormalizerProcessorStep(
|
||||
features=step.features,
|
||||
norm_map=step.norm_map,
|
||||
stats=step.stats,
|
||||
features=norm_step.features,
|
||||
norm_map=norm_step.norm_map,
|
||||
stats=norm_step.stats,
|
||||
device=config.device,
|
||||
dtype=torch.float32, # Deliberately configured as float32
|
||||
)
|
||||
@@ -422,13 +397,14 @@ def test_sac_processor_bfloat16_device_float32_normalizer():
|
||||
observation = {OBS_STATE: torch.randn(10, dtype=torch.float32)} # Start with float32
|
||||
action = torch.randn(5, dtype=torch.float32)
|
||||
transition = create_transition(observation, action)
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
# Process through full pipeline
|
||||
processed = preprocessor(transition)
|
||||
processed = preprocessor(batch)
|
||||
|
||||
# Verify: DeviceProcessor → bfloat16, NormalizerProcessor adapts → final output is bfloat16
|
||||
assert processed[TransitionKey.OBSERVATION][OBS_STATE].dtype == torch.bfloat16
|
||||
assert processed[TransitionKey.ACTION].dtype == torch.bfloat16
|
||||
assert processed[OBS_STATE].dtype == torch.bfloat16
|
||||
assert processed[TransitionKey.ACTION.value].dtype == torch.bfloat16
|
||||
|
||||
# Verify normalizer automatically adapted its internal state
|
||||
assert normalizer_step.dtype == torch.bfloat16
|
||||
|
||||
Reference in New Issue
Block a user