tangger/lerobot
1
0
Fork 1
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

chore: replace hard-coded action values with constants throughout all the source code (#2055)

Browse Source 
* chore: replace hard-coded 'action' values with constants throughout all the source code

* chore(tests): replace hard-coded action values with constants throughout all the test code
This commit is contained in:
Steven Palma
2025-09-26 13:33:18 +02:00
committed by GitHub
parent 9627765ce2
commit d2782cf66b
47 changed files with 269 additions and 255 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
tests/processor/test_batch_conversion.py
View File

@@ -2,7 +2,7 @@ import torch
from lerobot.processor import DataProcessorPipeline, TransitionKey
from lerobot.processor.converters import batch_to_transition, transition_to_batch
from lerobot.utils.constants import OBS_IMAGE, OBS_PREFIX, OBS_STATE
from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_PREFIX, OBS_STATE
def _dummy_batch():
@@ -11,7 +11,7 @@ def _dummy_batch():
f"{OBS_IMAGE}.left": torch.randn(1, 3, 128, 128),
f"{OBS_IMAGE}.right": torch.randn(1, 3, 128, 128),
OBS_STATE: torch.tensor([[0.1, 0.2, 0.3, 0.4]]),
"action": torch.tensor([[0.5]]),
ACTION: torch.tensor([[0.5]]),
"next.reward": 1.0,
"next.done": False,
"next.truncated": False,
@@ -37,7 +37,7 @@ def test_observation_grouping_roundtrip():
assert torch.allclose(batch_out[OBS_STATE], batch_in[OBS_STATE])
# Check other fields
assert torch.allclose(batch_out["action"], batch_in["action"])
assert torch.allclose(batch_out[ACTION], batch_in[ACTION])
assert batch_out["next.reward"] == batch_in["next.reward"]
assert batch_out["next.done"] == batch_in["next.done"]
assert batch_out["next.truncated"] == batch_in["next.truncated"]
@@ -50,7 +50,7 @@ def test_batch_to_transition_observation_grouping():
f"{OBS_IMAGE}.top": torch.randn(1, 3, 128, 128),
f"{OBS_IMAGE}.left": torch.randn(1, 3, 128, 128),
OBS_STATE: [1, 2, 3, 4],
"action": torch.tensor([0.1, 0.2, 0.3, 0.4]),
ACTION: torch.tensor([0.1, 0.2, 0.3, 0.4]),
"next.reward": 1.5,
"next.done": True,
"next.truncated": False,
@@ -114,7 +114,7 @@ def test_transition_to_batch_observation_flattening():
assert batch[OBS_STATE] == [1, 2, 3, 4]
# Check other fields are mapped to next.* format
assert batch["action"] == "action_data"
assert batch[ACTION] == "action_data"
assert batch["next.reward"] == 1.5
assert batch["next.done"]
assert not batch["next.truncated"]
@@ -124,7 +124,7 @@ def test_transition_to_batch_observation_flattening():
def test_no_observation_keys():
"""Test behavior when there are no observation.* keys."""
batch = {
"action": torch.tensor([1.0, 2.0]),
ACTION: torch.tensor([1.0, 2.0]),
"next.reward": 2.0,
"next.done": False,
"next.truncated": True,
@@ -145,7 +145,7 @@ def test_no_observation_keys():
# Round trip should work
reconstructed_batch = transition_to_batch(transition)
assert torch.allclose(reconstructed_batch["action"], torch.tensor([1.0, 2.0]))
assert torch.allclose(reconstructed_batch[ACTION], torch.tensor([1.0, 2.0]))
assert reconstructed_batch["next.reward"] == 2.0
assert not reconstructed_batch["next.done"]
assert reconstructed_batch["next.truncated"]
@@ -154,7 +154,7 @@ def test_no_observation_keys():
def test_minimal_batch():
"""Test with minimal batch containing only observation.* and action."""
batch = {OBS_STATE: "minimal_state", "action": torch.tensor([0.5])}
batch = {OBS_STATE: "minimal_state", ACTION: torch.tensor([0.5])}
transition = batch_to_transition(batch)
@@ -172,7 +172,7 @@ def test_minimal_batch():
# Round trip
reconstructed_batch = transition_to_batch(transition)
assert reconstructed_batch[OBS_STATE] == "minimal_state"
assert torch.allclose(reconstructed_batch["action"], torch.tensor([0.5]))
assert torch.allclose(reconstructed_batch[ACTION], torch.tensor([0.5]))
assert reconstructed_batch["next.reward"] == 0.0
assert not reconstructed_batch["next.done"]
assert not reconstructed_batch["next.truncated"]
@@ -196,7 +196,7 @@ def test_empty_batch():
# Round trip
reconstructed_batch = transition_to_batch(transition)
assert reconstructed_batch["action"] is None
assert reconstructed_batch[ACTION] is None
assert reconstructed_batch["next.reward"] == 0.0
assert not reconstructed_batch["next.done"]
assert not reconstructed_batch["next.truncated"]
@@ -209,7 +209,7 @@ def test_complex_nested_observation():
f"{OBS_IMAGE}.top": {"image": torch.randn(1, 3, 128, 128), "timestamp": 1234567890},
f"{OBS_IMAGE}.left": {"image": torch.randn(1, 3, 128, 128), "timestamp": 1234567891},
OBS_STATE: torch.randn(7),
"action": torch.randn(8),
ACTION: torch.randn(8),
"next.reward": 3.14,
"next.done": False,
"next.truncated": True,
@@ -237,7 +237,7 @@ def test_complex_nested_observation():
)
# Check action tensor
assert torch.allclose(batch["action"], reconstructed_batch["action"])
assert torch.allclose(batch[ACTION], reconstructed_batch[ACTION])
# Check other fields
assert batch["next.reward"] == reconstructed_batch["next.reward"]
@@ -266,7 +266,7 @@ def test_custom_converter():
batch = {
OBS_STATE: torch.randn(1, 4),
"action": torch.randn(1, 2),
ACTION: torch.randn(1, 2),
"next.reward": 1.0,
"next.done": False,
}
@@ -276,4 +276,4 @@ def test_custom_converter():
# Check the reward was doubled by our custom converter
assert result["next.reward"] == 2.0
assert torch.allclose(result[OBS_STATE], batch[OBS_STATE])
assert torch.allclose(result["action"], batch["action"])
assert torch.allclose(result[ACTION], batch[ACTION])

14
tests/processor/test_converters.py
View File

@@ -9,7 +9,7 @@ from lerobot.processor.converters import (
to_tensor,
transition_to_batch,
)
from lerobot.utils.constants import OBS_STATE, OBS_STR
from lerobot.utils.constants import ACTION, OBS_STATE, OBS_STR
# Tests for the unified to_tensor function
@@ -118,16 +118,16 @@ def test_to_tensor_dictionaries():
# Nested dictionary
nested = {
"action": {"mean": [0.1, 0.2], "std": [1.0, 2.0]},
ACTION: {"mean": [0.1, 0.2], "std": [1.0, 2.0]},
OBS_STR: {"mean": np.array([0.5, 0.6]), "count": 10},
}
result = to_tensor(nested)
assert isinstance(result, dict)
assert isinstance(result["action"], dict)
assert isinstance(result[ACTION], dict)
assert isinstance(result[OBS_STR], dict)
assert isinstance(result["action"]["mean"], torch.Tensor)
assert isinstance(result[ACTION]["mean"], torch.Tensor)
assert isinstance(result[OBS_STR]["mean"], torch.Tensor)
assert torch.allclose(result["action"]["mean"], torch.tensor([0.1, 0.2]))
assert torch.allclose(result[ACTION]["mean"], torch.tensor([0.1, 0.2]))
assert torch.allclose(result[OBS_STR]["mean"], torch.tensor([0.5, 0.6]))
@@ -200,7 +200,7 @@ def test_batch_to_transition_with_index_fields():
# Create batch with index and task_index fields
batch = {
OBS_STATE: torch.randn(1, 7),
"action": torch.randn(1, 4),
ACTION: torch.randn(1, 4),
"next.reward": 1.5,
"next.done": False,
"task": ["pick_cube"],
@@ -262,7 +262,7 @@ def test_batch_to_transition_without_index_fields():
# Batch without index/task_index
batch = {
OBS_STATE: torch.randn(1, 7),
"action": torch.randn(1, 4),
ACTION: torch.randn(1, 4),
"task": ["pick_cube"],
}

4
tests/processor/test_device_processor.py
View File

@@ -21,7 +21,7 @@ import torch
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.processor import DataProcessorPipeline, DeviceProcessorStep, TransitionKey
from lerobot.processor.converters import create_transition, identity_transition
from lerobot.utils.constants import OBS_IMAGE, OBS_STATE
from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE
def test_basic_functionality():
@@ -273,7 +273,7 @@ def test_features():
features = {
PipelineFeatureType.OBSERVATION: {OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(10,))},
PipelineFeatureType.ACTION: {"action": PolicyFeature(type=FeatureType.ACTION, shape=(5,))},
PipelineFeatureType.ACTION: {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(5,))},
}
result = processor.transform_features(features)

4
tests/processor/test_migration_detection.py
View File

@@ -25,7 +25,7 @@ from pathlib import Path
import pytest
from lerobot.processor.pipeline import DataProcessorPipeline, ProcessorMigrationError
from lerobot.utils.constants import OBS_STATE
from lerobot.utils.constants import ACTION, OBS_STATE
def test_is_processor_config_valid_configs():
@@ -113,7 +113,7 @@ def test_should_suggest_migration_with_model_config_only():
model_config = {
"type": "act",
"input_features": {OBS_STATE: {"shape": [7]}},
"output_features": {"action": {"shape": [7]}},
"output_features": {ACTION: {"shape": [7]}},
"hidden_dim": 256,
"n_obs_steps": 1,
"n_action_steps": 1,

140
tests/processor/test_normalize_processor.py
View File

@@ -29,7 +29,7 @@ from lerobot.processor import (
hotswap_stats,
)
from lerobot.processor.converters import create_transition, identity_transition, to_tensor
from lerobot.utils.constants import OBS_IMAGE, OBS_STATE, OBS_STR
from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE, OBS_STR
from lerobot.utils.utils import auto_select_torch_device
@@ -50,15 +50,15 @@ def test_numpy_conversion():
def test_tensor_conversion():
stats = {
"action": {
ACTION: {
"mean": torch.tensor([0.0, 0.0]),
"std": torch.tensor([1.0, 1.0]),
}
}
tensor_stats = to_tensor(stats)
assert tensor_stats["action"]["mean"].dtype == torch.float32
assert tensor_stats["action"]["std"].dtype == torch.float32
assert tensor_stats[ACTION]["mean"].dtype == torch.float32
assert tensor_stats[ACTION]["std"].dtype == torch.float32
def test_scalar_conversion():
@@ -212,12 +212,12 @@ def test_from_lerobot_dataset():
mock_dataset = Mock()
mock_dataset.meta.stats = {
OBS_IMAGE: {"mean": [0.5], "std": [0.2]},
"action": {"mean": [0.0], "std": [1.0]},
ACTION: {"mean": [0.0], "std": [1.0]},
}
features = {
OBS_IMAGE: PolicyFeature(FeatureType.VISUAL, (3, 96, 96)),
"action": PolicyFeature(FeatureType.ACTION, (1,)),
ACTION: PolicyFeature(FeatureType.ACTION, (1,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
@@ -228,7 +228,7 @@ def test_from_lerobot_dataset():
# Both observation and action statistics should be present in tensor stats
assert OBS_IMAGE in normalizer._tensor_stats
assert "action" in normalizer._tensor_stats
assert ACTION in normalizer._tensor_stats
def test_state_dict_save_load(observation_normalizer):
@@ -271,7 +271,7 @@ def action_stats_min_max():
def _create_action_features():
return {
"action": PolicyFeature(FeatureType.ACTION, (3,)),
ACTION: PolicyFeature(FeatureType.ACTION, (3,)),
}
@@ -291,7 +291,7 @@ def test_mean_std_unnormalization(action_stats_mean_std):
features = _create_action_features()
norm_map = _create_action_norm_map_mean_std()
unnormalizer = UnnormalizerProcessorStep(
features=features, norm_map=norm_map, stats={"action": action_stats_mean_std}
features=features, norm_map=norm_map, stats={ACTION: action_stats_mean_std}
)
normalized_action = torch.tensor([1.0, -0.5, 2.0])
@@ -309,7 +309,7 @@ def test_min_max_unnormalization(action_stats_min_max):
features = _create_action_features()
norm_map = _create_action_norm_map_min_max()
unnormalizer = UnnormalizerProcessorStep(
features=features, norm_map=norm_map, stats={"action": action_stats_min_max}
features=features, norm_map=norm_map, stats={ACTION: action_stats_min_max}
)
# Actions in [-1, 1]
@@ -335,7 +335,7 @@ def test_tensor_action_input(action_stats_mean_std):
features = _create_action_features()
norm_map = _create_action_norm_map_mean_std()
unnormalizer = UnnormalizerProcessorStep(
features=features, norm_map=norm_map, stats={"action": action_stats_mean_std}
features=features, norm_map=norm_map, stats={ACTION: action_stats_mean_std}
)
normalized_action = torch.tensor([1.0, -0.5, 2.0], dtype=torch.float32)
@@ -353,7 +353,7 @@ def test_none_action(action_stats_mean_std):
features = _create_action_features()
norm_map = _create_action_norm_map_mean_std()
unnormalizer = UnnormalizerProcessorStep(
features=features, norm_map=norm_map, stats={"action": action_stats_mean_std}
features=features, norm_map=norm_map, stats={ACTION: action_stats_mean_std}
)
transition = create_transition()
@@ -365,11 +365,11 @@ def test_none_action(action_stats_mean_std):
def test_action_from_lerobot_dataset():
mock_dataset = Mock()
mock_dataset.meta.stats = {"action": {"mean": [0.0], "std": [1.0]}}
features = {"action": PolicyFeature(FeatureType.ACTION, (1,))}
mock_dataset.meta.stats = {ACTION: {"mean": [0.0], "std": [1.0]}}
features = {ACTION: PolicyFeature(FeatureType.ACTION, (1,))}
norm_map = {FeatureType.ACTION: NormalizationMode.MEAN_STD}
unnormalizer = UnnormalizerProcessorStep.from_lerobot_dataset(mock_dataset, features, norm_map)
assert "mean" in unnormalizer._tensor_stats["action"]
assert "mean" in unnormalizer._tensor_stats[ACTION]
# Fixtures for NormalizerProcessorStep tests
@@ -384,7 +384,7 @@ def full_stats():
"min": np.array([0.0, -1.0]),
"max": np.array([1.0, 1.0]),
},
"action": {
ACTION: {
"mean": np.array([0.0, 0.0]),
"std": np.array([1.0, 2.0]),
},
@@ -395,7 +395,7 @@ def _create_full_features():
return {
OBS_IMAGE: PolicyFeature(FeatureType.VISUAL, (3, 96, 96)),
OBS_STATE: PolicyFeature(FeatureType.STATE, (2,)),
"action": PolicyFeature(FeatureType.ACTION, (2,)),
ACTION: PolicyFeature(FeatureType.ACTION, (2,)),
}
@@ -461,7 +461,7 @@ def test_processor_from_lerobot_dataset(full_stats):
assert processor.normalize_observation_keys == {OBS_IMAGE}
assert OBS_IMAGE in processor._tensor_stats
assert "action" in processor._tensor_stats
assert ACTION in processor._tensor_stats
def test_get_config(full_stats):
@@ -482,7 +482,7 @@ def test_get_config(full_stats):
"features": {
OBS_IMAGE: {"type": "VISUAL", "shape": (3, 96, 96)},
OBS_STATE: {"type": "STATE", "shape": (2,)},
"action": {"type": "ACTION", "shape": (2,)},
ACTION: {"type": "ACTION", "shape": (2,)},
},
"norm_map": {
"VISUAL": "MEAN_STD",
@@ -568,7 +568,7 @@ def test_missing_action_stats_no_error():
processor = UnnormalizerProcessorStep.from_lerobot_dataset(mock_dataset, features, norm_map)
# The tensor stats should not contain the 'action' key
assert "action" not in processor._tensor_stats
assert ACTION not in processor._tensor_stats
def test_serialization_roundtrip(full_stats):
@@ -676,9 +676,9 @@ def test_identity_normalization_observations():
def test_identity_normalization_actions():
"""Test that IDENTITY mode skips normalization for actions."""
features = {"action": PolicyFeature(FeatureType.ACTION, (2,))}
features = {ACTION: PolicyFeature(FeatureType.ACTION, (2,))}
norm_map = {FeatureType.ACTION: NormalizationMode.IDENTITY}
stats = {"action": {"mean": [0.0, 0.0], "std": [1.0, 2.0]}}
stats = {ACTION: {"mean": [0.0, 0.0], "std": [1.0, 2.0]}}
normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
@@ -729,9 +729,9 @@ def test_identity_unnormalization_observations():
def test_identity_unnormalization_actions():
"""Test that IDENTITY mode skips unnormalization for actions."""
features = {"action": PolicyFeature(FeatureType.ACTION, (2,))}
features = {ACTION: PolicyFeature(FeatureType.ACTION, (2,))}
norm_map = {FeatureType.ACTION: NormalizationMode.IDENTITY}
stats = {"action": {"min": [-1.0, -2.0], "max": [1.0, 2.0]}}
stats = {ACTION: {"min": [-1.0, -2.0], "max": [1.0, 2.0]}}
unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
@@ -748,7 +748,7 @@ def test_identity_with_missing_stats():
"""Test that IDENTITY mode works even when stats are missing."""
features = {
OBS_IMAGE: PolicyFeature(FeatureType.VISUAL, (3, 96, 96)),
"action": PolicyFeature(FeatureType.ACTION, (2,)),
ACTION: PolicyFeature(FeatureType.ACTION, (2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.IDENTITY,
@@ -784,7 +784,7 @@ def test_identity_mixed_with_other_modes():
features = {
OBS_IMAGE: PolicyFeature(FeatureType.VISUAL, (3,)),
OBS_STATE: PolicyFeature(FeatureType.STATE, (2,)),
"action": PolicyFeature(FeatureType.ACTION, (2,)),
ACTION: PolicyFeature(FeatureType.ACTION, (2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.IDENTITY,
@@ -794,7 +794,7 @@ def test_identity_mixed_with_other_modes():
stats = {
OBS_IMAGE: {"mean": [0.5, 0.5, 0.5], "std": [0.2, 0.2, 0.2]}, # Will be ignored
OBS_STATE: {"mean": [0.0, 0.0], "std": [1.0, 1.0]},
"action": {"min": [-1.0, -1.0], "max": [1.0, 1.0]},
ACTION: {"min": [-1.0, -1.0], "max": [1.0, 1.0]},
}
normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
@@ -862,7 +862,7 @@ def test_identity_roundtrip():
"""Test that IDENTITY normalization and unnormalization are true inverses."""
features = {
OBS_IMAGE: PolicyFeature(FeatureType.VISUAL, (3,)),
"action": PolicyFeature(FeatureType.ACTION, (2,)),
ACTION: PolicyFeature(FeatureType.ACTION, (2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.IDENTITY,
@@ -870,7 +870,7 @@ def test_identity_roundtrip():
}
stats = {
OBS_IMAGE: {"mean": [0.5, 0.5, 0.5], "std": [0.2, 0.2, 0.2]},
"action": {"min": [-1.0, -1.0], "max": [1.0, 1.0]},
ACTION: {"min": [-1.0, -1.0], "max": [1.0, 1.0]},
}
normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
@@ -893,7 +893,7 @@ def test_identity_config_serialization():
"""Test that IDENTITY mode is properly saved and loaded in config."""
features = {
OBS_IMAGE: PolicyFeature(FeatureType.VISUAL, (3,)),
"action": PolicyFeature(FeatureType.ACTION, (2,)),
ACTION: PolicyFeature(FeatureType.ACTION, (2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.IDENTITY,
@@ -901,7 +901,7 @@ def test_identity_config_serialization():
}
stats = {
OBS_IMAGE: {"mean": [0.5], "std": [0.2]},
"action": {"mean": [0.0, 0.0], "std": [1.0, 1.0]},
ACTION: {"mean": [0.0, 0.0], "std": [1.0, 1.0]},
}
normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
@@ -969,19 +969,19 @@ def test_hotswap_stats_basic_functionality():
# Create initial stats
initial_stats = {
OBS_IMAGE: {"mean": np.array([0.5, 0.5, 0.5]), "std": np.array([0.2, 0.2, 0.2])},
"action": {"mean": np.array([0.0, 0.0]), "std": np.array([1.0, 1.0])},
ACTION: {"mean": np.array([0.0, 0.0]), "std": np.array([1.0, 1.0])},
}
# Create new stats for hotswapping
new_stats = {
OBS_IMAGE: {"mean": np.array([0.3, 0.3, 0.3]), "std": np.array([0.1, 0.1, 0.1])},
"action": {"mean": np.array([0.1, 0.1]), "std": np.array([0.5, 0.5])},
ACTION: {"mean": np.array([0.1, 0.1]), "std": np.array([0.5, 0.5])},
}
# Create features and norm_map
features = {
OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128)),
"action": PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
@@ -1177,17 +1177,17 @@ def test_hotswap_stats_multiple_normalizer_types():
"""Test hotswap_stats with multiple normalizer and unnormalizer steps."""
initial_stats = {
OBS_IMAGE: {"mean": np.array([0.5]), "std": np.array([0.2])},
"action": {"min": np.array([-1.0]), "max": np.array([1.0])},
ACTION: {"min": np.array([-1.0]), "max": np.array([1.0])},
}
new_stats = {
OBS_IMAGE: {"mean": np.array([0.3]), "std": np.array([0.1])},
"action": {"min": np.array([-2.0]), "max": np.array([2.0])},
ACTION: {"min": np.array([-2.0]), "max": np.array([2.0])},
}
features = {
OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128)),
"action": PolicyFeature(type=FeatureType.ACTION, shape=(1,)),
ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(1,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
@@ -1232,7 +1232,7 @@ def test_hotswap_stats_with_different_data_types():
"min": 0, # int
"max": 1.0, # float
},
"action": {
ACTION: {
"mean": np.array([0.1, 0.2]), # numpy array
"std": torch.tensor([0.5, 0.6]), # torch tensor
},
@@ -1240,7 +1240,7 @@ def test_hotswap_stats_with_different_data_types():
features = {
OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128)),
"action": PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
@@ -1262,8 +1262,8 @@ def test_hotswap_stats_with_different_data_types():
assert isinstance(tensor_stats[OBS_IMAGE]["std"], torch.Tensor)
assert isinstance(tensor_stats[OBS_IMAGE]["min"], torch.Tensor)
assert isinstance(tensor_stats[OBS_IMAGE]["max"], torch.Tensor)
assert isinstance(tensor_stats["action"]["mean"], torch.Tensor)
assert isinstance(tensor_stats["action"]["std"], torch.Tensor)
assert isinstance(tensor_stats[ACTION]["mean"], torch.Tensor)
assert isinstance(tensor_stats[ACTION]["std"], torch.Tensor)
# Check values
torch.testing.assert_close(tensor_stats[OBS_IMAGE]["mean"], torch.tensor([0.3, 0.4, 0.5]))
@@ -1284,18 +1284,18 @@ def test_hotswap_stats_functional_test():
# Initial stats
initial_stats = {
OBS_IMAGE: {"mean": np.array([0.5, 0.4]), "std": np.array([0.2, 0.3])},
"action": {"mean": np.array([0.0, 0.0]), "std": np.array([1.0, 1.0])},
ACTION: {"mean": np.array([0.0, 0.0]), "std": np.array([1.0, 1.0])},
}
# New stats
new_stats = {
OBS_IMAGE: {"mean": np.array([0.3, 0.2]), "std": np.array([0.1, 0.2])},
"action": {"mean": np.array([0.1, -0.1]), "std": np.array([0.5, 0.5])},
ACTION: {"mean": np.array([0.1, -0.1]), "std": np.array([0.5, 0.5])},
}
features = {
OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(2, 2, 2)),
"action": PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
@@ -1324,18 +1324,18 @@ def test_hotswap_stats_functional_test():
rtol=1e-3,
atol=1e-3,
)
assert not torch.allclose(original_result["action"], new_result["action"], rtol=1e-3, atol=1e-3)
assert not torch.allclose(original_result[ACTION], new_result[ACTION], rtol=1e-3, atol=1e-3)
# Verify that the new processor is actually using the new stats by checking internal state
assert new_processor.steps[0].stats == new_stats
assert torch.allclose(new_processor.steps[0]._tensor_stats[OBS_IMAGE]["mean"], torch.tensor([0.3, 0.2]))
assert torch.allclose(new_processor.steps[0]._tensor_stats[OBS_IMAGE]["std"], torch.tensor([0.1, 0.2]))
assert torch.allclose(new_processor.steps[0]._tensor_stats["action"]["mean"], torch.tensor([0.1, -0.1]))
assert torch.allclose(new_processor.steps[0]._tensor_stats["action"]["std"], torch.tensor([0.5, 0.5]))
assert torch.allclose(new_processor.steps[0]._tensor_stats[ACTION]["mean"], torch.tensor([0.1, -0.1]))
assert torch.allclose(new_processor.steps[0]._tensor_stats[ACTION]["std"], torch.tensor([0.5, 0.5]))
# Test that normalization actually happens (output should not equal input)
assert not torch.allclose(new_result[OBS_STR][OBS_IMAGE], observation[OBS_IMAGE])
assert not torch.allclose(new_result["action"], action)
assert not torch.allclose(new_result[ACTION], action)
def test_zero_std_uses_eps():
@@ -1366,10 +1366,10 @@ def test_action_normalized_despite_normalize_observation_keys():
"""Action normalization is independent of normalize_observation_keys filter for observations."""
features = {
OBS_STATE: PolicyFeature(FeatureType.STATE, (1,)),
"action": PolicyFeature(FeatureType.ACTION, (2,)),
ACTION: PolicyFeature(FeatureType.ACTION, (2,)),
}
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])}}
stats = {ACTION: {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}}
normalizer = NormalizerProcessorStep(
features=features, norm_map=norm_map, stats=stats, normalize_observation_keys={OBS_STATE}
)
@@ -1426,9 +1426,9 @@ def test_unknown_observation_keys_ignored():
def test_batched_action_normalization():
features = {"action": PolicyFeature(FeatureType.ACTION, (2,))}
features = {ACTION: PolicyFeature(FeatureType.ACTION, (2,))}
norm_map = {FeatureType.ACTION: NormalizationMode.MEAN_STD}
stats = {"action": {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}}
stats = {ACTION: {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}}
normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
actions = torch.tensor([[1.0, -1.0], [3.0, 3.0]]) # first equals mean → zeros; second → [1, 1]
@@ -1453,12 +1453,12 @@ def test_complementary_data_preservation():
def test_roundtrip_normalize_unnormalize_non_identity():
features = {
OBS_STATE: PolicyFeature(FeatureType.STATE, (2,)),
"action": PolicyFeature(FeatureType.ACTION, (2,)),
ACTION: PolicyFeature(FeatureType.ACTION, (2,)),
}
norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD, FeatureType.ACTION: NormalizationMode.MIN_MAX}
stats = {
OBS_STATE: {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])},
"action": {"min": np.array([-2.0, 0.0]), "max": np.array([2.0, 4.0])},
ACTION: {"min": np.array([-2.0, 0.0]), "max": np.array([2.0, 4.0])},
}
normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
@@ -1530,18 +1530,18 @@ def test_stats_override_preservation_in_load_state_dict():
# Create original stats
original_stats = {
OBS_IMAGE: {"mean": np.array([0.5, 0.5, 0.5]), "std": np.array([0.2, 0.2, 0.2])},
"action": {"mean": np.array([0.0, 0.0]), "std": np.array([1.0, 1.0])},
ACTION: {"mean": np.array([0.0, 0.0]), "std": np.array([1.0, 1.0])},
}
# Create override stats (what user wants to use)
override_stats = {
OBS_IMAGE: {"mean": np.array([0.3, 0.3, 0.3]), "std": np.array([0.1, 0.1, 0.1])},
"action": {"mean": np.array([0.1, 0.1]), "std": np.array([0.5, 0.5])},
ACTION: {"mean": np.array([0.1, 0.1]), "std": np.array([0.5, 0.5])},
}
features = {
OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128)),
"action": PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
@@ -1601,12 +1601,12 @@ def test_stats_without_override_loads_normally():
"""
original_stats = {
OBS_IMAGE: {"mean": np.array([0.5, 0.5, 0.5]), "std": np.array([0.2, 0.2, 0.2])},
"action": {"mean": np.array([0.0, 0.0]), "std": np.array([1.0, 1.0])},
ACTION: {"mean": np.array([0.0, 0.0]), "std": np.array([1.0, 1.0])},
}
features = {
OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128)),
"action": PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
@@ -1674,7 +1674,7 @@ def test_pipeline_from_pretrained_with_stats_overrides():
# Create test data
features = {
OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 32, 32)),
"action": PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
@@ -1683,12 +1683,12 @@ def test_pipeline_from_pretrained_with_stats_overrides():
original_stats = {
OBS_IMAGE: {"mean": np.array([0.5, 0.5, 0.5]), "std": np.array([0.2, 0.2, 0.2])},
"action": {"mean": np.array([0.0, 0.0]), "std": np.array([1.0, 1.0])},
ACTION: {"mean": np.array([0.0, 0.0]), "std": np.array([1.0, 1.0])},
}
override_stats = {
OBS_IMAGE: {"mean": np.array([0.3, 0.3, 0.3]), "std": np.array([0.1, 0.1, 0.1])},
"action": {"mean": np.array([0.1, 0.1]), "std": np.array([0.5, 0.5])},
ACTION: {"mean": np.array([0.1, 0.1]), "std": np.array([0.5, 0.5])},
}
# Create and save a pipeline with the original stats
@@ -1751,8 +1751,8 @@ def test_pipeline_from_pretrained_with_stats_overrides():
# The critical part was verified above: loaded_normalizer.stats == override_stats
# This confirms that override stats are preserved during load_state_dict.
# Let's just verify the pipeline processes data successfully.
assert "action" in override_result
assert isinstance(override_result["action"], torch.Tensor)
assert ACTION in override_result
assert isinstance(override_result[ACTION], torch.Tensor)
def test_dtype_adaptation_device_processor_bfloat16_normalizer_float32():
@@ -1812,7 +1812,7 @@ def test_stats_reconstruction_after_load_state_dict():
features = {
OBS_IMAGE: PolicyFeature(FeatureType.VISUAL, (3, 96, 96)),
OBS_STATE: PolicyFeature(FeatureType.STATE, (2,)),
"action": PolicyFeature(FeatureType.ACTION, (2,)),
ACTION: PolicyFeature(FeatureType.ACTION, (2,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
@@ -1828,7 +1828,7 @@ def test_stats_reconstruction_after_load_state_dict():
"min": np.array([0.0, -1.0]),
"max": np.array([1.0, 1.0]),
},
"action": {
ACTION: {
"mean": np.array([0.0, 0.0]),
"std": np.array([1.0, 2.0]),
},
@@ -1852,15 +1852,15 @@ def test_stats_reconstruction_after_load_state_dict():
# Check that all expected keys are present
assert OBS_IMAGE in new_normalizer.stats
assert OBS_STATE in new_normalizer.stats
assert "action" in new_normalizer.stats
assert ACTION in new_normalizer.stats
# Check that values are correct (converted back from tensors)
np.testing.assert_allclose(new_normalizer.stats[OBS_IMAGE]["mean"], [0.5, 0.5, 0.5])
np.testing.assert_allclose(new_normalizer.stats[OBS_IMAGE]["std"], [0.2, 0.2, 0.2])
np.testing.assert_allclose(new_normalizer.stats[OBS_STATE]["min"], [0.0, -1.0])
np.testing.assert_allclose(new_normalizer.stats[OBS_STATE]["max"], [1.0, 1.0])
np.testing.assert_allclose(new_normalizer.stats["action"]["mean"], [0.0, 0.0])
np.testing.assert_allclose(new_normalizer.stats["action"]["std"], [1.0, 2.0])
np.testing.assert_allclose(new_normalizer.stats[ACTION]["mean"], [0.0, 0.0])
np.testing.assert_allclose(new_normalizer.stats[ACTION]["std"], [1.0, 2.0])
# Test that methods that depend on self.stats work correctly after loading
# This would fail before the bug fix because self.stats was empty
@@ -1876,7 +1876,7 @@ def test_stats_reconstruction_after_load_state_dict():
new_stats = {
OBS_IMAGE: {"mean": [0.3, 0.3, 0.3], "std": [0.1, 0.1, 0.1]},
OBS_STATE: {"min": [-1.0, -2.0], "max": [2.0, 2.0]},
"action": {"mean": [0.1, 0.1], "std": [0.5, 0.5]},
ACTION: {"mean": [0.1, 0.1], "std": [0.5, 0.5]},
}
pipeline = DataProcessorPipeline([new_normalizer])

28
tests/processor/test_pipeline.py
View File

@@ -35,7 +35,7 @@ from lerobot.processor import (
TransitionKey,
)
from lerobot.processor.converters import create_transition, identity_transition
from lerobot.utils.constants import OBS_IMAGE, OBS_IMAGES, OBS_STATE
from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_IMAGES, OBS_STATE
from tests.conftest import assert_contract_is_typed
@@ -257,7 +257,7 @@ def test_step_through_with_dict():
batch = {
OBS_IMAGE: None,
"action": None,
ACTION: None,
"next.reward": 0.0,
"next.done": False,
"next.truncated": False,
@@ -1842,7 +1842,7 @@ def test_save_load_with_custom_converter_functions():
# Verify it uses default converters by checking with standard batch format
batch = {
OBS_IMAGE: torch.randn(1, 3, 32, 32),
"action": torch.randn(1, 7),
ACTION: torch.randn(1, 7),
"next.reward": torch.tensor([1.0]),
"next.done": torch.tensor([False]),
"next.truncated": torch.tensor([False]),
@@ -2094,11 +2094,11 @@ def test_aggregate_joint_action_only():
patterns=["action.j1.pos", "action.j2.pos"],
)
# Expect only "action" with joint names
assert "action" in out and OBS_STATE not in out
assert out["action"]["dtype"] == "float32"
assert set(out["action"]["names"]) == {"j1.pos", "j2.pos"}
assert out["action"]["shape"] == (len(out["action"]["names"]),)
# Expect only ACTION with joint names
assert ACTION in out and OBS_STATE not in out
assert out[ACTION]["dtype"] == "float32"
assert set(out[ACTION]["names"]) == {"j1.pos", "j2.pos"}
assert out[ACTION]["shape"] == (len(out[ACTION]["names"]),)
def test_aggregate_ee_action_and_observation_with_videos():
@@ -2113,9 +2113,9 @@ def test_aggregate_ee_action_and_observation_with_videos():
)
# Action should pack only EE names
assert "action" in out
assert set(out["action"]["names"]) == {"ee.x", "ee.y"}
assert out["action"]["dtype"] == "float32"
assert ACTION in out
assert set(out[ACTION]["names"]) == {"ee.x", "ee.y"}
assert out[ACTION]["dtype"] == "float32"
# Observation state should pack both ee.x and j1.pos as a vector
assert OBS_STATE in out
@@ -2140,10 +2140,10 @@ def test_aggregate_both_action_types():
patterns=["action.ee", "action.j1", "action.j2.pos"],
)
assert "action" in out
assert ACTION in out
expected = {"ee.x", "ee.y", "j1.pos", "j2.pos"}
assert set(out["action"]["names"]) == expected
assert out["action"]["shape"] == (len(expected),)
assert set(out[ACTION]["names"]) == expected
assert out[ACTION]["shape"] == (len(expected),)
def test_aggregate_images_when_use_videos_false():

15
tests/processor/test_policy_robot_bridge.py
View File

@@ -28,6 +28,7 @@ from lerobot.processor import (
RobotActionToPolicyActionProcessorStep,
)
from lerobot.processor.converters import identity_transition
from lerobot.utils.constants import ACTION
from tests.conftest import assert_contract_is_typed
@@ -134,8 +135,8 @@ def test_robot_to_policy_transform_features():
transformed = processor.transform_features(features)
assert "action" in transformed[PipelineFeatureType.ACTION]
action_feature = transformed[PipelineFeatureType.ACTION]["action"]
assert ACTION in transformed[PipelineFeatureType.ACTION]
action_feature = transformed[PipelineFeatureType.ACTION][ACTION]
assert action_feature.type == FeatureType.ACTION
assert action_feature.shape == (3,)
@@ -251,7 +252,7 @@ def test_policy_to_robot_transform_features():
features = {
PipelineFeatureType.ACTION: {
"action": {"type": FeatureType.ACTION, "shape": (2,)},
ACTION: {"type": FeatureType.ACTION, "shape": (2,)},
"other_data": {"type": FeatureType.ENV, "shape": (1,)},
}
}
@@ -266,7 +267,7 @@ def test_policy_to_robot_transform_features():
assert motor_feature.type == FeatureType.ACTION
assert motor_feature.shape == (1,)
assert "action" in transformed[PipelineFeatureType.ACTION]
assert ACTION in transformed[PipelineFeatureType.ACTION]
assert "other_data" in transformed[PipelineFeatureType.ACTION]
@@ -447,8 +448,8 @@ def test_robot_to_policy_features_contract(policy_feature_factory):
assert_contract_is_typed(out)
assert "action" in out[PipelineFeatureType.ACTION]
action_feature = out[PipelineFeatureType.ACTION]["action"]
assert ACTION in out[PipelineFeatureType.ACTION]
action_feature = out[PipelineFeatureType.ACTION][ACTION]
assert action_feature.type == FeatureType.ACTION
assert action_feature.shape == (2,)
@@ -458,7 +459,7 @@ def test_policy_to_robot_features_contract(policy_feature_factory):
processor = PolicyActionToRobotActionProcessorStep(motor_names=["m1", "m2", "m3"])
features = {
PipelineFeatureType.ACTION: {
"action": policy_feature_factory(FeatureType.ACTION, (3,)),
ACTION: policy_feature_factory(FeatureType.ACTION, (3,)),
"other": policy_feature_factory(FeatureType.ENV, (1,)),
}
}

4
tests/processor/test_rename_processor.py
View File

@@ -28,7 +28,7 @@ from lerobot.processor import (
)
from lerobot.processor.converters import create_transition, identity_transition
from lerobot.processor.rename_processor import rename_stats
from lerobot.utils.constants import OBS_IMAGE, OBS_IMAGES, OBS_STATE
from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_IMAGES, OBS_STATE
from tests.conftest import assert_contract_is_typed
@@ -488,7 +488,7 @@ def test_features_chained_processors(policy_feature_factory):
def test_rename_stats_basic():
orig = {
OBS_STATE: {"mean": np.array([0.0]), "std": np.array([1.0])},
"action": {"mean": np.array([0.0])},
ACTION: {"mean": np.array([0.0])},
}
mapping = {OBS_STATE: "observation.robot_state"}
renamed = rename_stats(orig, mapping)

6
tests/processor/test_tokenizer_processor.py
View File

@@ -11,7 +11,7 @@ import torch
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep, TransitionKey
from lerobot.processor.converters import create_transition, identity_transition
from lerobot.utils.constants import OBS_IMAGE, OBS_LANGUAGE, OBS_STATE
from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_LANGUAGE, OBS_STATE
from tests.utils import require_package
@@ -504,14 +504,14 @@ def test_features_basic():
input_features = {
PipelineFeatureType.OBSERVATION: {OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(10,))},
PipelineFeatureType.ACTION: {"action": PolicyFeature(type=FeatureType.ACTION, shape=(5,))},
PipelineFeatureType.ACTION: {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(5,))},
}
output_features = processor.transform_features(input_features)
# Check that original features are preserved
assert OBS_STATE in output_features[PipelineFeatureType.OBSERVATION]
assert "action" in output_features[PipelineFeatureType.ACTION]
assert ACTION in output_features[PipelineFeatureType.ACTION]
# Check that tokenized features are added
assert f"{OBS_LANGUAGE}.tokens" in output_features[PipelineFeatureType.OBSERVATION]