Remove previous pi0 and rename pi0_openpi and pi05_openpi

tests/policies/pi0_pi05/test_pi05_openpi.py

@@ -13,7 +13,7 @@ pytestmark = pytest.mark.skipif(
     reason="This test requires local OpenPI installation and is not meant for CI",
 )
 
-from lerobot.policies.pi05_openpi import PI05OpenPIConfig, PI05OpenPIPolicy  # noqa: E402
+from lerobot.policies.pi05 import PI05Config, PI05Policy  # noqa: E402
 from tests.utils import require_cuda  # noqa: E402
 
 
@@ -22,7 +22,7 @@ def test_pi05_model_architecture():
     """Test that pi05=True creates the correct model architecture."""
 
     # Create config
-    config = PI05OpenPIConfig(
+    config = PI05Config(
         max_action_dim=7,
         max_state_dim=14,
         dtype="float32",
@@ -73,7 +73,7 @@ def test_pi05_model_architecture():
     }
 
     # Instantiate policy
-    policy = PI05OpenPIPolicy(config, dataset_stats)
+    policy = PI05Policy(config, dataset_stats)
 
     # Verify pi05 model components exist
     # Check that time_mlp layers exist (for AdaRMS conditioning)
@@ -104,7 +104,7 @@ def test_pi05_forward_pass():
     """Test forward pass with"""
 
     # Create config
-    config = PI05OpenPIConfig(
+    config = PI05Config(
         max_action_dim=7,
         max_state_dim=14,
         dtype="float32",
@@ -150,7 +150,7 @@ def test_pi05_forward_pass():
     }
 
     # Instantiate policy
-    policy = PI05OpenPIPolicy(config, dataset_stats)
+    policy = PI05Policy(config, dataset_stats)
 
     # Create test batch
     batch_size = 2

tests/policies/pi0_pi05/test_pi0_openpi.py

@@ -14,7 +14,7 @@ pytestmark = pytest.mark.skipif(
 )
 
 from lerobot.policies.factory import make_policy_config  # noqa: E402
-from lerobot.policies.pi0_openpi import PI0OpenPIConfig, PI0OpenPIPolicy  # noqa: E402
+from lerobot.policies.pi0 import PI0OpenPIConfig, PI0OpenPIPolicy  # noqa: E402
 from tests.utils import require_cuda  # noqa: E402
 
 
@@ -96,7 +96,7 @@ def test_config_creation():
     """Test policy config creation through factory."""
     try:
         config = make_policy_config(
-            policy_type="pi0_openpi",
+            policy_type="pi0",
             max_action_dim=7,
             max_state_dim=14,
         )

tests/policies/pi0_pi05/test_pi0_original_vs_lerobot.py

@@ -21,7 +21,7 @@ from openpi.models_pytorch import preprocessing_pytorch as openpi_preprocessing
 from openpi.models_pytorch.pi0_pytorch import PI0Pytorch  # noqa: E402
 from transformers import AutoTokenizer  # noqa: E402
 
-from lerobot.policies.pi0_openpi import PI0OpenPIConfig, PI0OpenPIPolicy  # noqa: E402
+from lerobot.policies.pi0 import PI0Config, PI0Policy  # noqa: E402
 
 DUMMY_ACTION_DIM = 32
 DUMMY_STATE_DIM = 32
@@ -68,9 +68,7 @@ class PI0BaseOriginalConfig:
 def instantiate_lerobot_pi0(from_pretrained: bool = False):
     if from_pretrained:
         # Load the policy first
-        policy = PI0OpenPIPolicy.from_pretrained(
-            pretrained_name_or_path="pepijn223/pi0_base_fp32", strict=True
-        )
+        policy = PI0Policy.from_pretrained(pretrained_name_or_path="pepijn223/pi0_base_fp32", strict=True)
         # Then reinitialize the normalization with proper stats
         from lerobot.policies.normalize import Normalize, Unnormalize
 
@@ -84,10 +82,8 @@ def instantiate_lerobot_pi0(from_pretrained: bool = False):
             policy.config.output_features, policy.config.normalization_mapping, DUMMY_DATASET_STATS
         )
     else:
-        config = PI0OpenPIConfig(
-            max_action_dim=DUMMY_ACTION_DIM, max_state_dim=DUMMY_STATE_DIM, dtype="float32"
-        )
-        policy = PI0OpenPIPolicy(config, DUMMY_DATASET_STATS)
+        config = PI0Config(max_action_dim=DUMMY_ACTION_DIM, max_state_dim=DUMMY_STATE_DIM, dtype="float32")
+        policy = PI0Policy(config, DUMMY_DATASET_STATS)
     policy.to(DEVICE)
     return policy
 

tests/policies/pi0_pi05/test_pi0_pi05_hub.py

@@ -18,8 +18,8 @@ pytestmark = pytest.mark.skipif(
     reason="This test requires HuggingFace authentication and is not meant for CI",
 )
 
-from lerobot.policies.pi0_openpi import PI0OpenPIPolicy  # noqa: E402
-from lerobot.policies.pi05_openpi.modeling_pi05openpi import PI05OpenPIPolicy  # noqa: E402
+from lerobot.policies.pi0 import PI0Policy  # noqa: E402
+from lerobot.policies.pi05.modeling_pi05openpi import PI05Policy  # noqa: E402
 
 
 def create_dummy_stats(config):
@@ -48,13 +48,13 @@ def create_dummy_stats(config):
 # Test data for all 6 base models
 MODEL_TEST_PARAMS = [
     # PI0 models
-    ("pepijn223/pi0_base_fp32", "PI0", PI0OpenPIPolicy),
-    ("pepijn223/pi0_droid_fp32", "PI0", PI0OpenPIPolicy),
-    ("pepijn223/pi0_libero_fp32", "PI0", PI0OpenPIPolicy),
+    ("pepijn223/pi0_base_fp32", "PI0", PI0Policy),
+    ("pepijn223/pi0_droid_fp32", "PI0", PI0Policy),
+    ("pepijn223/pi0_libero_fp32", "PI0", PI0Policy),
     # PI0.5 models
-    ("pepijn223/pi05_base_fp32", "PI0.5", PI05OpenPIPolicy),
-    ("pepijn223/pi05_droid_fp32", "PI0.5", PI05OpenPIPolicy),
-    ("pepijn223/pi05_libero_fp32", "PI0.5", PI05OpenPIPolicy),
+    ("pepijn223/pi05_base_fp32", "PI0.5", PI05Policy),
+    ("pepijn223/pi05_droid_fp32", "PI0.5", PI05Policy),
+    ("pepijn223/pi05_libero_fp32", "PI0.5", PI05Policy),
 ]
 
 
@@ -65,7 +65,7 @@ def test_all_base_models_hub_loading(model_id, model_type, policy_class):
     Args:
         model_id: HuggingFace model ID (e.g., "pepijn223/pi0_base_fp32")
         model_type: Model type ("PI0" or "PI0.5")
-        policy_class: Policy class to use (PI0OpenPIPolicy or PI05OpenPIPolicy)
+        policy_class: Policy class to use (PI0Policy or PI05Policy)
     """
     print(f"\n{'=' * 80}")
     print(f"Testing {model_type} model: {model_id}")