admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-21 00:42:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into feat/add_pi_conflicts_main

Browse Source
This commit is contained in:
Steven Palma
2025-09-26 13:38:31 +02:00
parent 18db813e2d d2782cf66b
commit 9ab7be7abb
68 changed files with 963 additions and 865 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/datasets/test_compute_stats.py
+17 -16
View File
@@ -29,6 +29,7 @@ from lerobot.datasets.compute_stats import (
sample_images,
sample_indices,
)
from lerobot.utils.constants import OBS_IMAGE, OBS_STATE
def mock_load_image_as_numpy(path, dtype, channel_first):
@@ -143,21 +144,21 @@ def test_get_feature_stats_single_value():
def test_compute_episode_stats():
episode_data = {
"observation.image": [f"image_{i}.jpg" for i in range(100)],
"observation.state": np.random.rand(100, 10),
OBS_IMAGE: [f"image_{i}.jpg" for i in range(100)],
OBS_STATE: np.random.rand(100, 10),
}
features = {
"observation.image": {"dtype": "image"},
"observation.state": {"dtype": "numeric"},
OBS_IMAGE: {"dtype": "image"},
OBS_STATE: {"dtype": "numeric"},
}
with patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy):
stats = compute_episode_stats(episode_data, features)
assert "observation.image" in stats and "observation.state" in stats
assert stats["observation.image"]["count"].item() == 100
assert stats["observation.state"]["count"].item() == 100
assert stats["observation.image"]["mean"].shape == (3, 1, 1)
assert OBS_IMAGE in stats and OBS_STATE in stats
assert stats[OBS_IMAGE]["count"].item() == 100
assert stats[OBS_STATE]["count"].item() == 100
assert stats[OBS_IMAGE]["mean"].shape == (3, 1, 1)
def test_assert_type_and_shape_valid():
@@ -231,38 +232,38 @@ def test_aggregate_feature_stats():
def test_aggregate_stats():
all_stats = [
{
"observation.image": {
OBS_IMAGE: {
"min": [1, 2, 3],
"max": [10, 20, 30],
"mean": [5.5, 10.5, 15.5],
"std": [2.87, 5.87, 8.87],
"count": 10,
},
"observation.state": {"min": 1, "max": 10, "mean": 5.5, "std": 2.87, "count": 10},
OBS_STATE: {"min": 1, "max": 10, "mean": 5.5, "std": 2.87, "count": 10},
"extra_key_0": {"min": 5, "max": 25, "mean": 15, "std": 6, "count": 6},
},
{
"observation.image": {
OBS_IMAGE: {
"min": [2, 1, 0],
"max": [15, 10, 5],
"mean": [8.5, 5.5, 2.5],
"std": [3.42, 2.42, 1.42],
"count": 15,
},
"observation.state": {"min": 2, "max": 15, "mean": 8.5, "std": 3.42, "count": 15},
OBS_STATE: {"min": 2, "max": 15, "mean": 8.5, "std": 3.42, "count": 15},
"extra_key_1": {"min": 0, "max": 20, "mean": 10, "std": 5, "count": 5},
},
]
expected_agg_stats = {
"observation.image": {
OBS_IMAGE: {
"min": [1, 1, 0],
"max": [15, 20, 30],
"mean": [7.3, 7.5, 7.7],
"std": [3.5317, 4.8267, 8.5581],
"count": 25,
},
"observation.state": {
OBS_STATE: {
"min": 1,
"max": 15,
"mean": 7.3,
@@ -290,7 +291,7 @@ def test_aggregate_stats():
for fkey, stats in ep_stats.items():
for k in stats:
stats[k] = np.array(stats[k], dtype=np.int64 if k == "count" else np.float32)
if fkey == "observation.image" and k != "count":
if fkey == OBS_IMAGE and k != "count":
stats[k] = stats[k].reshape(3, 1, 1) # for normalization on image channels
else:
stats[k] = stats[k].reshape(1)
@@ -299,7 +300,7 @@ def test_aggregate_stats():
for fkey, stats in expected_agg_stats.items():
for k in stats:
stats[k] = np.array(stats[k], dtype=np.int64 if k == "count" else np.float32)
if fkey == "observation.image" and k != "count":
if fkey == OBS_IMAGE and k != "count":
stats[k] = stats[k].reshape(3, 1, 1) # for normalization on image channels
else:
stats[k] = stats[k].reshape(1)
tests/datasets/test_dataset_utils.py
+18 -17
View File
@@ -21,6 +21,7 @@ from huggingface_hub import DatasetCard
from lerobot.datasets.push_dataset_to_hub.utils import calculate_episode_data_index
from lerobot.datasets.utils import combine_feature_dicts, create_lerobot_dataset_card, hf_transform_to_torch
from lerobot.utils.constants import ACTION, OBS_IMAGES
def test_default_parameters():
@@ -58,14 +59,14 @@ def test_calculate_episode_data_index():
def test_merge_simple_vectors():
g1 = {
"action": {
ACTION: {
"dtype": "float32",
"shape": (2,),
"names": ["ee.x", "ee.y"],
}
}
g2 = {
"action": {
ACTION: {
"dtype": "float32",
"shape": (2,),
"names": ["ee.y", "ee.z"],
@@ -74,36 +75,36 @@ def test_merge_simple_vectors():
out = combine_feature_dicts(g1, g2)
assert "action" in out
assert out["action"]["dtype"] == "float32"
assert ACTION in out
assert out[ACTION]["dtype"] == "float32"
# Names merged with preserved order and de-dupuplication
assert out["action"]["names"] == ["ee.x", "ee.y", "ee.z"]
assert out[ACTION]["names"] == ["ee.x", "ee.y", "ee.z"]
# Shape correctly recomputed from names length
assert out["action"]["shape"] == (3,)
assert out[ACTION]["shape"] == (3,)
def test_merge_multiple_groups_order_and_dedup():
g1 = {"action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}}
g2 = {"action": {"dtype": "float32", "shape": (2,), "names": ["b", "c"]}}
g3 = {"action": {"dtype": "float32", "shape": (3,), "names": ["a", "c", "d"]}}
g1 = {ACTION: {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}}
g2 = {ACTION: {"dtype": "float32", "shape": (2,), "names": ["b", "c"]}}
g3 = {ACTION: {"dtype": "float32", "shape": (3,), "names": ["a", "c", "d"]}}
out = combine_feature_dicts(g1, g2, g3)
assert out["action"]["names"] == ["a", "b", "c", "d"]
assert out["action"]["shape"] == (4,)
assert out[ACTION]["names"] == ["a", "b", "c", "d"]
assert out[ACTION]["shape"] == (4,)
def test_non_vector_last_wins_for_images():
# Non-vector (images) with same name should be overwritten by the last image specified
g1 = {
"observation.images.front": {
f"{OBS_IMAGES}.front": {
"dtype": "image",
"shape": (3, 480, 640),
"names": ["channels", "height", "width"],
}
}
g2 = {
"observation.images.front": {
f"{OBS_IMAGES}.front": {
"dtype": "image",
"shape": (3, 720, 1280),
"names": ["channels", "height", "width"],
@@ -111,13 +112,13 @@ def test_non_vector_last_wins_for_images():
}
out = combine_feature_dicts(g1, g2)
assert out["observation.images.front"]["shape"] == (3, 720, 1280)
assert out["observation.images.front"]["dtype"] == "image"
assert out[f"{OBS_IMAGES}.front"]["shape"] == (3, 720, 1280)
assert out[f"{OBS_IMAGES}.front"]["dtype"] == "image"
def test_dtype_mismatch_raises():
g1 = {"action": {"dtype": "float32", "shape": (1,), "names": ["a"]}}
g2 = {"action": {"dtype": "float64", "shape": (1,), "names": ["b"]}}
g1 = {ACTION: {"dtype": "float32", "shape": (1,), "names": ["a"]}}
g2 = {ACTION: {"dtype": "float64", "shape": (1,), "names": ["b"]}}
with pytest.raises(ValueError, match="dtype mismatch for 'action'"):
_ = combine_feature_dicts(g1, g2)
tests/datasets/test_datasets.py
+13 -12
View File
@@ -46,6 +46,7 @@ from lerobot.datasets.utils import (
from lerobot.envs.factory import make_env_config
from lerobot.policies.factory import make_policy_config
from lerobot.robots import make_robot_from_config
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR
from tests.fixtures.constants import DUMMY_CHW, DUMMY_HWC, DUMMY_REPO_ID
from tests.mocks.mock_robot import MockRobotConfig
from tests.utils import require_x86_64_kernel
@@ -74,8 +75,8 @@ def test_same_attributes_defined(tmp_path, lerobot_dataset_factory):
"""
# Instantiate both ways
robot = make_robot_from_config(MockRobotConfig())
action_features = hw_to_dataset_features(robot.action_features, "action", True)
obs_features = hw_to_dataset_features(robot.observation_features, "observation", True)
action_features = hw_to_dataset_features(robot.action_features, ACTION, True)
obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR, True)
dataset_features = {**action_features, **obs_features}
root_create = tmp_path / "create"
dataset_create = LeRobotDataset.create(
@@ -392,12 +393,12 @@ def test_factory(env_name, repo_id, policy_name):
item = dataset[0]
keys_ndim_required = [
("action", 1, True),
(ACTION, 1, True),
("episode_index", 0, True),
("frame_index", 0, True),
("timestamp", 0, True),
# TODO(rcadene): should we rename it agent_pos?
("observation.state", 1, True),
(OBS_STATE, 1, True),
("next.reward", 0, False),
("next.done", 0, False),
]
@@ -662,12 +663,12 @@ def test_check_cached_episodes_sufficient(tmp_path, lerobot_dataset_factory):
def test_update_chunk_settings(tmp_path, empty_lerobot_dataset_factory):
"""Test the update_chunk_settings functionality for both LeRobotDataset and LeRobotDatasetMetadata."""
features = {
"observation.state": {
OBS_STATE: {
"dtype": "float32",
"shape": (6,),
"names": ["shoulder_pan", "shoulder_lift", "elbow", "wrist_1", "wrist_2", "wrist_3"],
},
"action": {
ACTION: {
"dtype": "float32",
"shape": (6,),
"names": ["shoulder_pan", "shoulder_lift", "elbow", "wrist_1", "wrist_2", "wrist_3"],
@@ -769,12 +770,12 @@ def test_update_chunk_settings(tmp_path, empty_lerobot_dataset_factory):
def test_update_chunk_settings_video_dataset(tmp_path):
"""Test update_chunk_settings with a video dataset to ensure video-specific logic works."""
features = {
"observation.images.cam": {
f"{OBS_IMAGES}.cam": {
"dtype": "video",
"shape": (480, 640, 3),
"names": ["height", "width", "channels"],
},
"action": {"dtype": "float32", "shape": (6,), "names": ["j1", "j2", "j3", "j4", "j5", "j6"]},
ACTION: {"dtype": "float32", "shape": (6,), "names": ["j1", "j2", "j3", "j4", "j5", "j6"]},
}
# Create video dataset
@@ -841,7 +842,7 @@ def test_multi_episode_metadata_consistency(tmp_path, empty_lerobot_dataset_fact
"""Test episode metadata consistency across multiple episodes."""
features = {
"state": {"dtype": "float32", "shape": (3,), "names": ["x", "y", "z"]},
"action": {"dtype": "float32", "shape": (2,), "names": ["v", "w"]},
ACTION: {"dtype": "float32", "shape": (2,), "names": ["v", "w"]},
}
dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features, use_videos=False)
@@ -851,7 +852,7 @@ def test_multi_episode_metadata_consistency(tmp_path, empty_lerobot_dataset_fact
for episode_idx in range(num_episodes):
for _ in range(frames_per_episode[episode_idx]):
dataset.add_frame({"state": torch.randn(3), "action": torch.randn(2), "task": tasks[episode_idx]})
dataset.add_frame({"state": torch.randn(3), ACTION: torch.randn(2), "task": tasks[episode_idx]})
dataset.save_episode()
# Load and validate episode metadata
@@ -926,7 +927,7 @@ def test_statistics_metadata_validation(tmp_path, empty_lerobot_dataset_factory)
"""Test that statistics are properly computed and stored for all features."""
features = {
"state": {"dtype": "float32", "shape": (2,), "names": ["pos", "vel"]},
"action": {"dtype": "float32", "shape": (1,), "names": ["force"]},
ACTION: {"dtype": "float32", "shape": (1,), "names": ["force"]},
}
dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features, use_videos=False)
@@ -940,7 +941,7 @@ def test_statistics_metadata_validation(tmp_path, empty_lerobot_dataset_factory)
for frame_idx in range(frames_per_episode[episode_idx]):
state_data = torch.tensor([frame_idx * 0.1, frame_idx * 0.2], dtype=torch.float32)
action_data = torch.tensor([frame_idx * 0.05], dtype=torch.float32)
dataset.add_frame({"state": state_data, "action": action_data, "task": "stats_test"})
dataset.add_frame({"state": state_data, ACTION: action_data, "task": "stats_test"})
dataset.save_episode()
loaded_dataset = LeRobotDataset(dataset.repo_id, root=dataset.root)
tests/datasets/test_streaming.py
+3 -2
View File
@@ -19,6 +19,7 @@ import torch
from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
from lerobot.datasets.utils import safe_shard
from lerobot.utils.constants import ACTION
from tests.fixtures.constants import DUMMY_REPO_ID
@@ -234,7 +235,7 @@ def test_frames_with_delta_consistency(tmp_path, lerobot_dataset_factory, state_
delta_timestamps = {
camera_key: state_deltas,
"state": state_deltas,
"action": action_deltas,
ACTION: action_deltas,
}
ds = lerobot_dataset_factory(
@@ -319,7 +320,7 @@ def test_frames_with_delta_consistency_with_shards(
delta_timestamps = {
camera_key: state_deltas,
"state": state_deltas,
"action": action_deltas,
ACTION: action_deltas,
}
ds = lerobot_dataset_factory(