feat(processor): multiple improvements to the pipeline porting (#1749)

* [Port codebase pipeline] General fixes for RL and scripts (#1748)

* Refactor dataset configuration in documentation and codebase

- Updated dataset configuration keys from `dataset_root` to `root` and `num_episodes` to `num_episodes_to_record` for consistency.
- Adjusted replay episode handling by renaming `episode` to `replay_episode`.
- Enhanced documentation
- added specific processor to transform from policy actions to delta actions

* Added Robot action to tensor processor
Added new processor script for dealing with gym specific action processing

* removed RobotAction2Tensor processor; imrpoved choosing observations in actor

* nit in delta action

* added missing reset functions to kinematics

* Adapt teleoperate and replay to pipeline similar to record

* refactor(processors): move to inheritance (#1750)

* fix(teleoperator): improvements phone implementation (#1752)

* fix(teleoperator): protect shared state in phone implementation

* refactor(teleop): separate classes in phone

* fix: solve breaking changes (#1753)

* refactor(policies): multiple improvements (#1754)

* refactor(processor): simpler logic in device processor (#1755)

* refactor(processor): euclidean distance in delta action processor (#1757)

* refactor(processor): improvements to joint observations processor migration (#1758)

* refactor(processor): improvements to tokenizer migration (#1759)

* refactor(processor): improvements to tokenizer migration

* fix(tests): tokenizer tests regression from #1750

* fix(processors): fix float comparison and config in hil processors (#1760)

* chore(teleop): remove unnecessary callbacks in KeyboardEndEffectorTeleop (#1761)

* refactor(processor): improvements normalize pipeline migration (#1756)

* refactor(processor): several improvements normalize processor step

* refactor(processor): more improvements normalize processor

* refactor(processor): more changes to normalizer

* refactor(processor): take a different approach to DRY

* refactor(processor): final design

* chore(record): revert comment and continue deleted (#1764)

* refactor(examples): pipeline phone examples (#1769)

* refactor(examples): phone teleop + teleop script

* refactor(examples): phone replay + replay

* chore(examples): rename phone example files & folders

* feat(processor): fix improvements to the pipeline porting (#1796)

* refactor(processor): enhance tensor device handling in normalization process (#1795)

* refactor(tests): remove unsupported device detection test for complementary data (#1797)

* chore(tests): update ToBatchProcessor test (#1798)

* refactor(tests): remove in-place mutation tests for actions and complementary data in batch processor

* test(tests): add tests for action and task processing in batch processor

* add names for android and ios phone (#1799)

* use _tensor_stats in normalize processor (#1800)

* fix(normalize_processor): correct device reference for tensor epsilon handling (#1801)

* add point 5 add missing feature contracts (#1806)

* Fix PR comments 1452 (#1807)

* use key to determine image

* Address rest of PR comments

* use PolicyFeatures in transform_features

---------

Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>

---------

Co-authored-by: Michel Aractingi <michel.aractingi@huggingface.co>
Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com>
Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
This commit is contained in:
Steven Palma
2025-08-31 20:38:52 +02:00
committed by GitHub
parent 35c5d43255
commit ce665160ae
55 changed files with 1549 additions and 2024 deletions
+11 -11
View File
@@ -23,7 +23,7 @@ import torch
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.constants import ACTION, OBS_STATE
from lerobot.policies.act.configuration_act import ACTConfig
from lerobot.policies.act.processor_act import make_act_processor
from lerobot.policies.act.processor_act import make_act_pre_post_processors
from lerobot.processor import (
DeviceProcessor,
NormalizerProcessor,
@@ -78,7 +78,7 @@ def test_make_act_processor_basic():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_act_processor(config, stats)
preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
# Check processor names
assert preprocessor.name == "robot_preprocessor"
@@ -102,7 +102,7 @@ def test_act_processor_normalization():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_act_processor(config, stats)
preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
# Create test data
observation = {OBS_STATE: torch.randn(7)}
@@ -131,7 +131,7 @@ def test_act_processor_cuda():
config.device = "cuda"
stats = create_default_stats()
preprocessor, postprocessor = make_act_processor(config, stats)
preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
# Create CPU data
observation = {OBS_STATE: torch.randn(7)}
@@ -160,7 +160,7 @@ def test_act_processor_accelerate_scenario():
config.device = "cuda:0"
stats = create_default_stats()
preprocessor, postprocessor = make_act_processor(config, stats)
preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
# Simulate Accelerate: data already on GPU
device = torch.device("cuda:0")
@@ -183,7 +183,7 @@ def test_act_processor_multi_gpu():
config.device = "cuda:0"
stats = create_default_stats()
preprocessor, postprocessor = make_act_processor(config, stats)
preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
# Simulate data on different GPU (like in multi-GPU training)
device = torch.device("cuda:1")
@@ -203,7 +203,7 @@ def test_act_processor_without_stats():
"""Test ACT processor creation without dataset statistics."""
config = create_default_config()
preprocessor, postprocessor = make_act_processor(config, dataset_stats=None)
preprocessor, postprocessor = make_act_pre_post_processors(config, dataset_stats=None)
# Should still create processors, but normalization won't have stats
assert preprocessor is not None
@@ -223,7 +223,7 @@ def test_act_processor_save_and_load():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_act_processor(config, stats)
preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
with tempfile.TemporaryDirectory() as tmpdir:
# Save preprocessor
@@ -249,7 +249,7 @@ def test_act_processor_device_placement_preservation():
# Test with CPU config
config.device = "cpu"
preprocessor, _ = make_act_processor(config, stats)
preprocessor, _ = make_act_pre_post_processors(config, stats)
# Process CPU data
observation = {OBS_STATE: torch.randn(7)}
@@ -269,7 +269,7 @@ def test_act_processor_mixed_precision():
stats = create_default_stats()
# Modify the device processor to use float16
preprocessor, postprocessor = make_act_processor(config, stats)
preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
# Replace DeviceProcessor with one that uses float16
for i, step in enumerate(preprocessor.steps):
@@ -294,7 +294,7 @@ def test_act_processor_batch_consistency():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_act_processor(config, stats)
preprocessor, postprocessor = make_act_pre_post_processors(config, stats)
# Test single sample (unbatched)
observation = {OBS_STATE: torch.randn(7)}
+46 -39
View File
@@ -603,24 +603,6 @@ def test_action_dtype_preservation():
assert result[TransitionKey.ACTION].shape == (1, 4)
def test_action_in_place_mutation():
"""Test that the processor mutates the transition in place for actions."""
processor = ToBatchProcessor()
action = torch.randn(4)
transition = create_transition(action=action)
# Store reference to original transition
original_transition = transition
# Process
result = processor(transition)
# Should be the same object (in-place mutation)
assert result is original_transition
assert result[TransitionKey.ACTION].shape == (1, 4)
def test_empty_action_tensor():
"""Test handling of empty action tensors."""
processor = ToBatchProcessor()
@@ -851,27 +833,6 @@ def test_task_comprehensive_string_cases():
processed_comp_data = result[TransitionKey.COMPLEMENTARY_DATA]
assert processed_comp_data["task"] == task_list
assert isinstance(processed_comp_data["task"], list)
assert processed_comp_data["task"] is task_list # Should be same object (in-place)
def test_task_in_place_mutation():
"""Test that the processor mutates complementary_data in place for tasks."""
processor = ToBatchProcessor()
complementary_data = {"task": "sort_objects"}
transition = create_transition(complementary_data=complementary_data)
# Store reference to original transition and complementary_data
original_transition = transition
original_comp_data = complementary_data
# Process
result = processor(transition)
# Should be the same objects (in-place mutation)
assert result is original_transition
assert result[TransitionKey.COMPLEMENTARY_DATA] is original_comp_data
assert original_comp_data["task"] == ["sort_objects"]
def test_task_preserves_other_keys():
@@ -1127,3 +1088,49 @@ def test_empty_index_tensor():
# Should remain unchanged (already 1D)
assert result[TransitionKey.COMPLEMENTARY_DATA]["index"].shape == (0,)
def test_action_processing_creates_new_transition():
"""Test that the processor creates a new transition object with correctly processed action."""
processor = ToBatchProcessor()
action = torch.randn(4)
transition = create_transition(action=action)
# Store reference to original transition
original_transition = transition
# Process
result = processor(transition)
# Should be a different object (functional design, not in-place mutation)
assert result is not original_transition
# Original transition should remain unchanged
assert original_transition[TransitionKey.ACTION].shape == (4,)
# Result should have correctly processed action with batch dimension
assert result[TransitionKey.ACTION].shape == (1, 4)
assert torch.equal(result[TransitionKey.ACTION][0], action)
def test_task_processing_creates_new_transition():
"""Test that the processor creates a new transition object with correctly processed task."""
processor = ToBatchProcessor()
complementary_data = {"task": "sort_objects"}
transition = create_transition(complementary_data=complementary_data)
# Store reference to original transition and complementary_data
original_transition = transition
original_comp_data = complementary_data
# Process
result = processor(transition)
# Should be different transition object (functional design)
assert result is not original_transition
# But complementary_data is the same reference (current implementation behavior)
assert result[TransitionKey.COMPLEMENTARY_DATA] is original_comp_data
# The task should be processed correctly (wrapped in list)
assert result[TransitionKey.COMPLEMENTARY_DATA]["task"] == ["sort_objects"]
# Original complementary data is also modified (current behavior)
assert original_comp_data["task"] == ["sort_objects"]
+11 -11
View File
@@ -23,7 +23,7 @@ import torch
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE
from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
from lerobot.policies.diffusion.processor_diffusion import make_diffusion_processor
from lerobot.policies.diffusion.processor_diffusion import make_diffusion_pre_post_processors
from lerobot.processor import (
DeviceProcessor,
NormalizerProcessor,
@@ -81,7 +81,7 @@ def test_make_diffusion_processor_basic():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_diffusion_processor(config, stats)
preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
# Check processor names
assert preprocessor.name == "robot_preprocessor"
@@ -105,7 +105,7 @@ def test_diffusion_processor_with_images():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_diffusion_processor(config, stats)
preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
# Create test data with images
observation = {
@@ -131,7 +131,7 @@ def test_diffusion_processor_cuda():
config.device = "cuda"
stats = create_default_stats()
preprocessor, postprocessor = make_diffusion_processor(config, stats)
preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
# Create CPU data
observation = {
@@ -164,7 +164,7 @@ def test_diffusion_processor_accelerate_scenario():
config.device = "cuda:0"
stats = create_default_stats()
preprocessor, postprocessor = make_diffusion_processor(config, stats)
preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
# Simulate Accelerate: data already on GPU
device = torch.device("cuda:0")
@@ -191,7 +191,7 @@ def test_diffusion_processor_multi_gpu():
config.device = "cuda:0"
stats = create_default_stats()
preprocessor, postprocessor = make_diffusion_processor(config, stats)
preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
# Simulate data on different GPU
device = torch.device("cuda:1")
@@ -215,7 +215,7 @@ def test_diffusion_processor_without_stats():
"""Test Diffusion processor creation without dataset statistics."""
config = create_default_config()
preprocessor, postprocessor = make_diffusion_processor(config, dataset_stats=None)
preprocessor, postprocessor = make_diffusion_pre_post_processors(config, dataset_stats=None)
# Should still create processors
assert preprocessor is not None
@@ -238,7 +238,7 @@ def test_diffusion_processor_save_and_load():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_diffusion_processor(config, stats)
preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
with tempfile.TemporaryDirectory() as tmpdir:
# Save preprocessor
@@ -269,7 +269,7 @@ def test_diffusion_processor_mixed_precision():
stats = create_default_stats()
# Create processor
preprocessor, postprocessor = make_diffusion_processor(config, stats)
preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
# Replace DeviceProcessor with one that uses float16
for i, step in enumerate(preprocessor.steps):
@@ -298,7 +298,7 @@ def test_diffusion_processor_identity_normalization():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_diffusion_processor(config, stats)
preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
# Create test data
image_value = torch.rand(3, 224, 224) * 255 # Large values
@@ -322,7 +322,7 @@ def test_diffusion_processor_batch_consistency():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_diffusion_processor(config, stats)
preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats)
# Test with different batch sizes
for batch_size in [1, 8, 32]:
+48 -315
View File
@@ -25,7 +25,6 @@ from lerobot.processor.normalize_processor import (
UnnormalizerProcessor,
_convert_stats_to_tensors,
hotswap_stats,
rename_stats,
)
from lerobot.processor.pipeline import IdentityProcessor, RobotProcessor, TransitionKey
@@ -182,7 +181,10 @@ def test_selective_normalization(observation_stats):
features = _create_observation_features()
norm_map = _create_observation_norm_map()
normalizer = NormalizerProcessor(
features=features, norm_map=norm_map, stats=observation_stats, normalize_keys={"observation.image"}
features=features,
norm_map=norm_map,
stats=observation_stats,
normalize_observation_keys={"observation.image"},
)
observation = {
@@ -243,6 +245,7 @@ def test_from_lerobot_dataset():
def test_state_dict_save_load(observation_normalizer):
# Save state
state_dict = observation_normalizer.state_dict()
print("State dict:", state_dict)
# Create new normalizer and load state
features = _create_observation_features()
@@ -464,10 +467,10 @@ def test_processor_from_lerobot_dataset(full_stats):
norm_map = _create_full_norm_map()
processor = NormalizerProcessor.from_lerobot_dataset(
mock_dataset, features, norm_map, normalize_keys={"observation.image"}
mock_dataset, features, norm_map, normalize_observation_keys={"observation.image"}
)
assert processor.normalize_keys == {"observation.image"}
assert processor.normalize_observation_keys == {"observation.image"}
assert "observation.image" in processor._tensor_stats
assert "action" in processor._tensor_stats
@@ -476,12 +479,16 @@ def test_get_config(full_stats):
features = _create_full_features()
norm_map = _create_full_norm_map()
processor = NormalizerProcessor(
features=features, norm_map=norm_map, stats=full_stats, normalize_keys={"observation.image"}, eps=1e-6
features=features,
norm_map=norm_map,
stats=full_stats,
normalize_observation_keys={"observation.image"},
eps=1e-6,
)
config = processor.get_config()
expected_config = {
"normalize_keys": ["observation.image"],
"normalize_observation_keys": ["observation.image"],
"eps": 1e-6,
"features": {
"observation.image": {"type": "VISUAL", "shape": (3, 96, 96)},
@@ -580,7 +587,11 @@ def test_serialization_roundtrip(full_stats):
features = _create_full_features()
norm_map = _create_full_norm_map()
original_processor = NormalizerProcessor(
features=features, norm_map=norm_map, stats=full_stats, normalize_keys={"observation.image"}, eps=1e-6
features=features,
norm_map=norm_map,
stats=full_stats,
normalize_observation_keys={"observation.image"},
eps=1e-6,
)
# Get config (serialization)
@@ -591,7 +602,7 @@ def test_serialization_roundtrip(full_stats):
features=config["features"],
norm_map=config["norm_map"],
stats=full_stats,
normalize_keys=set(config["normalize_keys"]),
normalize_observation_keys=set(config["normalize_observation_keys"]),
eps=config["eps"],
)
@@ -939,31 +950,31 @@ def test_identity_config_serialization():
assert torch.allclose(result1[TransitionKey.ACTION], result2[TransitionKey.ACTION])
def test_unsupported_normalization_mode_error():
"""Test that unsupported normalization modes raise appropriate errors."""
features = {"observation.state": PolicyFeature(FeatureType.STATE, (2,))}
# def test_unsupported_normalization_mode_error():
# """Test that unsupported normalization modes raise appropriate errors."""
# features = {"observation.state": PolicyFeature(FeatureType.STATE, (2,))}
# Create an invalid norm_map (this would never happen in practice, but tests error handling)
from enum import Enum
# # Create an invalid norm_map (this would never happen in practice, but tests error handling)
# from enum import Enum
class InvalidMode(str, Enum):
INVALID = "INVALID"
# class InvalidMode(str, Enum):
# INVALID = "INVALID"
# We can't actually pass an invalid enum to the processor due to type checking,
# but we can test the error by manipulating the norm_map after creation
norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD}
stats = {"observation.state": {"mean": [0.0, 0.0], "std": [1.0, 1.0]}}
# # We can't actually pass an invalid enum to the processor due to type checking,
# # but we can test the error by manipulating the norm_map after creation
# norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD}
# stats = {"observation.state": {"mean": [0.0, 0.0], "std": [1.0, 1.0]}}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
# normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
# Manually inject an invalid mode to test error handling
normalizer.norm_map[FeatureType.STATE] = "INVALID_MODE"
# # Manually inject an invalid mode to test error handling
# normalizer.norm_map[FeatureType.STATE] = "INVALID_MODE"
observation = {"observation.state": torch.tensor([1.0, -0.5])}
transition = create_transition(observation=observation)
# observation = {"observation.state": torch.tensor([1.0, -0.5])}
# transition = create_transition(observation=observation)
with pytest.raises(ValueError, match="Unsupported normalization mode"):
normalizer(transition)
# with pytest.raises(ValueError, match="Unsupported normalization mode"):
# normalizer(transition)
def test_hotswap_stats_basic_functionality():
@@ -1149,11 +1160,15 @@ def test_hotswap_stats_preserves_other_attributes():
"observation.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128)),
}
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
normalize_keys = {"observation.image"}
normalize_observation_keys = {"observation.image"}
eps = 1e-6
normalizer = NormalizerProcessor(
features=features, norm_map=norm_map, stats=initial_stats, normalize_keys=normalize_keys, eps=eps
features=features,
norm_map=norm_map,
stats=initial_stats,
normalize_observation_keys=normalize_observation_keys,
eps=eps,
)
robot_processor = RobotProcessor(steps=[normalizer])
@@ -1164,7 +1179,7 @@ def test_hotswap_stats_preserves_other_attributes():
new_normalizer = new_processor.steps[0]
assert new_normalizer.features == features
assert new_normalizer.norm_map == norm_map
assert new_normalizer.normalize_keys == normalize_keys
assert new_normalizer.normalize_observation_keys == normalize_observation_keys
assert new_normalizer.eps == eps
# But stats should be updated
@@ -1270,273 +1285,6 @@ def test_hotswap_stats_with_different_data_types():
torch.testing.assert_close(tensor_stats["observation.image"]["max"], torch.tensor(1.0))
def test_normalization_info_tracking():
"""Test that normalization info is tracked in complementary_data."""
features = {
"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96)),
"observation.state": PolicyFeature(FeatureType.STATE, (2,)),
"action": PolicyFeature(FeatureType.ACTION, (2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
FeatureType.STATE: NormalizationMode.MIN_MAX,
FeatureType.ACTION: NormalizationMode.IDENTITY,
}
stats = {
"observation.image": {
"mean": np.array([0.5, 0.5, 0.5]),
"std": np.array([0.2, 0.2, 0.2]),
},
"observation.state": {
"min": np.array([0.0, -1.0]),
"max": np.array([1.0, 1.0]),
},
"action": {
"mean": np.array([0.0, 0.0]),
"std": np.array([1.0, 1.0]),
},
}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]),
"observation.state": torch.tensor([0.5, 0.0]),
}
action = torch.tensor([1.0, -0.5])
transition = create_transition(observation=observation, action=action)
# Process the transition
normalized_transition = normalizer(transition)
# Check that normalization info is added
comp_data = normalized_transition.get(TransitionKey.COMPLEMENTARY_DATA)
assert comp_data is not None
assert "normalized_keys" in comp_data
norm_info = comp_data["normalized_keys"]
assert norm_info["observation.image"] == "MEAN_STD"
assert norm_info["observation.state"] == "MIN_MAX"
assert norm_info["action"] == "IDENTITY"
def test_unnormalization_info_tracking():
"""Test that unnormalization info is tracked in complementary_data."""
features = {
"observation.image": PolicyFeature(FeatureType.VISUAL, (3,)),
"action": PolicyFeature(FeatureType.ACTION, (2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
FeatureType.ACTION: NormalizationMode.MIN_MAX,
}
stats = {
"observation.image": {
"mean": np.array([0.5, 0.5, 0.5]),
"std": np.array([0.2, 0.2, 0.2]),
},
"action": {
"min": np.array([-1.0, -1.0]),
"max": np.array([1.0, 1.0]),
},
}
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
observation = {"observation.image": torch.tensor([0.7, 0.5, 0.3])}
action = torch.tensor([0.0, -0.5])
transition = create_transition(observation=observation, action=action)
# Process the transition
unnormalized_transition = unnormalizer(transition)
# Check that unnormalization info is added
comp_data = unnormalized_transition.get(TransitionKey.COMPLEMENTARY_DATA)
assert comp_data is not None
assert "unnormalized_keys" in comp_data
unnorm_info = comp_data["unnormalized_keys"]
assert unnorm_info["observation.image"] == "MEAN_STD"
assert unnorm_info["action"] == "MIN_MAX"
def test_normalization_info_with_missing_stats():
"""Test normalization info when stats are missing for some keys."""
features = {
"observation.image": PolicyFeature(FeatureType.VISUAL, (3,)),
"observation.state": PolicyFeature(FeatureType.STATE, (2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
FeatureType.STATE: NormalizationMode.MIN_MAX,
}
# Only provide stats for image, not state
stats = {
"observation.image": {
"mean": np.array([0.5, 0.5, 0.5]),
"std": np.array([0.2, 0.2, 0.2]),
},
}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]),
"observation.state": torch.tensor([0.5, 0.0]),
}
transition = create_transition(observation=observation)
# Process the transition
normalized_transition = normalizer(transition)
# Check that only keys with stats are in normalization info
comp_data = normalized_transition.get(TransitionKey.COMPLEMENTARY_DATA)
assert comp_data is not None
assert "normalized_keys" in comp_data
norm_info = comp_data["normalized_keys"]
assert norm_info["observation.image"] == "MEAN_STD"
# State should not be in the normalization info since it has no stats
assert "observation.state" not in norm_info
def test_normalization_info_with_selective_keys():
"""Test normalization info with selective normalization."""
features = {
"observation.image": PolicyFeature(FeatureType.VISUAL, (3,)),
"observation.state": PolicyFeature(FeatureType.STATE, (2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
FeatureType.STATE: NormalizationMode.MIN_MAX,
}
stats = {
"observation.image": {
"mean": np.array([0.5, 0.5, 0.5]),
"std": np.array([0.2, 0.2, 0.2]),
},
"observation.state": {
"min": np.array([0.0, -1.0]),
"max": np.array([1.0, 1.0]),
},
}
# Only normalize image
normalizer = NormalizerProcessor(
features=features, norm_map=norm_map, stats=stats, normalize_keys={"observation.image"}
)
observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]),
"observation.state": torch.tensor([0.5, 0.0]),
}
transition = create_transition(observation=observation)
# Process the transition
normalized_transition = normalizer(transition)
# Check that only selected keys are in normalization info
comp_data = normalized_transition.get(TransitionKey.COMPLEMENTARY_DATA)
assert comp_data is not None
assert "normalized_keys" in comp_data
norm_info = comp_data["normalized_keys"]
assert norm_info["observation.image"] == "MEAN_STD"
# State should not be in the normalization info since it wasn't in normalize_keys
assert "observation.state" not in norm_info
def test_normalization_info_preserved_in_pipeline():
"""Test that normalization info is preserved when using RobotProcessor pipeline."""
features = {
"observation.image": PolicyFeature(FeatureType.VISUAL, (3,)),
"action": PolicyFeature(FeatureType.ACTION, (2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
FeatureType.ACTION: NormalizationMode.MIN_MAX,
}
stats = {
"observation.image": {
"mean": np.array([0.5, 0.5, 0.5]),
"std": np.array([0.2, 0.2, 0.2]),
},
"action": {
"min": np.array([-1.0, -1.0]),
"max": np.array([1.0, 1.0]),
},
}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
# Create pipeline
pipeline = RobotProcessor([normalizer, unnormalizer])
observation = {"observation.image": torch.tensor([0.7, 0.5, 0.3])}
action = torch.tensor([0.5, -0.5])
transition = create_transition(observation=observation, action=action)
# Process through pipeline
result = pipeline(transition)
# Check that both normalization and unnormalization info are present
comp_data = result.get(TransitionKey.COMPLEMENTARY_DATA)
assert comp_data is not None
assert "normalized_keys" in comp_data
assert "unnormalized_keys" in comp_data
# Check normalization info
norm_info = comp_data["normalized_keys"]
assert norm_info["observation.image"] == "MEAN_STD"
assert norm_info["action"] == "MIN_MAX"
# Check unnormalization info
unnorm_info = comp_data["unnormalized_keys"]
assert unnorm_info["observation.image"] == "MEAN_STD"
assert unnorm_info["action"] == "MIN_MAX"
def test_normalization_info_empty_transition():
"""Test that no normalization info is added for empty transitions."""
features = {
"observation.image": PolicyFeature(FeatureType.VISUAL, (3,)),
"action": PolicyFeature(FeatureType.ACTION, (2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
FeatureType.ACTION: NormalizationMode.MIN_MAX,
}
stats = {
"observation.image": {"mean": [0.5], "std": [0.2]},
"action": {"min": [-1.0], "max": [1.0]},
}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
# Empty transition
transition = create_transition()
# Process the transition
normalized_transition = normalizer(transition)
# Check that no normalization info is added
comp_data = normalized_transition.get(TransitionKey.COMPLEMENTARY_DATA)
assert comp_data is None or "normalized_keys" not in comp_data
def test_hotswap_stats_functional_test():
"""Test that hotswapped processor actually works functionally."""
# Create test data
@@ -1631,8 +1379,8 @@ def test_min_equals_max_maps_to_minus_one():
assert torch.allclose(out[TransitionKey.OBSERVATION]["observation.state"], torch.tensor([-1.0]))
def test_action_normalized_despite_normalize_keys():
"""Action normalization is independent of normalize_keys filter for observations."""
def test_action_normalized_despite_normalize_observation_keys():
"""Action normalization is independent of normalize_observation_keys filter for observations."""
features = {
"observation.state": PolicyFeature(FeatureType.STATE, (1,)),
"action": PolicyFeature(FeatureType.ACTION, (2,)),
@@ -1640,7 +1388,7 @@ def test_action_normalized_despite_normalize_keys():
norm_map = {FeatureType.STATE: NormalizationMode.IDENTITY, FeatureType.ACTION: NormalizationMode.MEAN_STD}
stats = {"action": {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}}
normalizer = NormalizerProcessor(
features=features, norm_map=norm_map, stats=stats, normalize_keys={"observation.state"}
features=features, norm_map=norm_map, stats=stats, normalize_observation_keys={"observation.state"}
)
transition = create_transition(
@@ -1680,19 +1428,6 @@ def test_unnormalize_observations_mean_std_and_min_max():
assert torch.allclose(out_mm, torch.tensor([1.0, 0.0])) # mid of [0,2] and [-2,2]
def test_rename_stats_basic():
orig = {
"observation.state": {"mean": np.array([0.0]), "std": np.array([1.0])},
"action": {"mean": np.array([0.0])},
}
mapping = {"observation.state": "observation.robot_state"}
renamed = rename_stats(orig, mapping)
assert "observation.robot_state" in renamed and "observation.state" not in renamed
# Ensure deep copy: mutate original and verify renamed unaffected
orig["observation.state"]["mean"][0] = 42.0
assert renamed["observation.robot_state"]["mean"][0] != 42.0
def test_unknown_observation_keys_ignored():
features = {"observation.state": PolicyFeature(FeatureType.STATE, (1,))}
norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD}
@@ -1705,8 +1440,6 @@ def test_unknown_observation_keys_ignored():
# Unknown key should pass through unchanged and not be tracked
assert torch.allclose(out[TransitionKey.OBSERVATION]["observation.unknown"], obs["observation.unknown"])
comp = out.get(TransitionKey.COMPLEMENTARY_DATA) or {}
assert "normalized_keys" in comp and "observation.unknown" not in comp["normalized_keys"]
def test_batched_action_normalization():
@@ -1731,7 +1464,7 @@ def test_complementary_data_preservation():
tr = create_transition(observation={"observation.state": torch.tensor([1.0])}, complementary_data=comp)
out = normalizer(tr)
new_comp = out[TransitionKey.COMPLEMENTARY_DATA]
assert new_comp["existing"] == 123 and "normalized_keys" in new_comp
assert new_comp["existing"] == 123
def test_roundtrip_normalize_unnormalize_non_identity():
+6 -6
View File
@@ -23,7 +23,7 @@ import torch
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE
from lerobot.policies.pi0.configuration_pi0 import PI0Config
from lerobot.policies.pi0.processor_pi0 import Pi0NewLineProcessor, make_pi0_processor
from lerobot.policies.pi0.processor_pi0 import Pi0NewLineProcessor, make_pi0_pre_post_processors
from lerobot.processor import (
DeviceProcessor,
NormalizerProcessor,
@@ -84,7 +84,7 @@ def test_make_pi0_processor_basic():
stats = create_default_stats()
with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor"):
preprocessor, postprocessor = make_pi0_processor(config, stats)
preprocessor, postprocessor = make_pi0_pre_post_processors(config, stats)
# Check processor names
assert preprocessor.name == "robot_preprocessor"
@@ -183,7 +183,7 @@ def test_pi0_processor_cuda():
return features
with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor):
preprocessor, postprocessor = make_pi0_processor(config, stats)
preprocessor, postprocessor = make_pi0_pre_post_processors(config, stats)
# Create CPU data
observation = {
@@ -233,7 +233,7 @@ def test_pi0_processor_accelerate_scenario():
return features
with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor):
preprocessor, postprocessor = make_pi0_processor(config, stats)
preprocessor, postprocessor = make_pi0_pre_post_processors(config, stats)
# Simulate Accelerate: data already on GPU and batched
device = torch.device("cuda:0")
@@ -284,7 +284,7 @@ def test_pi0_processor_multi_gpu():
return features
with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor):
preprocessor, postprocessor = make_pi0_processor(config, stats)
preprocessor, postprocessor = make_pi0_pre_post_processors(config, stats)
# Simulate data on different GPU
device = torch.device("cuda:1")
@@ -310,7 +310,7 @@ def test_pi0_processor_without_stats():
# Mock the tokenizer processor
with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor"):
preprocessor, postprocessor = make_pi0_processor(config, dataset_stats=None)
preprocessor, postprocessor = make_pi0_pre_post_processors(config, dataset_stats=None)
# Should still create processors
assert preprocessor is not None
+14
View File
@@ -21,6 +21,7 @@ import torch
from lerobot.configs.types import FeatureType
from lerobot.processor import ProcessorStepRegistry, RenameProcessor, RobotProcessor, TransitionKey
from lerobot.processor.rename_processor import rename_stats
from tests.conftest import assert_contract_is_typed
@@ -465,3 +466,16 @@ def test_features_chained_processors(policy_feature_factory):
assert out["observation.image"] == spec["img"]
assert out["extra"] == spec["extra"]
assert_contract_is_typed(out)
def test_rename_stats_basic():
orig = {
"observation.state": {"mean": np.array([0.0]), "std": np.array([1.0])},
"action": {"mean": np.array([0.0])},
}
mapping = {"observation.state": "observation.robot_state"}
renamed = rename_stats(orig, mapping)
assert "observation.robot_state" in renamed and "observation.state" not in renamed
# Ensure deep copy: mutate original and verify renamed unaffected
orig["observation.state"]["mean"][0] = 42.0
assert renamed["observation.robot_state"]["mean"][0] != 42.0
+11 -11
View File
@@ -23,7 +23,7 @@ import torch
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.constants import ACTION, OBS_STATE
from lerobot.policies.sac.configuration_sac import SACConfig
from lerobot.policies.sac.processor_sac import make_sac_processor
from lerobot.policies.sac.processor_sac import make_sac_pre_post_processors
from lerobot.processor import (
DeviceProcessor,
NormalizerProcessor,
@@ -78,7 +78,7 @@ def test_make_sac_processor_basic():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_sac_processor(config, stats)
preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
# Check processor names
assert preprocessor.name == "robot_preprocessor"
@@ -102,7 +102,7 @@ def test_sac_processor_normalization_modes():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_sac_processor(config, stats)
preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
# Create test data
observation = {OBS_STATE: torch.randn(10) * 2} # Larger values to test normalization
@@ -133,7 +133,7 @@ def test_sac_processor_cuda():
config.device = "cuda"
stats = create_default_stats()
preprocessor, postprocessor = make_sac_processor(config, stats)
preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
# Create CPU data
observation = {OBS_STATE: torch.randn(10)}
@@ -162,7 +162,7 @@ def test_sac_processor_accelerate_scenario():
config.device = "cuda:0"
stats = create_default_stats()
preprocessor, postprocessor = make_sac_processor(config, stats)
preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
# Simulate Accelerate: data already on GPU
device = torch.device("cuda:0")
@@ -185,7 +185,7 @@ def test_sac_processor_multi_gpu():
config.device = "cuda:0"
stats = create_default_stats()
preprocessor, postprocessor = make_sac_processor(config, stats)
preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
# Simulate data on different GPU
device = torch.device("cuda:1")
@@ -205,7 +205,7 @@ def test_sac_processor_without_stats():
"""Test SAC processor creation without dataset statistics."""
config = create_default_config()
preprocessor, postprocessor = make_sac_processor(config, dataset_stats=None)
preprocessor, postprocessor = make_sac_pre_post_processors(config, dataset_stats=None)
# Should still create processors
assert preprocessor is not None
@@ -225,7 +225,7 @@ def test_sac_processor_save_and_load():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_sac_processor(config, stats)
preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
with tempfile.TemporaryDirectory() as tmpdir:
# Save preprocessor
@@ -252,7 +252,7 @@ def test_sac_processor_mixed_precision():
stats = create_default_stats()
# Create processor
preprocessor, postprocessor = make_sac_processor(config, stats)
preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
# Replace DeviceProcessor with one that uses float16
for i, step in enumerate(preprocessor.steps):
@@ -277,7 +277,7 @@ def test_sac_processor_batch_data():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_sac_processor(config, stats)
preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
# Test with batched data
batch_size = 32
@@ -298,7 +298,7 @@ def test_sac_processor_edge_cases():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_sac_processor(config, stats)
preprocessor, postprocessor = make_sac_pre_post_processors(config, stats)
# Test with empty observation
transition = create_transition(observation={}, action=torch.randn(5))
+9 -6
View File
@@ -23,7 +23,10 @@ import torch
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE
from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
from lerobot.policies.smolvla.processor_smolvla import SmolVLANewLineProcessor, make_smolvla_processor
from lerobot.policies.smolvla.processor_smolvla import (
SmolVLANewLineProcessor,
make_smolvla_pre_post_processors,
)
from lerobot.processor import (
DeviceProcessor,
NormalizerProcessor,
@@ -86,7 +89,7 @@ def test_make_smolvla_processor_basic():
stats = create_default_stats()
with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor"):
preprocessor, postprocessor = make_smolvla_processor(config, stats)
preprocessor, postprocessor = make_smolvla_pre_post_processors(config, stats)
# Check processor names
assert preprocessor.name == "robot_preprocessor"
@@ -185,7 +188,7 @@ def test_smolvla_processor_cuda():
return features
with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor):
preprocessor, postprocessor = make_smolvla_processor(config, stats)
preprocessor, postprocessor = make_smolvla_pre_post_processors(config, stats)
# Create CPU data
observation = {
@@ -235,7 +238,7 @@ def test_smolvla_processor_accelerate_scenario():
return features
with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor):
preprocessor, postprocessor = make_smolvla_processor(config, stats)
preprocessor, postprocessor = make_smolvla_pre_post_processors(config, stats)
# Simulate Accelerate: data already on GPU and batched
device = torch.device("cuda:0")
@@ -286,7 +289,7 @@ def test_smolvla_processor_multi_gpu():
return features
with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor):
preprocessor, postprocessor = make_smolvla_processor(config, stats)
preprocessor, postprocessor = make_smolvla_pre_post_processors(config, stats)
# Simulate data on different GPU
device = torch.device("cuda:1")
@@ -312,7 +315,7 @@ def test_smolvla_processor_without_stats():
# Mock the tokenizer processor
with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor"):
preprocessor, postprocessor = make_smolvla_processor(config, dataset_stats=None)
preprocessor, postprocessor = make_smolvla_pre_post_processors(config, dataset_stats=None)
# Should still create processors
assert preprocessor is not None
+11 -11
View File
@@ -23,7 +23,7 @@ import torch
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE
from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
from lerobot.policies.tdmpc.processor_tdmpc import make_tdmpc_processor
from lerobot.policies.tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
from lerobot.processor import (
DeviceProcessor,
NormalizerProcessor,
@@ -81,7 +81,7 @@ def test_make_tdmpc_processor_basic():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_tdmpc_processor(config, stats)
preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
# Check processor names
assert preprocessor.name == "robot_preprocessor"
@@ -105,7 +105,7 @@ def test_tdmpc_processor_normalization():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_tdmpc_processor(config, stats)
preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
# Create test data
observation = {
@@ -138,7 +138,7 @@ def test_tdmpc_processor_cuda():
config.device = "cuda"
stats = create_default_stats()
preprocessor, postprocessor = make_tdmpc_processor(config, stats)
preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
# Create CPU data
observation = {
@@ -171,7 +171,7 @@ def test_tdmpc_processor_accelerate_scenario():
config.device = "cuda:0"
stats = create_default_stats()
preprocessor, postprocessor = make_tdmpc_processor(config, stats)
preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
# Simulate Accelerate: data already on GPU
device = torch.device("cuda:0")
@@ -198,7 +198,7 @@ def test_tdmpc_processor_multi_gpu():
config.device = "cuda:0"
stats = create_default_stats()
preprocessor, postprocessor = make_tdmpc_processor(config, stats)
preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
# Simulate data on different GPU
device = torch.device("cuda:1")
@@ -222,7 +222,7 @@ def test_tdmpc_processor_without_stats():
"""Test TDMPC processor creation without dataset statistics."""
config = create_default_config()
preprocessor, postprocessor = make_tdmpc_processor(config, dataset_stats=None)
preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, dataset_stats=None)
# Should still create processors
assert preprocessor is not None
@@ -245,7 +245,7 @@ def test_tdmpc_processor_save_and_load():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_tdmpc_processor(config, stats)
preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
with tempfile.TemporaryDirectory() as tmpdir:
# Save preprocessor
@@ -276,7 +276,7 @@ def test_tdmpc_processor_mixed_precision():
stats = create_default_stats()
# Create processor
preprocessor, postprocessor = make_tdmpc_processor(config, stats)
preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
# Replace DeviceProcessor with one that uses float16
for i, step in enumerate(preprocessor.steps):
@@ -305,7 +305,7 @@ def test_tdmpc_processor_batch_data():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_tdmpc_processor(config, stats)
preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
# Test with batched data
batch_size = 64
@@ -330,7 +330,7 @@ def test_tdmpc_processor_edge_cases():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_tdmpc_processor(config, stats)
preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats)
# Test with only state observation (no image)
observation = {OBS_STATE: torch.randn(12)}
+58 -38
View File
@@ -98,7 +98,11 @@ def test_basic_tokenization(mock_auto_tokenizer):
processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=10)
transition = create_transition(complementary_data={"task": "pick up the red cube"})
transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])},
action=torch.tensor([0.1, 0.2]),
complementary_data={"task": "pick up the red cube"},
)
result = processor(transition)
@@ -126,7 +130,11 @@ def test_basic_tokenization_with_tokenizer_object():
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10)
transition = create_transition(complementary_data={"task": "pick up the red cube"})
transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])},
action=torch.tensor([0.1, 0.2]),
complementary_data={"task": "pick up the red cube"},
)
result = processor(transition)
@@ -156,7 +164,11 @@ def test_list_of_strings_tokenization(mock_auto_tokenizer):
processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=8)
transition = create_transition(complementary_data={"task": ["pick up cube", "place on table"]})
transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])},
action=torch.tensor([0.1, 0.2]),
complementary_data={"task": ["pick up cube", "place on table"]},
)
result = processor(transition)
@@ -180,7 +192,11 @@ def test_custom_keys(mock_auto_tokenizer):
processor = TokenizerProcessor(tokenizer_name="test-tokenizer", task_key="instruction", max_length=5)
transition = create_transition(complementary_data={"instruction": "move forward"})
transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])},
action=torch.tensor([0.1, 0.2]),
complementary_data={"instruction": "move forward"},
)
result = processor(transition)
@@ -421,7 +437,11 @@ def test_save_and_load_pretrained_with_tokenizer_name(mock_auto_tokenizer):
loaded_processor = RobotProcessor.from_pretrained(temp_dir)
# Test that loaded processor works
transition = create_transition(complementary_data={"instruction": "test instruction"})
transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])},
action=torch.tensor([0.1, 0.2]),
complementary_data={"instruction": "test instruction"},
)
result = loaded_processor(transition)
assert TransitionKey.OBSERVATION in result
@@ -448,7 +468,11 @@ def test_save_and_load_pretrained_with_tokenizer_object():
)
# Test that loaded processor works
transition = create_transition(complementary_data={"instruction": "test instruction"})
transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])},
action=torch.tensor([0.1, 0.2]),
complementary_data={"instruction": "test instruction"},
)
result = loaded_processor(transition)
assert TransitionKey.OBSERVATION in result
@@ -569,7 +593,11 @@ def test_tokenization_parameters(mock_auto_tokenizer):
padding_side="left",
)
transition = create_transition(complementary_data={"task": "test task"})
transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])},
action=torch.tensor([0.1, 0.2]),
complementary_data={"task": "test task"},
)
processor(transition)
@@ -592,12 +620,14 @@ def test_preserves_other_complementary_data(mock_auto_tokenizer):
processor = TokenizerProcessor(tokenizer_name="test-tokenizer")
transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])},
action=torch.tensor([0.1, 0.2]),
complementary_data={
"task": "test task",
"episode_id": 123,
"timestamp": 456.789,
"other_field": {"nested": "data"},
}
},
)
result = processor(transition)
@@ -624,7 +654,11 @@ def test_deterministic_tokenization(mock_auto_tokenizer):
processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=10)
transition = create_transition(complementary_data={"task": "consistent test"})
transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])},
action=torch.tensor([0.1, 0.2]),
complementary_data={"task": "consistent test"},
)
result1 = processor(transition)
result2 = processor(transition)
@@ -648,7 +682,11 @@ def test_empty_string_task(mock_auto_tokenizer):
processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=8)
transition = create_transition(complementary_data={"task": ""})
transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])},
action=torch.tensor([0.1, 0.2]),
complementary_data={"task": ""},
)
result = processor(transition)
@@ -669,7 +707,11 @@ def test_very_long_task(mock_auto_tokenizer):
processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=5, truncation=True)
long_task = " ".join(["word"] * 100) # Very long task
transition = create_transition(complementary_data={"task": long_task})
transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])},
action=torch.tensor([0.1, 0.2]),
complementary_data={"task": long_task},
)
result = processor(transition)
@@ -714,7 +756,11 @@ def test_custom_padding_side(mock_auto_tokenizer):
# Test left padding
processor_left = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=10, padding_side="left")
transition = create_transition(complementary_data={"task": "test task"})
transition = create_transition(
observation={"state": torch.tensor([1.0, 2.0])},
action=torch.tensor([0.1, 0.2]),
complementary_data={"task": "test task"},
)
processor_left(transition)
assert tracking_tokenizer.padding_side_calls[-1] == "left"
@@ -873,32 +919,6 @@ def test_device_detection_from_action():
assert attention_mask.device.type == "cuda"
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
@require_package("transformers")
def test_device_detection_from_complementary_data():
"""Test that device is detected from tensors in complementary_data."""
mock_tokenizer = MockTokenizer(vocab_size=100)
processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10)
# Create transition with tensor in complementary_data
transition = create_transition(
observation={"metadata": {"key": "value"}}, # No tensors
complementary_data={
"task": "comp data test",
"index": torch.tensor([42]).cuda(), # Tensor in complementary_data
},
)
result = processor(transition)
# Check that tokenized tensors match complementary_data tensor's device
tokens = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
attention_mask = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.attention_mask"]
assert tokens.device.type == "cuda"
assert attention_mask.device.type == "cuda"
@require_package("transformers")
def test_device_detection_preserves_dtype():
"""Test that device detection doesn't affect dtype of tokenized tensors."""
+11 -11
View File
@@ -23,7 +23,7 @@ import torch
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE
from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
from lerobot.policies.vqbet.processor_vqbet import make_vqbet_processor
from lerobot.policies.vqbet.processor_vqbet import make_vqbet_pre_post_processors
from lerobot.processor import (
DeviceProcessor,
NormalizerProcessor,
@@ -81,7 +81,7 @@ def test_make_vqbet_processor_basic():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_vqbet_processor(config, stats)
preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
# Check processor names
assert preprocessor.name == "robot_preprocessor"
@@ -105,7 +105,7 @@ def test_vqbet_processor_with_images():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_vqbet_processor(config, stats)
preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
# Create test data with images and states
observation = {
@@ -131,7 +131,7 @@ def test_vqbet_processor_cuda():
config.device = "cuda"
stats = create_default_stats()
preprocessor, postprocessor = make_vqbet_processor(config, stats)
preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
# Create CPU data
observation = {
@@ -164,7 +164,7 @@ def test_vqbet_processor_accelerate_scenario():
config.device = "cuda:0"
stats = create_default_stats()
preprocessor, postprocessor = make_vqbet_processor(config, stats)
preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
# Simulate Accelerate: data already on GPU and batched
device = torch.device("cuda:0")
@@ -191,7 +191,7 @@ def test_vqbet_processor_multi_gpu():
config.device = "cuda:0"
stats = create_default_stats()
preprocessor, postprocessor = make_vqbet_processor(config, stats)
preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
# Simulate data on different GPU
device = torch.device("cuda:1")
@@ -215,7 +215,7 @@ def test_vqbet_processor_without_stats():
"""Test VQBeT processor creation without dataset statistics."""
config = create_default_config()
preprocessor, postprocessor = make_vqbet_processor(config, dataset_stats=None)
preprocessor, postprocessor = make_vqbet_pre_post_processors(config, dataset_stats=None)
# Should still create processors
assert preprocessor is not None
@@ -238,7 +238,7 @@ def test_vqbet_processor_save_and_load():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_vqbet_processor(config, stats)
preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
with tempfile.TemporaryDirectory() as tmpdir:
# Save preprocessor
@@ -269,7 +269,7 @@ def test_vqbet_processor_mixed_precision():
stats = create_default_stats()
# Create processor
preprocessor, postprocessor = make_vqbet_processor(config, stats)
preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
# Replace DeviceProcessor with one that uses float16
for i, step in enumerate(preprocessor.steps):
@@ -298,7 +298,7 @@ def test_vqbet_processor_large_batch():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_vqbet_processor(config, stats)
preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
# Test with large batch
batch_size = 128
@@ -323,7 +323,7 @@ def test_vqbet_processor_sequential_processing():
config = create_default_config()
stats = create_default_stats()
preprocessor, postprocessor = make_vqbet_processor(config, stats)
preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats)
# Process multiple samples sequentially
results = []