initial commit

2026-05-21 11:39:50 +00:00 · 2025-09-10 19:44:41 +02:00
parent d602e8169c
commit bd5c264c49
12 changed files with 4728 additions and 0 deletions
@@ -0,0 +1,109 @@
+#!/usr/bin/env python
+
+"""Test script to verify PI0OpenPI policy integration with LeRobot."""
+
+import torch
+
+from lerobot.policies.factory import make_policy_config
+from lerobot.policies.pi0_openpi import PI0OpenPIConfig, PI0OpenPIPolicy
+
+
+def test_policy_instantiation():
+    """Test basic policy instantiation."""
+    print("Testing PI0OpenPI policy instantiation...")
+
+    # Create config
+    config = PI0OpenPIConfig(action_dim=7, state_dim=14, device="cpu", dtype="float32")
+
+    # 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),
+        },
+    }
+
+    # Instantiate policy
+    policy = PI0OpenPIPolicy(config, dataset_stats)
+    print(f"Policy created successfully: {policy.name}")
+
+    # Test forward pass with dummy data
+    batch_size = 1
+    device = policy.device if hasattr(policy, "device") else "cpu"
+    batch = {
+        "observation.state": torch.randn(batch_size, 14, dtype=torch.float32, device=device),
+        "action": torch.randn(batch_size, config.action_horizon, 7, dtype=torch.float32, device=device),
+        "observation.images.base_0_rgb": torch.rand(
+            batch_size, 3, 224, 224, dtype=torch.float32, device=device
+        ),  # Use rand for [0,1] range
+        "task": ["Pick up the object"] * batch_size,
+    }
+
+    print("\nTesting forward pass...")
+    try:
+        loss, loss_dict = policy.forward(batch)
+        print(f"✓ Forward pass successful. Loss: {loss_dict['loss']:.4f}")
+    except Exception as e:
+        print(f"✗ Forward pass failed: {e}")
+        return False
+
+    print("\nTesting action prediction...")
+    try:
+        with torch.no_grad():
+            action = policy.select_action(batch)
+        print(f"✓ Action prediction successful. Action shape: {action.shape}")
+    except Exception as e:
+        print(f"✗ Action prediction failed: {e}")
+        return False
+
+    return True
+
+
+def test_config_creation():
+    """Test policy config creation through factory."""
+    print("\nTesting config creation through factory...")
+
+    try:
+        config = make_policy_config(
+            policy_type="pi0_openpi",
+            action_dim=7,
+            state_dim=14,
+        )
+        print("✓ Config created successfully through factory")
+        print(f"  Config type: {type(config).__name__}")
+        print(f"  PaliGemma variant: {config.paligemma_variant}")
+        print(f"  Action expert variant: {config.action_expert_variant}")
+        return True
+    except Exception as e:
+        print(f"✗ Config creation failed: {e}")
+        return False
+
+
+def main():
+    """Run all tests."""
+    print("=" * 60)
+    print("PI0OpenPI Policy Integration Test")
+    print("=" * 60)
+
+    # Test config creation
+    config_test = test_config_creation()
+
+    print("\n" + "-" * 60)
+
+    # Test policy instantiation
+    policy_test = test_policy_instantiation()
+
+    print("\n" + "=" * 60)
+    if config_test and policy_test:
+        print("✓ All tests passed!")
+    else:
+        print("✗ Some tests failed.")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()