feat: expose UR policy config in Gradio

app.py

@@ -21,7 +21,7 @@ except ImportError:  # Local and dedicated-GPU environments omit this helper.
 
 from artifacts import resolve_checkpoint_path, resolve_model_id
 from inference import ACTION_LABELS, run_prediction
-from model_loader import MODEL_MANAGER
+from model_loader import DEFAULT_POLICY_CONFIG, MODEL_MANAGER, POLICY_CONFIGS
 
 
 def _gradio_integer(value, name: str) -> int:
@@ -39,6 +39,7 @@ def _gradio_integer(value, name: str) -> int:
 def predict_ui(
     model_id,
     checkpoint_path,
+    config_name,
     fixed_image,
     wrist_image,
     instruction,
@@ -53,7 +54,7 @@ def predict_ui(
 ):
     try:
         trial = _gradio_integer(trial_index, "trial index")
-        policy = MODEL_MANAGER.get(model_id, checkpoint_path)
+        policy = MODEL_MANAGER.get(model_id, checkpoint_path, config_name)
         result = run_prediction(
             policy,
             fixed_image,
@@ -64,7 +65,12 @@ def predict_ui(
             model_id,
             checkpoint_path,
         )
-        return result.actions, result.json_path, result.status
+        uses_discrete_state = config_name == "pi05_ur_demo_state"
+        status = (
+            f"{result.status} Config={config_name} "
+            f"(discrete_state_input={uses_discrete_state})."
+        )
+        return result.actions, result.json_path, status
     except Exception as exc:
         gc.collect()
         try:
@@ -98,6 +104,11 @@ def build_demo():
                 label="Checkpoint path",
                 placeholder="checkpoints/30000",
             )
+            config_name = gr.Dropdown(
+                choices=list(POLICY_CONFIGS),
+                value=DEFAULT_POLICY_CONFIG,
+                label="Policy config",
+            )
         with gr.Row():
             fixed_image = gr.Image(type="pil", label="Fixed camera")
             wrist_image = gr.Image(type="pil", label="Wrist camera")
@@ -106,7 +117,11 @@ def build_demo():
             placeholder="e.g. pick up the object and place it in the tray",
             lines=2,
         )
-        gr.Markdown("### Current state — metres/radians, followed by gripper state")
+        gr.Markdown(
+            "### Current state — metres/radians, followed by gripper state\n"
+            "These values are discrete state conditioning only when "
+            "`pi05_ur_demo_state` is selected."
+        )
         with gr.Row():
             tcp_x = gr.Number(value=0.0, label="TCP x")
             tcp_y = gr.Number(value=0.0, label="TCP y")
@@ -128,6 +143,7 @@ def build_demo():
             inputs=[
                 model_id,
                 checkpoint_path,
+                config_name,
                 fixed_image,
                 wrist_image,
                 instruction,

tests/test_app.py

@@ -11,21 +11,33 @@ class AppTests(unittest.TestCase):
             "Result", (),
             {"actions": "table", "json_path": "/tmp/result.json", "status": "done"},
         )()
-        with mock.patch.object(app.MODEL_MANAGER, "get", return_value=object()), mock.patch.object(
+        with mock.patch.object(app.MODEL_MANAGER, "get", return_value=object()) as get_model, mock.patch.object(
             app, "run_prediction", return_value=result
         ):
             actual = app.predict_ui(
-                "owner/model", "checkpoint", object(), object(), "task",
+                "owner/model", "checkpoint", "pi05_ur_demo_no_state",
+                object(), object(), "task",
                 1, 2, 3, 4, 5, 6, 0, 0,
             )
-        self.assertEqual(actual, ("table", "/tmp/result.json", "done"))
+        self.assertEqual(
+            actual,
+            (
+                "table",
+                "/tmp/result.json",
+                "done Config=pi05_ur_demo_no_state (discrete_state_input=False).",
+            ),
+        )
+        get_model.assert_called_once_with(
+            "owner/model", "checkpoint", "pi05_ur_demo_no_state"
+        )
 
     def test_predict_ui_turns_exceptions_into_status(self):
         import app
 
         with mock.patch.object(app.MODEL_MANAGER, "get", side_effect=RuntimeError("load failed")):
             table, output_file, status = app.predict_ui(
-                "model", "checkpoint", object(), object(), "task",
+                "model", "checkpoint", "pi05_ur_demo_state",
+                object(), object(), "task",
                 0, 0, 0, 0, 0, 0, 0, 0,
             )
         self.assertIsNone(table)
@@ -40,6 +52,8 @@ class AppTests(unittest.TestCase):
         ):
             self.assertIn(label, source)
         self.assertIn("default_concurrency_limit=1", source)
+        self.assertIn('label="Policy config"', source)
+        self.assertIn("value=DEFAULT_POLICY_CONFIG", source)
 
 
 if __name__ == "__main__":