Merge remote-tracking branch 'upstream/main' into add_drop_last_keyframes

tests/conftest.py

@@ -1,3 +1,18 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from .utils import DEVICE
 
 

tests/scripts/save_dataset_to_safetensors.py

@@ -1,3 +1,18 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 """
 This script provides a utility for saving a dataset as safetensors files for the purpose of testing backward compatibility
 when updating the data format. It uses the `PushtDataset` to create a DataLoader and saves selected frame from the

tests/scripts/save_policy_to_safetensor.py

@@ -1,3 +1,18 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import shutil
 from pathlib import Path
 

tests/test_available.py

@@ -1,3 +1,18 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import importlib
 
 import gymnasium as gym
@@ -15,7 +30,7 @@ from tests.utils import require_env
 def test_available_env_task(env_name: str, task_name: list):
     """
     This test verifies that all environments listed in `lerobot/__init__.py` can
-    be sucessfully imported — if they're installed — and that their
+    be successfully imported — if they're installed — and that their
     `available_tasks_per_env` are valid.
     """
     package_name = f"gym_{env_name}"

tests/test_datasets.py

@@ -1,3 +1,18 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import json
 import logging
 from copy import deepcopy
@@ -41,7 +56,7 @@ def test_factory(env_name, repo_id, policy_name):
     )
     dataset = make_dataset(cfg)
     delta_timestamps = dataset.delta_timestamps
-    image_keys = dataset.image_keys
+    camera_keys = dataset.camera_keys
 
     item = dataset[0]
 
@@ -71,7 +86,7 @@ def test_factory(env_name, repo_id, policy_name):
         else:
             assert item[key].ndim == ndim, f"{key}"
 
-        if key in image_keys:
+        if key in camera_keys:
             assert item[key].dtype == torch.float32, f"{key}"
             # TODO(rcadene): we assume for now that image normalization takes place in the model
             assert item[key].max() <= 1.0, f"{key}"

tests/test_envs.py

@@ -1,3 +1,18 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import importlib
 
 import gymnasium as gym

tests/test_examples.py

@@ -1,8 +1,25 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 # TODO(aliberts): Mute logging for these tests
 import subprocess
 import sys
 from pathlib import Path
 
+from tests.utils import require_package
+
 
 def _find_and_replace(text: str, finds_and_replaces: list[tuple[str, str]]) -> str:
     for f, r in finds_and_replaces:
@@ -21,6 +38,7 @@ def test_example_1():
     assert Path("outputs/examples/1_load_lerobot_dataset/episode_0.mp4").exists()
 
 
+@require_package("gym_pusht")
 def test_examples_3_and_2():
     """
     Train a model with example 3, check the outputs.
@@ -46,7 +64,7 @@ def test_examples_3_and_2():
     # Pass empty globals to allow dictionary comprehension https://stackoverflow.com/a/32897127/4391249.
     exec(file_contents, {})
 
-    for file_name in ["model.safetensors", "config.json", "config.yaml"]:
+    for file_name in ["model.safetensors", "config.json"]:
         assert Path(f"outputs/train/example_pusht_diffusion/{file_name}").exists()
 
     path = "examples/2_evaluate_pretrained_policy.py"
@@ -58,16 +76,16 @@ def test_examples_3_and_2():
     file_contents = _find_and_replace(
         file_contents,
         [
-            ('pretrained_policy_name = "lerobot/diffusion_pusht"', ""),
-            ("pretrained_policy_path = Path(snapshot_download(pretrained_policy_name))", ""),
+            ('pretrained_policy_path = Path(snapshot_download("lerobot/diffusion_pusht"))', ""),
             (
                 '# pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")',
                 'pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")',
             ),
-            ('"eval.n_episodes=10"', '"eval.n_episodes=1"'),
-            ('"eval.batch_size=10"', '"eval.batch_size=1"'),
-            ('"device=cuda"', '"device=cpu"'),
+            ('device = torch.device("cuda")', 'device = torch.device("cpu")'),
+            ("step += 1", "break"),
         ],
     )
 
-    assert Path("outputs/train/example_pusht_diffusion").exists()
+    exec(file_contents, {})
+
+    assert Path("outputs/eval/example_pusht_diffusion/rollout.mp4").exists()

tests/test_policies.py

@@ -1,3 +1,18 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import inspect
 from pathlib import Path
 
