simplify the test logic, and improve comments

Author: AdilZouitine

tests/artifacts/image_transforms/save_image_transforms_to_safetensors.py

@@ -28,7 +28,6 @@
 from lerobot.common.utils.random_utils import seeded_context
 
 ARTIFACT_DIR = Path("tests/artifacts/image_transforms")
-# TODO: Fix lerobot/aloha_mobile_shrimp dataset
 DATASET_REPO_ID = "lerobot/aloha_static_cups_open"
 
 
@@ -43,16 +42,16 @@ def save_default_config_transform(original_frame: torch.Tensor, output_dir: Path
 
 
 def save_single_transforms(original_frame: torch.Tensor, output_dir: Path):
-    transforms = [
+    transforms = {
         ("ColorJitter", "brightness", [(0.5, 0.5), (2.0, 2.0)]),
         ("ColorJitter", "contrast", [(0.5, 0.5), (2.0, 2.0)]),
         ("ColorJitter", "saturation", [(0.5, 0.5), (2.0, 2.0)]),
         ("ColorJitter", "hue", [(-0.25, -0.25), (0.25, 0.25)]),
         ("SharpnessJitter", "sharpness", [(0.5, 0.5), (2.0, 2.0)]),
-    ]
+    }
 
     frames = {"original_frame": original_frame}
-    for tf_type, tf_name, min_max_values in transforms:
+    for tf_type, tf_name, min_max_values in transforms.items():
         for min_max in min_max_values:
             tf_cfg = ImageTransformConfig(type=tf_type, kwargs={tf_name: min_max})
             tf = make_transform_from_config(tf_cfg)

tests/datasets/test_image_transforms.py

@@ -254,15 +254,13 @@ 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):
-    # This test depends on the behavior of torch.multinomial, which changed in PyTorch 2.7.0.
-    # The test artifacts (default_transforms.safetensors) were generated using PyTorch >= 2.7.0.
-    # For more details, see:
-    # - PyTorch issue: https://github.com/pytorch/pytorch/issues/154031
-    # - LeRobot PR: https://github.com/huggingface/lerobot/pull/1127
-    # If running with PyTorch < 2.7.0, this test is expected to fail.
-    if version.parse(torch.__version__) < version.parse("2.7.0"):
-        pytest.skip(f"Skipping test with PyTorch {torch.__version__}. Requires PyTorch >= 2.7.0")
+    # 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)

tests/policies/test_policies.py

@@ -415,9 +415,7 @@ def test_backward_compatibility(ds_repo_id: str, policy_name: str, policy_kwargs
     https://github.com/huggingface/lerobot/pull/1127.
 
     """
-    # Skip the test for act policy if PyTorch version is below 2.7
-    # The MultiheadSelfAttention randomness changed between PyTorch 2.6 and 2.7
-    # The artifacts were generated with PyTorch 2.7, so the test will break with act policy if PyTorch is below 2.7
+    # 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")