admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-05-26 05:59:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

require cuda in tests

Browse Source
This commit is contained in:
Jade Choghari
2025-11-25 22:51:16 +01:00
parent 81cf4d8ed5
commit 825146d218
1 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/policies/xvla/test_xvla_original_vs_lerobot.py
+4
View File
@@ -34,6 +34,7 @@ from lerobot.policies.xvla.modeling_xvla import XVLAPolicy
from lerobot.policies.xvla.processor_xvla import make_xvla_pre_post_processors from lerobot.policies.xvla.processor_xvla import make_xvla_pre_post_processors
from lerobot.processor import PolicyAction, PolicyProcessorPipeline # noqa: E402 from lerobot.processor import PolicyAction, PolicyProcessorPipeline # noqa: E402
from lerobot.utils.constants import OBS_IMAGES, OBS_STATE # noqa: E402 from lerobot.utils.constants import OBS_IMAGES, OBS_STATE # noqa: E402
from tests.utils import require_cuda # noqa: E402
# Constants # Constants
DUMMY_ACTION_DIM = 7 # Standard robot arm action dimension DUMMY_ACTION_DIM = 7 # Standard robot arm action dimension
@@ -164,6 +165,7 @@ def preprocessor(xvla_components):
return xvla_components[1] return xvla_components[1]
@require_cuda
def test_xvla_preprocessor_alignment(policy, preprocessor): def test_xvla_preprocessor_alignment(policy, preprocessor):
"""Test that LeRobot XVLA preprocessor produces expected outputs.""" """Test that LeRobot XVLA preprocessor produces expected outputs."""
print("\n" + "=" * 80) print("\n" + "=" * 80)
@@ -210,6 +212,7 @@ def test_xvla_preprocessor_alignment(policy, preprocessor):
print("\nAll preprocessor outputs have correct shapes!") print("\nAll preprocessor outputs have correct shapes!")
@require_cuda
def test_xvla_action_generation(policy, preprocessor): def test_xvla_action_generation(policy, preprocessor):
"""Test XVLA LeRobot implementation generates expected actions.""" """Test XVLA LeRobot implementation generates expected actions."""
print("\n" + "=" * 80) print("\n" + "=" * 80)
@@ -287,6 +290,7 @@ def test_xvla_action_generation(policy, preprocessor):
print(f"\nSuccess: Actions match expected values within tolerance ({tolerance})!") print(f"\nSuccess: Actions match expected values within tolerance ({tolerance})!")
@require_cuda
def test_xvla_inference_reproducibility(policy, preprocessor): def test_xvla_inference_reproducibility(policy, preprocessor):
"""Test that XVLA inference is reproducible with the same seed.""" """Test that XVLA inference is reproducible with the same seed."""
print("\n" + "=" * 80) print("\n" + "=" * 80)