Move test to specific folder

tests/policies/pi0_pi05/test_pi05_openpi.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python
+
+"""Test script to verify PI0.5 (pi05) support in PI0OpenPI policy, only meant to be run locally!"""
+
+import os
+
+import pytest
+import torch
+
+# Skip this entire module in CI
+pytestmark = pytest.mark.skipif(
+    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
+    reason="This test requires local OpenPI installation and is not meant for CI",
+)
+
+from lerobot.policies.pi05_openpi import PI05OpenPIConfig, PI05OpenPIPolicy  # noqa: E402
+from tests.utils import require_cuda  # noqa: E402
+
+
+@require_cuda
+def test_pi05_model_architecture():
+    """Test that pi05=True creates the correct model architecture."""
+
+    # Create config
+    config = PI05OpenPIConfig(
+        max_action_dim=7,
+        max_state_dim=14,
+        dtype="float32",
+    )
+
+    # Set up input_features and output_features in the config
+    from lerobot.configs.types import FeatureType, PolicyFeature
+
+    config.input_features = {
+        "observation.state": PolicyFeature(
+            type=FeatureType.STATE,
+            shape=(14,),
+        ),
+        "observation.images.base_0_rgb": PolicyFeature(
+            type=FeatureType.VISUAL,
+            shape=(3, 224, 224),
+        ),
+    }
+
+    config.output_features = {
+        "action": PolicyFeature(
+            type=FeatureType.ACTION,
+            shape=(7,),
+        ),
+    }
+
+    assert config.tokenizer_max_length == 200, (
+        f"Expected tokenizer_max_length=200 for pi05, got {config.tokenizer_max_length}"
+    )
+    assert config.discrete_state_input == True, (  # noqa: E712
+        f"Expected discrete_state_input=True for pi05, got {config.discrete_state_input}"
+    )
+
+    # Create dummy dataset stats
+    dataset_stats = {
+        "observation.state": {
+            "mean": torch.zeros(14),
+            "std": torch.ones(14),
+        },
+        "action": {
+            "mean": torch.zeros(7),
+            "std": torch.ones(7),
+        },
+        "observation.images.base_0_rgb": {
+            "mean": torch.zeros(3, 224, 224),
+            "std": torch.ones(3, 224, 224),
+        },
+    }
+
+    # Instantiate policy
+    policy = PI05OpenPIPolicy(config, dataset_stats)
+
+    # Verify pi05 model components exist
+    # Check that time_mlp layers exist (for AdaRMS conditioning)
+    assert hasattr(policy.model, "time_mlp_in"), "Missing time_mlp_in layer for pi05"
+    assert hasattr(policy.model, "time_mlp_out"), "Missing time_mlp_out layer for pi05"
+
+    # Check that action_time_mlp layers don't exist (pi0 only)
+    assert not hasattr(policy.model, "action_time_mlp_in"), "action_time_mlp_in should not exist in pi05 mode"
+    assert not hasattr(policy.model, "action_time_mlp_out"), (
+        "action_time_mlp_out should not exist in pi05 mode"
+    )
+
+    # Check that state_proj doesn't exist in pi05 mode
+    assert not hasattr(policy.model, "state_proj"), "state_proj should not exist in pi05 mode"
+
+    # Check AdaRMS configuration in the underlying model
+    adarms_config = policy.model.paligemma_with_expert.paligemma.config.text_config.use_adarms
+    assert adarms_config == False, f"PaliGemma should not use AdaRMS, got {adarms_config}"  # noqa: E712
+
+    adarms_expert_config = policy.model.paligemma_with_expert.gemma_expert.config.use_adarms
+    assert adarms_expert_config == True, (  # noqa: E712
+        f"Action expert should use AdaRMS in pi05, got {adarms_expert_config}"
+    )
+
+
+@require_cuda
+def test_pi05_forward_pass():
+    """Test forward pass with"""
+
+    # Create config
+    config = PI05OpenPIConfig(
+        max_action_dim=7,
+        max_state_dim=14,
+        dtype="float32",
+        chunk_size=16,  # Shorter chunk_size for testing
+        n_action_steps=16,  # Shorter action steps for testing
+    )
+
+    # Set up input_features and output_features in the config
+    from lerobot.configs.types import FeatureType, PolicyFeature
+
+    config.input_features = {
+        "observation.state": PolicyFeature(
+            type=FeatureType.STATE,
+            shape=(14,),
+        ),
+        "observation.images.base_0_rgb": PolicyFeature(
+            type=FeatureType.VISUAL,
+            shape=(3, 224, 224),
+        ),
+    }
+
+    config.output_features = {
+        "action": PolicyFeature(
+            type=FeatureType.ACTION,
+            shape=(7,),
+        ),
+    }
+
+    # Create dummy dataset stats
+    dataset_stats = {
+        "observation.state": {
+            "mean": torch.zeros(14),
+            "std": torch.ones(14),
+        },
+        "action": {
+            "mean": torch.zeros(7),
+            "std": torch.ones(7),
+        },
+        "observation.images.base_0_rgb": {
+            "mean": torch.zeros(3, 224, 224),
+            "std": torch.ones(3, 224, 224),
+        },
+    }
+
+    # Instantiate policy
+    policy = PI05OpenPIPolicy(config, dataset_stats)
+
+    # Create test batch
+    batch_size = 2
+    device = next(policy.parameters()).device
+    batch = {
+        "observation.state": torch.randn(batch_size, 14, dtype=torch.float32, device=device),
+        "action": torch.randn(batch_size, config.chunk_size, 7, dtype=torch.float32, device=device),
+        "observation.images.base_0_rgb": torch.rand(
+            batch_size, 3, 224, 224, dtype=torch.float32, device=device
+        ),
+        "task": ["Pick up the object"] * batch_size,
+    }
+
+    # Test forward pass
+    try:
+        loss, loss_dict = policy.forward(batch)
+        print(f"Forward pass successful. Loss: {loss_dict['loss']:.4f}")
+        assert not torch.isnan(loss), "Loss is NaN"
+        assert loss.item() >= 0, "Loss should be non-negative"
+    except Exception as e:
+        print(f"Forward pass failed: {e}")
+        raise
+
+    # Test action prediction
+    try:
+        with torch.no_grad():
+            action = policy.select_action(batch)
+        print(f"Action prediction successful. Action shape: {action.shape}")
+        # When batch_size > 1, select_action returns (batch_size, action_dim)
+        assert action.shape == (batch_size, 7), f"Expected action shape ({batch_size}, 7), got {action.shape}"
+        assert not torch.isnan(action).any(), "Action contains NaN values"
+    except Exception as e:
+        print(f"Action prediction failed: {e}")
+        raise