From 70802668a042319dab14116481b7431b50c9e022 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Fri, 13 Mar 2026 13:24:50 +0100 Subject: [PATCH] test(rewards): restore missing CUDA and mixed precision classifier processor tests --- tests/rewards/test_classifier_processor.py | 160 ++++++++++++++++++++- 1 file changed, 153 insertions(+), 7 deletions(-) diff --git a/tests/rewards/test_classifier_processor.py b/tests/rewards/test_classifier_processor.py index e01515929..c54d80b0e 100644 --- a/tests/rewards/test_classifier_processor.py +++ b/tests/rewards/test_classifier_processor.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python - # Copyright 2025 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -17,6 +15,7 @@ import tempfile +import pytest import torch from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature @@ -46,7 +45,7 @@ def create_default_config(): config.normalization_mapping = { FeatureType.STATE: NormalizationMode.MEAN_STD, FeatureType.VISUAL: NormalizationMode.IDENTITY, - FeatureType.ACTION: NormalizationMode.IDENTITY, + FeatureType.ACTION: NormalizationMode.IDENTITY, # No normalization for classifier output } config.device = "cpu" return config @@ -56,8 +55,8 @@ def create_default_stats(): """Create default dataset statistics for testing.""" return { OBS_STATE: {"mean": torch.zeros(10), "std": torch.ones(10)}, - OBS_IMAGE: {}, - "reward": {}, + OBS_IMAGE: {}, # No normalization for images + "reward": {}, # No normalization for classifier output } @@ -68,14 +67,17 @@ def test_make_classifier_processor_basic(): preprocessor, postprocessor = make_classifier_processor(config, stats) + # Check processor names assert preprocessor.name == "classifier_preprocessor" assert postprocessor.name == "classifier_postprocessor" + # Check steps in preprocessor assert len(preprocessor.steps) == 3 - assert isinstance(preprocessor.steps[0], NormalizerProcessorStep) - assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) + assert isinstance(preprocessor.steps[0], NormalizerProcessorStep) # For input features + assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) # For output features assert isinstance(preprocessor.steps[2], DeviceProcessorStep) + # Check steps in postprocessor assert len(postprocessor.steps) == 2 assert isinstance(postprocessor.steps[0], DeviceProcessorStep) assert isinstance(postprocessor.steps[1], IdentityProcessorStep) @@ -88,6 +90,7 @@ def test_classifier_processor_normalization(): preprocessor, postprocessor = make_classifier_processor(config, stats) + # Create test data observation = { OBS_STATE: torch.randn(10), OBS_IMAGE: torch.randn(3, 224, 224), @@ -96,22 +99,118 @@ def test_classifier_processor_normalization(): transition = create_transition(observation, action) batch = transition_to_batch(transition) + # Process through preprocessor processed = preprocessor(batch) + # Check that data is processed assert processed[OBS_STATE].shape == (10,) assert processed[OBS_IMAGE].shape == (3, 224, 224) assert processed[TransitionKey.ACTION.value].shape == (1,) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_classifier_processor_cuda(): + """Test Classifier processor with CUDA device.""" + config = create_default_config() + config.device = "cuda" + stats = create_default_stats() + + preprocessor, postprocessor = make_classifier_processor(config, stats) + + # Create CPU data + observation = { + OBS_STATE: torch.randn(10), + OBS_IMAGE: torch.randn(3, 224, 224), + } + action = torch.randn(1) + transition = create_transition(observation, action) + batch = transition_to_batch(transition) + + # Process through preprocessor + + processed = preprocessor(batch) + + # Check that data is on CUDA + assert processed[OBS_STATE].device.type == "cuda" + assert processed[OBS_IMAGE].device.type == "cuda" + assert processed[TransitionKey.ACTION.value].device.type == "cuda" + + # Process through postprocessor + postprocessed = postprocessor(processed[TransitionKey.ACTION.value]) + + # Check that output is back on CPU + assert postprocessed.device.type == "cpu" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_classifier_processor_accelerate_scenario(): + """Test Classifier processor in simulated Accelerate scenario.""" + config = create_default_config() + config.device = "cuda:0" + stats = create_default_stats() + + preprocessor, postprocessor = make_classifier_processor(config, stats) + + # Simulate Accelerate: data already on GPU + device = torch.device("cuda:0") + observation = { + OBS_STATE: torch.randn(10).to(device), + OBS_IMAGE: torch.randn(3, 224, 224).to(device), + } + action = torch.randn(1).to(device) + transition = create_transition(observation, action) + batch = transition_to_batch(transition) + + # Process through preprocessor + + processed = preprocessor(batch) + + # Check that data stays on same GPU + assert processed[OBS_STATE].device == device + assert processed[OBS_IMAGE].device == device + assert processed[TransitionKey.ACTION.value].device == device + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") +def test_classifier_processor_multi_gpu(): + """Test Classifier processor with multi-GPU setup.""" + config = create_default_config() + config.device = "cuda:0" + stats = create_default_stats() + + preprocessor, postprocessor = make_classifier_processor(config, stats) + + # Simulate data on different GPU + device = torch.device("cuda:1") + observation = { + OBS_STATE: torch.randn(10).to(device), + OBS_IMAGE: torch.randn(3, 224, 224).to(device), + } + action = torch.randn(1).to(device) + transition = create_transition(observation, action) + batch = transition_to_batch(transition) + + # Process through preprocessor + + processed = preprocessor(batch) + + # Check that data stays on cuda:1 + assert processed[OBS_STATE].device == device + assert processed[OBS_IMAGE].device == device + assert processed[TransitionKey.ACTION.value].device == device + + def test_classifier_processor_without_stats(): """Test Classifier processor creation without dataset statistics.""" config = create_default_config() preprocessor, postprocessor = make_classifier_processor(config, dataset_stats=None) + # Should still create processors assert preprocessor is not None assert postprocessor is not None + # Process should still work observation = { OBS_STATE: torch.randn(10), OBS_IMAGE: torch.randn(3, 224, 224), @@ -132,12 +231,15 @@ def test_classifier_processor_save_and_load(): preprocessor, postprocessor = make_classifier_processor(config, stats) with tempfile.TemporaryDirectory() as tmpdir: + # Save preprocessor preprocessor.save_pretrained(tmpdir) + # Load preprocessor loaded_preprocessor = DataProcessorPipeline.from_pretrained( tmpdir, config_filename="classifier_preprocessor.json" ) + # Test that loaded processor works observation = { OBS_STATE: torch.randn(10), OBS_IMAGE: torch.randn(3, 224, 224), @@ -152,6 +254,43 @@ def test_classifier_processor_save_and_load(): assert processed[TransitionKey.ACTION.value].shape == (1,) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_classifier_processor_mixed_precision(): + """Test Classifier processor with mixed precision.""" + config = create_default_config() + config.device = "cuda" + stats = create_default_stats() + + preprocessor, postprocessor = make_classifier_processor(config, stats) + + # Replace DeviceProcessorStep with one that uses float16 + modified_steps = [] + for step in preprocessor.steps: + if isinstance(step, DeviceProcessorStep): + modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16")) + else: + modified_steps.append(step) + preprocessor.steps = modified_steps + + # Create test data + observation = { + OBS_STATE: torch.randn(10, dtype=torch.float32), + OBS_IMAGE: torch.randn(3, 224, 224, dtype=torch.float32), + } + action = torch.randn(1, dtype=torch.float32) + transition = create_transition(observation, action) + batch = transition_to_batch(transition) + + # Process through preprocessor + + processed = preprocessor(batch) + + # Check that data is converted to float16 + assert processed[OBS_STATE].dtype == torch.float16 + assert processed[OBS_IMAGE].dtype == torch.float16 + assert processed[TransitionKey.ACTION.value].dtype == torch.float16 + + def test_classifier_processor_batch_data(): """Test Classifier processor with batched data.""" config = create_default_config() @@ -159,6 +298,7 @@ def test_classifier_processor_batch_data(): preprocessor, postprocessor = make_classifier_processor(config, stats) + # Test with batched data batch_size = 16 observation = { OBS_STATE: torch.randn(batch_size, 10), @@ -168,8 +308,11 @@ def test_classifier_processor_batch_data(): transition = create_transition(observation, action) batch = transition_to_batch(transition) + # Process through preprocessor + processed = preprocessor(batch) + # Check that batch dimension is preserved assert processed[OBS_STATE].shape == (batch_size, 10) assert processed[OBS_IMAGE].shape == (batch_size, 3, 224, 224) assert processed[TransitionKey.ACTION.value].shape == (batch_size, 1) @@ -182,11 +325,14 @@ def test_classifier_processor_postprocessor_identity(): preprocessor, postprocessor = make_classifier_processor(config, stats) + # Create test data for postprocessor reward = torch.tensor([[0.8], [0.3], [0.9]]) transition = create_transition(action=reward) _ = transition_to_batch(transition) + # Process through postprocessor processed = postprocessor(reward) + # IdentityProcessor should leave values unchanged (except device) assert torch.allclose(processed.cpu(), reward.cpu()) assert processed.device.type == "cpu"