@@ -49,6 +64,14 @@ def test_get_policy_and_config_classes(policy_name: str):
             "act",
             ["env.task=AlohaTransferCube-v0", "dataset_repo_id=lerobot/aloha_sim_transfer_cube_scripted"],
         ),
+        # Note: these parameters also need custom logic in the test function for overriding the Hydra config.
+        (
+            "aloha",
+            "diffusion",
+            ["env.task=AlohaInsertion-v0", "dataset_repo_id=lerobot/aloha_sim_insertion_human"],
+        ),
+        # Note: these parameters also need custom logic in the test function for overriding the Hydra config.
+        ("pusht", "act", ["env.task=PushT-v0", "dataset_repo_id=lerobot/pusht"]),
     ],
 )
 @require_env
@@ -72,6 +95,31 @@ def test_policy(env_name, policy_name, extra_overrides):
         + extra_overrides,
     )
 
+    # Additional config override logic.
+    if env_name == "aloha" and policy_name == "diffusion":
+        for keys in [
+            ("training", "delta_timestamps"),
+            ("policy", "input_shapes"),
+            ("policy", "input_normalization_modes"),
+        ]:
+            dct = dict(cfg[keys[0]][keys[1]])
+            dct["observation.images.top"] = dct["observation.image"]
+            del dct["observation.image"]
+            cfg[keys[0]][keys[1]] = dct
+        cfg.override_dataset_stats = None
+
+    # Additional config override logic.
+    if env_name == "pusht" and policy_name == "act":
+        for keys in [
+            ("policy", "input_shapes"),
+            ("policy", "input_normalization_modes"),
+        ]:
+            dct = dict(cfg[keys[0]][keys[1]])
+            dct["observation.image"] = dct["observation.images.top"]
+            del dct["observation.images.top"]
+            cfg[keys[0]][keys[1]] = dct
+        cfg.override_dataset_stats = None
+
     # Check that we can make the policy object.
     dataset = make_dataset(cfg)
     policy = make_policy(hydra_cfg=cfg, dataset_stats=dataset.stats)
@@ -236,7 +284,7 @@ def test_normalize(insert_temporal_dim):
 @pytest.mark.parametrize(
     "env_name, policy_name, extra_overrides",
     [
-        # ("xarm", "tdmpc", ["policy.n_action_repeats=2"]),
+        ("xarm", "tdmpc", []),
         (
             "pusht",
             "diffusion",

tests/test_utils.py

@@ -0,0 +1,38 @@
+import random
+from typing import Callable
+
+import numpy as np
+import pytest
+import torch
+
+from lerobot.common.utils.utils import seeded_context, set_global_seed
+
+
+@pytest.mark.parametrize(
+    "rand_fn",
+    [
+        random.random,
+        np.random.random,
+        lambda: torch.rand(1).item(),
+    ]
+    + [lambda: torch.rand(1, device="cuda")]
+    if torch.cuda.is_available()
+    else [],
+)
+def test_seeding(rand_fn: Callable[[], int]):
+    set_global_seed(0)
+    a = rand_fn()
+    with seeded_context(1337):
+        c = rand_fn()
+    b = rand_fn()
+    set_global_seed(0)
+    a_ = rand_fn()
+    b_ = rand_fn()
+    # Check that `set_global_seed` lets us reproduce a and b.
+    assert a_ == a
+    # Additionally, check that the `seeded_context` didn't interrupt the global RNG.
+    assert b_ == b
+    set_global_seed(1337)
+    c_ = rand_fn()
+    # Check that `seeded_context` and `global_seed` give the same reproducibility.
+    assert c_ == c

tests/test_visualize_dataset.py

@@ -1,3 +1,18 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import pytest
 
 from lerobot.scripts.visualize_dataset import visualize_dataset

tests/utils.py

@@ -1,4 +1,20 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import platform
+from functools import wraps
 
 import pytest
 import torch
@@ -61,7 +77,6 @@ def require_env(func):
     Decorator that skips the test if the required environment package is not installed.
     As it need 'env_name' in args, it also checks whether it is provided as an argument.
     """
-    from functools import wraps
 
     @wraps(func)
     def wrapper(*args, **kwargs):
@@ -82,3 +97,20 @@ def require_env(func):
         return func(*args, **kwargs)
 
     return wrapper
+
+
+def require_package(package_name):
+    """
+    Decorator that skips the test if the specified package is not installed.
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            if not is_package_available(package_name):
+                pytest.skip(f"{package_name} not installed")
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator