[Fix] Unpin torch beyond 2.6.0 & torchcodec beyond 0.2.1 (#1127)

AdilZouitine · AdilZouitine · commit ecb962fd7bef · 2025-08-10T19:24:47.000+02:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -68,8 +68,8 @@ dependencies = [
     "pyzmq>=26.2.1",
     "rerun-sdk>=0.21.0",
     "termcolor>=2.4.0",
-    "torch>=2.2.1,<2.7",
-    "torchcodec==0.2.1; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')",
+    "torch>=2.2.1",
+    "torchcodec>=0.2.1; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')",
     "torchvision>=0.21.0",
     "wandb>=0.16.3",
     "zarr>=2.17.0",
diff --git a/tests/artifacts/image_transforms/default_transforms.safetensors b/tests/artifacts/image_transforms/default_transforms.safetensors
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0389a716d51c1c615fb2a3bfa386d89f00b0deca08c4fa21b23e020a939d0213
+oid sha256:6b1e600768a8771c5fe650e038a1193597e3810f032041b2a0d021e4496381c1
 size 3686488
diff --git a/tests/artifacts/image_transforms/save_image_transforms_to_safetensors.py b/tests/artifacts/image_transforms/save_image_transforms_to_safetensors.py
@@ -28,7 +28,7 @@
 from lerobot.common.utils.random_utils import seeded_context
 
 ARTIFACT_DIR = Path("tests/artifacts/image_transforms")
-DATASET_REPO_ID = "lerobot/aloha_mobile_shrimp"
+DATASET_REPO_ID = "lerobot/aloha_static_cups_open"
 
 
 def save_default_config_transform(original_frame: torch.Tensor, output_dir: Path):
diff --git a/tests/artifacts/image_transforms/single_transforms.safetensors b/tests/artifacts/image_transforms/single_transforms.safetensors
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0dc691503e7d90b2086bb408e89a65f772ce5ee6e3562ef8c127bcb09bd90851
+oid sha256:9d4ebab73eabddc58879a4e770289d19e00a1a4cf2fa5fa33cd3a3246992bc90
 size 40551392
diff --git a/tests/artifacts/policies/aloha_sim_insertion_human_act_/actions.safetensors b/tests/artifacts/policies/aloha_sim_insertion_human_act_/actions.safetensors
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cc67af1d60f95d84c98d6c9ebd648990e0f0705368bd6b72d2b39533950b0179
+oid sha256:f3e4c8e85e146b043fd4e4984947c2a6f01627f174a19f18b5914cf690579d77
 size 5104
diff --git a/tests/artifacts/policies/aloha_sim_insertion_human_act_/grad_stats.safetensors b/tests/artifacts/policies/aloha_sim_insertion_human_act_/grad_stats.safetensors
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:64518cf652105d15f5fd2cfc13d0681f66a4ec4797dc5d5dc2f7b0d91fe5dfd6
+oid sha256:1a7a8b1a457149109f843c32bcbb047d09de2201847b9b79f7501b447f77ecf4
 size 31672
diff --git a/tests/artifacts/policies/aloha_sim_insertion_human_act_/output_dict.safetensors b/tests/artifacts/policies/aloha_sim_insertion_human_act_/output_dict.safetensors
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:32b6d14fab4244b5140adb345e47f662b6739c04974e04b21c3127caa988abbb
+oid sha256:5e6ce85296b2009e7c2060d336c0429b1c7197d9adb159e7df0ba18003067b36
 size 68
diff --git a/tests/artifacts/policies/aloha_sim_insertion_human_act_/param_stats.safetensors b/tests/artifacts/policies/aloha_sim_insertion_human_act_/param_stats.safetensors
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e1904ef0338f7b6efdec70ec235ee931b5751008bf4eb433edb0b3fa0838a4f1
+oid sha256:9b5f557e30aead3731c38cbd85af8c706395d8689a918ad88805b5a886245603
 size 33400
diff --git a/tests/artifacts/policies/aloha_sim_insertion_human_act_1000_steps/actions.safetensors b/tests/artifacts/policies/aloha_sim_insertion_human_act_1000_steps/actions.safetensors
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fa544a97f00bf46393a09b006b44c2499bbf7d177782360a8c21cacbf200c07a
+oid sha256:2e6625cabfeb4800abc80252cf9112a9271c154edd01eb291658f143c951610b
 size 515400
diff --git a/tests/artifacts/policies/aloha_sim_insertion_human_act_1000_steps/grad_stats.safetensors b/tests/artifacts/policies/aloha_sim_insertion_human_act_1000_steps/grad_stats.safetensors
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:83c7a8ae912300b5cedba31904f7ba22542059fd60dd86548a95e415713f719e
+oid sha256:224b5fa4828aa88171b68c036e8919c1eae563e2113f03b6461eadf5bf8525a6
 size 31672
diff --git a/tests/artifacts/policies/aloha_sim_insertion_human_act_1000_steps/output_dict.safetensors b/tests/artifacts/policies/aloha_sim_insertion_human_act_1000_steps/output_dict.safetensors
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5a010633237b3a1141603c65174c551daa9e7b4c474af5a1376d73e5425bfb5d
+oid sha256:016d2fa8fe5f58017dfd46f4632fdc19dfd751e32a2c7cde2077c6f95546d6bd
 size 68
diff --git a/tests/artifacts/policies/aloha_sim_insertion_human_act_1000_steps/param_stats.safetensors b/tests/artifacts/policies/aloha_sim_insertion_human_act_1000_steps/param_stats.safetensors
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ec8b5c440e9fcec190c9be48b28ebb79f82ae63626afe7c811e4bb0c3dd08842
+oid sha256:021562ee3e4814425e367ed0c144d6fbe2eb28838247085716cf0b58fd69a075
 size 33400
diff --git a/tests/datasets/test_image_transforms.py b/tests/datasets/test_image_transforms.py
@@ -16,6 +16,7 @@
 
 import pytest
 import torch
+from packaging import version
 from safetensors.torch import load_file
 from torchvision.transforms import v2
 from torchvision.transforms.v2 import functional as F  # noqa: N812
@@ -253,7 +254,14 @@ def test_backward_compatibility_single_transforms(
 
 
 @require_x86_64_kernel
+@pytest.mark.skipif(
+    version.parse(torch.__version__) < version.parse("2.7.0"),
+    reason="Test artifacts were generated with PyTorch >= 2.7.0 which has different multinomial behavior",
+)
 def test_backward_compatibility_default_config(img_tensor, default_transforms):
+    # NOTE: PyTorch versions have different randomness, it might break this test.
+    # See this PR: https://github.com/huggingface/lerobot/pull/1127.
+
     cfg = ImageTransformsConfig(enable=True)
     default_tf = ImageTransforms(cfg)
 
diff --git a/tests/optim/test_schedulers.py b/tests/optim/test_schedulers.py
@@ -37,7 +37,6 @@ def test_diffuser_scheduler(optimizer):
         "base_lrs": [0.001],
         "last_epoch": 1,
         "lr_lambdas": [None],
-        "verbose": False,
     }
     assert scheduler.state_dict() == expected_state_dict
 
@@ -56,7 +55,6 @@ def test_vqbet_scheduler(optimizer):
         "base_lrs": [0.001],
         "last_epoch": 1,
         "lr_lambdas": [None],
-        "verbose": False,
     }
     assert scheduler.state_dict() == expected_state_dict
 
@@ -77,7 +75,6 @@ def test_cosine_decay_with_warmup_scheduler(optimizer):
         "base_lrs": [0.001],
         "last_epoch": 1,
         "lr_lambdas": [None],
-        "verbose": False,
     }
     assert scheduler.state_dict() == expected_state_dict
 
diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py
@@ -20,6 +20,7 @@
 import einops
 import pytest
 import torch
+from packaging import version
 from safetensors.torch import load_file
 
 from lerobot import available_policies
@@ -408,7 +409,16 @@ def test_backward_compatibility(ds_repo_id: str, policy_name: str, policy_kwargs
         4. Check that this test now passes.
         5. Remember to restore `tests/scripts/save_policy_to_safetensors.py` to its original state.
         6. Remember to stage and commit the resulting changes to `tests/artifacts`.
+
+    NOTE: If the test does not pass, and you don't change the policy, it is likely that the test artifact
+    is out of date. For example, some PyTorch versions have different randomness, see this PR:
+    https://github.com/huggingface/lerobot/pull/1127.
+
     """
+    # NOTE: ACT policy has different randomness, after PyTorch 2.7.0
+    if policy_name == "act" and version.parse(torch.__version__) < version.parse("2.7.0"):
+        pytest.skip(f"Skipping act policy test with PyTorch {torch.__version__}. Requires PyTorch >= 2.7.0")
+
     ds_name = ds_repo_id.split("/")[-1]
     artifact_dir = Path("tests/artifacts/policies") / f"{ds_name}_{policy_name}_{file_name_extra}"
     saved_output_dict = load_file(artifact_dir / "output_dict.safetensors")
