feat: select UR policy configuration

model_loader.py

@@ -14,6 +14,16 @@ from artifacts import (
     normalize_model_id,
 )
 
+POLICY_CONFIGS = ("pi05_ur_demo_no_state", "pi05_ur_demo_state")
+DEFAULT_POLICY_CONFIG = POLICY_CONFIGS[0]
+
+
+def normalize_policy_config(value: str) -> str:
+    if value not in POLICY_CONFIGS:
+        choices = ", ".join(POLICY_CONFIGS)
+        raise ValueError(f"unsupported policy config {value!r}; choose one of: {choices}")
+    return value
+
 
 class ModelUnavailableError(RuntimeError):
     """Raised when the requested policy cannot be initialized."""
@@ -31,30 +41,31 @@ def _release_gpu_memory() -> None:
 
 
 class ModelManager:
-    def __init__(self, loader: Callable[[str, str], object] | None = None):
+    def __init__(self, loader: Callable[[str, str, str], object] | None = None):
         self._loader = loader or self._load_default
         self._lock = threading.Lock()
         self._value = None
-        self._active_key: tuple[str, str] | None = None
+        self._active_key: tuple[str, str, str] | None = None
         self._error: str | None = None
 
     @property
-    def active_key(self) -> tuple[str, str] | None:
+    def active_key(self) -> tuple[str, str, str] | None:
         return self._active_key
 
     @property
     def health_message(self) -> str:
         if self._value is not None and self._active_key is not None:
-            model_id, checkpoint_path = self._active_key
-            return f"Model ready: {model_id}/{checkpoint_path}."
+            model_id, checkpoint_path, config_name = self._active_key
+            return f"Model ready: {model_id}/{checkpoint_path} ({config_name})."
         if self._error:
             return f"Model unavailable: {self._error}"
         return "Model has not been loaded yet."
 
-    def get(self, model_id: str, checkpoint_path: str):
+    def get(self, model_id: str, checkpoint_path: str, config_name: str):
         key = (
             normalize_model_id(model_id),
             normalize_checkpoint_path(checkpoint_path),
+            normalize_policy_config(config_name),
         )
         if self._value is not None and self._active_key == key:
             return self._value
@@ -71,14 +82,14 @@ class ModelManager:
             except Exception as exc:
                 _release_gpu_memory()
                 detail = str(exc) or exc.__class__.__name__
-                self._error = f"{key[0]}/{key[1]}: {detail}"
+                self._error = f"{key[0]}/{key[1]} ({key[2]}): {detail}"
                 raise ModelUnavailableError(self._error) from exc
             self._value = value
             self._active_key = key
             return value
 
     @staticmethod
-    def _load_default(model_id: str, checkpoint_path: str):
+    def _load_default(model_id: str, checkpoint_path: str, config_name: str):
         import torch
 
         if not torch.cuda.is_available():
@@ -90,7 +101,7 @@ class ModelManager:
         from openpi.training import config as openpi_config
 
         paths = download_checkpoint(model_id, checkpoint_path)
-        config = openpi_config.get_config("pi05_ur_demo_state")
+        config = openpi_config.get_config(config_name)
         return policy_config.create_trained_policy(
             config,
             paths.checkpoint,

tests/test_model_loader.py

@@ -7,13 +7,16 @@ class ModelManagerTests(unittest.TestCase):
 
         calls = []
         manager = ModelManager(
-            loader=lambda model, path: calls.append((model, path)) or object()
+            loader=lambda model, path, config: calls.append((model, path, config)) or object()
         )
-        first = manager.get("owner/model", "a")
-        self.assertIs(manager.get("owner/model", "a"), first)
-        second = manager.get("owner/model", "b")
+        first = manager.get("owner/model", "a", "pi05_ur_demo_no_state")
+        self.assertIs(manager.get("owner/model", "a", "pi05_ur_demo_no_state"), first)
+        second = manager.get("owner/model", "a", "pi05_ur_demo_state")
         self.assertIsNot(second, first)
-        self.assertEqual(calls, [("owner/model", "a"), ("owner/model", "b")])
+        self.assertEqual(calls, [
+            ("owner/model", "a", "pi05_ur_demo_no_state"),
+            ("owner/model", "a", "pi05_ur_demo_state"),
+        ])
 
     def test_failed_load_is_reported_and_can_retry(self):
         from model_loader import ModelManager, ModelUnavailableError
@@ -28,7 +31,7 @@ class ModelManagerTests(unittest.TestCase):
         manager = ModelManager(loader=loader)
         for _ in range(2):
             with self.assertRaisesRegex(ModelUnavailableError, "bad checkpoint"):
-                manager.get("owner/model", "a")
+                manager.get("owner/model", "a", "pi05_ur_demo_no_state")
         self.assertEqual(attempts, 2)
         self.assertIn("bad checkpoint", manager.health_message)
 
@@ -37,7 +40,14 @@ class ModelManagerTests(unittest.TestCase):
 
         manager = ModelManager(loader=lambda *_: self.fail("loader was called"))
         with self.assertRaises(ValueError):
-            manager.get("", "checkpoint")
+            manager.get("", "checkpoint", "pi05_ur_demo_no_state")
+
+    def test_unsupported_policy_config_is_rejected_before_loader(self):
+        from model_loader import ModelManager
+
+        manager = ModelManager(loader=lambda *_: self.fail("loader was called"))
+        with self.assertRaisesRegex(ValueError, "unsupported policy config"):
+            manager.get("owner/model", "checkpoint", "pi05_libero")
 
 
 if __name__ == "__main__":