test(rewards): restore missing CUDA and mixed precision classifier processor tests

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: {},
+        OBS_IMAGE: {},  # No normalization for images
-        "reward": {},
+        "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[0], NormalizerProcessorStep)  # For input features
-    assert isinstance(preprocessor.steps[1], NormalizerProcessorStep)
+    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"