docs: plan pi05 UR Gradio Space

docs/superpowers/plans/2026-07-22-pi05-ur-gradio.md

@@ -0,0 +1,494 @@
+# π₀.₅ UR Hugging Face Gradio Implementation Plan
+
+> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
+
+**Goal:** Build a self-contained Hugging Face Gradio Space that loads a state-conditioned `pi05_ur_demo_state` checkpoint and predicts ten 7-dimensional UR actions from two RGB images, a task instruction, and current TCP/gripper state.
+
+**Architecture:** Vendor the proven local OpenPI runtime into the Space, while keeping Space-specific artifact resolution, model lifecycle, inference validation, and Gradio wiring in focused modules. Load one CUDA policy lazily, cache it by Hugging Face repository/checkpoint identity, and keep all CPU tests independent of model downloads.
+
+**Tech Stack:** Python 3.11, Gradio 6.20.0, Hugging Face Hub, OpenPI/JAX/Flax, NumPy, pandas, Pillow, pytest.
+
+## Global Constraints
+
+- Support only `pi05_ur_demo_state`; do not expose LIBERO, DROID, no-state inference, websocket serving, or robot control.
+- Input state order is exactly `x, y, z, roll, pitch, yaw, gripper`.
+- Output shape is exactly `(10, 7)` with labels `dx`, `dy`, `dz`, `droll`, `dpitch`, `dyaw`, `gripper`.
+- Use Python 3.11 and Gradio 6.20.0 in Hugging Face Space metadata.
+- Model loading is lazy, CUDA-only, thread-safe, and cached for one `(model_id, checkpoint_path)` key.
+- CPU tests must not download checkpoints or require CUDA.
+
+---
+
+### Task 1: Space metadata and vendored OpenPI runtime
+
+**Files:**
+- Modify: `README.md`
+- Create: `requirements.txt`
+- Create: `openpi_runtime/openpi/**` by copying `../VLA/openpi/src/openpi/**`
+- Create: `openpi_runtime/openpi_client/**` by copying `../VLA/openpi/packages/openpi-client/src/openpi_client/**`
+- Test: `tests/test_space_config.py`
+
+**Interfaces:**
+- Consumes: local source trees `../VLA/openpi/src/openpi` and `../VLA/openpi/packages/openpi-client/src/openpi_client`.
+- Produces: importable `openpi.training.config`, `openpi.policies.policy_config`, and `openpi_client`; Space dependency and metadata contract.
+
+- [ ] **Step 1: Write the failing metadata test**
+
+```python
+from pathlib import Path
+
+
+ROOT = Path(__file__).parents[1]
+
+
+def test_space_metadata_and_runtime_files():
+    readme = (ROOT / "README.md").read_text()
+    requirements = (ROOT / "requirements.txt").read_text()
+    assert "sdk_version: 6.20.0" in readme
+    assert "python_version: '3.11'" in readme
+    assert "app_file: app.py" in readme
+    assert "gradio==6.20.0" in requirements
+    assert (ROOT / "openpi_runtime/openpi/policies/ur_policy.py").is_file()
+    assert (ROOT / "openpi_runtime/openpi_client/base_policy.py").is_file()
+
+
+def test_vendored_config_contains_state_conditioned_ur_policy():
+    source = (ROOT / "openpi_runtime/openpi/training/config.py").read_text()
+    assert 'name="pi05_ur_demo_state"' in source
+    assert "LeRobotURDataConfig" in source
+```
+
+- [ ] **Step 2: Run the test to verify it fails**
+
+Run: `pytest tests/test_space_config.py -v`
+
+Expected: FAIL because `requirements.txt` and `openpi_runtime` do not exist and the README still selects Python 3.12.
+
+- [ ] **Step 3: Vendor runtime and add pinned deployment dependencies**
+
+Run these bulk-copy commands from `pi0.5`:
+
+```bash
+mkdir -p openpi_runtime
+rsync -a --exclude='*_test.py' --exclude='__pycache__' ../VLA/openpi/src/openpi/ openpi_runtime/openpi/
+rsync -a --exclude='*_test.py' --exclude='__pycache__' ../VLA/openpi/packages/openpi-client/src/openpi_client/ openpi_runtime/openpi_client/
+```
+
+Change README metadata to `python_version: '3.11'`. Create `requirements.txt` from the OpenPI runtime pins, including:
+
+```text
+gradio==6.20.0
+huggingface-hub==0.36.2
+pandas==2.2.3
+pillow==11.1.0
+numpy==1.26.4
+augmax>=0.3.4
+dm-tree>=0.1.8
+einops>=0.8.0
+equinox>=0.11.8
+flax==0.10.2
+fsspec[gcs]>=2024.6.0
+jax[cuda12]==0.5.3
+jaxtyping==0.2.36
+ml-collections==1.0.0
+numpydantic>=1.6.6
+opencv-python-headless>=4.10.0.84
+orbax-checkpoint==0.11.13
+sentencepiece>=0.2.0
+torch==2.7.1
+typing-extensions>=4.12.2
+tyro>=0.9.5
+wandb>=0.19.1
+filelock>=3.16.1
+beartype==0.19.0
+chex==0.1.90
+treescope>=0.1.7
+transformers==4.53.2
+rich>=14.0.0
+polars>=1.30.0
+```
+
+- [ ] **Step 4: Run metadata tests and import smoke test**
+
+Run: `pytest tests/test_space_config.py -v && PYTHONPATH=openpi_runtime python -c "from openpi.training import config; assert config.get_config('pi05_ur_demo_state').model.action_horizon == 10"`
+
+Expected: tests PASS and the import command exits zero.
+
+- [ ] **Step 5: Commit**
+
+```bash
+git add README.md requirements.txt openpi_runtime tests/test_space_config.py
+git commit -m "build: add OpenPI runtime for UR Space"
+```
+
+---
+
+### Task 2: Artifact identity, download, and validation
+
+**Files:**
+- Create: `artifacts.py`
+- Test: `tests/test_artifacts.py`
+
+**Interfaces:**
+- Consumes: environment variables `PI05_MODEL_ID` and `PI05_CHECKPOINT_PATH`; `huggingface_hub.snapshot_download`.
+- Produces: `ArtifactPaths(checkpoint: Path)`, `resolve_model_id() -> str`, `resolve_checkpoint_path() -> str`, `normalize_model_id(str) -> str`, `normalize_checkpoint_path(str) -> str`, and `download_checkpoint(str, str) -> ArtifactPaths`.
+
+- [ ] **Step 1: Write failing validation and download tests**
+
+```python
+from pathlib import Path
+import pytest
+import artifacts
+
+
+def test_checkpoint_path_rejects_absolute_and_parent_paths():
+    for value in ("/tmp/model", "../model", "a/../../model", ""):
+        with pytest.raises(ValueError):
+            artifacts.normalize_checkpoint_path(value)
+
+
+def test_download_checkpoint_validates_required_layout(tmp_path, monkeypatch):
+    checkpoint = tmp_path / "checkpoints/30000"
+    (checkpoint / "params").mkdir(parents=True)
+    (checkpoint / "assets/ur_demo").mkdir(parents=True)
+    (checkpoint / "assets/ur_demo/norm_stats.json").write_text("{}")
+    monkeypatch.setattr(artifacts, "snapshot_download", lambda **_: str(tmp_path))
+    paths = artifacts.download_checkpoint("owner/model", "checkpoints/30000")
+    assert paths.checkpoint == checkpoint
+
+
+def test_download_checkpoint_reports_missing_files(tmp_path, monkeypatch):
+    monkeypatch.setattr(artifacts, "snapshot_download", lambda **_: str(tmp_path))
+    with pytest.raises(FileNotFoundError, match="params|model.safetensors"):
+        artifacts.download_checkpoint("owner/model", "checkpoint")
+```
+
+- [ ] **Step 2: Run the tests to verify they fail**
+
+Run: `pytest tests/test_artifacts.py -v`
+
+Expected: FAIL with `ModuleNotFoundError: No module named 'artifacts'`.
+
+- [ ] **Step 3: Implement artifact helpers**
+
+Implement strict non-empty model IDs, POSIX-relative checkpoint paths without `..`, environment defaults, snapshot download, and checks for either `params/` or `model.safetensors` plus `assets/ur_demo/norm_stats.json`:
+
+```python
+@dataclass(frozen=True)
+class ArtifactPaths:
+    checkpoint: Path
+
+
+def download_checkpoint(model_id: str, checkpoint_path: str) -> ArtifactPaths:
+    model_id = normalize_model_id(model_id)
+    relative = normalize_checkpoint_path(checkpoint_path)
+    root = Path(snapshot_download(repo_id=model_id))
+    checkpoint = root.joinpath(*PurePosixPath(relative).parts)
+    if not (checkpoint / "params").exists() and not (checkpoint / "model.safetensors").is_file():
+        raise FileNotFoundError(f"checkpoint has neither params/ nor model.safetensors: {checkpoint}")
+    stats = checkpoint / "assets/ur_demo/norm_stats.json"
+    if not stats.is_file():
+        raise FileNotFoundError(f"UR normalization statistics not found: {stats}")
+    return ArtifactPaths(checkpoint)
+```
+
+- [ ] **Step 4: Run artifact tests**
+
+Run: `pytest tests/test_artifacts.py -v`
+
+Expected: PASS.
+
+- [ ] **Step 5: Commit**
+
+```bash
+git add artifacts.py tests/test_artifacts.py
+git commit -m "feat: resolve pi05 checkpoint artifacts"
+```
+
+---
+
+### Task 3: UR inference adapter and result serialization
+
+**Files:**
+- Create: `inference.py`
+- Test: `tests/test_inference.py`
+
+**Interfaces:**
+- Consumes: an object exposing `infer(observation: dict) -> dict`; two PIL/array images; instruction; seven numeric state values; trial index; artifact identity.
+- Produces: `ACTION_LABELS`, `ACTION_HORIZON`, `PredictionResult(actions: pandas.DataFrame, json_path: str, status: str)`, `set_seed(int)`, and `run_prediction(...) -> PredictionResult`.
+
+- [ ] **Step 1: Write failing mapping, validation, shape, and JSON tests**
+
+```python
+import json
+import numpy as np
+from PIL import Image
+import pytest
+from inference import ACTION_LABELS, run_prediction
+
+
+class FakePolicy:
+    def __init__(self, actions=None):
+        self.observation = None
+        self.actions = np.ones((10, 7)) if actions is None else actions
+
+    def infer(self, observation):
+        self.observation = observation
+        return {"actions": self.actions}
+
+
+def test_prediction_maps_ur_observation_and_writes_json():
+    policy = FakePolicy()
+    image = Image.new("RGB", (8, 6), "red")
+    result = run_prediction(policy, image, image, "pick up", [1, 2, 3, 4, 5, 6, 0], 2, "owner/model", "checkpoint")
+    assert tuple(result.actions.columns) == ACTION_LABELS
+    assert policy.observation["observation/state"].tolist() == [1, 2, 3, 4, 5, 6, 0]
+    assert policy.observation["observation/image"].dtype == np.uint8
+    assert policy.observation["prompt"] == "pick up"
+    assert json.loads(open(result.json_path).read())["seed"] == 44
+
+
+def test_prediction_rejects_invalid_state_and_action_shape():
+    image = Image.new("RGB", (8, 6))
+    with pytest.raises(ValueError, match="seven"):
+        run_prediction(FakePolicy(), image, image, "task", [1, 2], 0, "m", "c")
+    with pytest.raises(RuntimeError, match=r"\(10, 7\)"):
+        run_prediction(FakePolicy(np.zeros((8, 7))), image, image, "task", [0] * 7, 0, "m", "c")
+```
+
+- [ ] **Step 2: Run tests to verify they fail**
+
+Run: `pytest tests/test_inference.py -v`
+
+Expected: FAIL with `ModuleNotFoundError: No module named 'inference'`.
+
+- [ ] **Step 3: Implement validated inference**
+
+Define `ACTION_LABELS = ("dx", "dy", "dz", "droll", "dpitch", "dyaw", "gripper")`, `ACTION_HORIZON = 10`, and `BASE_SEED = 42`. Validate both images, trimmed instruction, integer non-negative trial index, exactly seven finite float state values, and exact action shape. Convert images with `Image.fromarray(np.asarray(value)).convert("RGB")`, seed Python/NumPy/Torch/JAX-compatible randomness, invoke `policy.infer`, build a pandas table, and serialize inputs and actions to a temporary JSON file.
+
+- [ ] **Step 4: Run inference tests**
+
+Run: `pytest tests/test_inference.py -v`
+
+Expected: PASS.
+
+- [ ] **Step 5: Commit**
+
+```bash
+git add inference.py tests/test_inference.py
+git commit -m "feat: add validated UR action inference"
+```
+
+---
+
+### Task 4: Thread-safe lazy OpenPI model manager
+
+**Files:**
+- Create: `model_loader.py`
+- Test: `tests/test_model_loader.py`
+
+**Interfaces:**
+- Consumes: `artifacts.download_checkpoint`, vendored `_config.get_config("pi05_ur_demo_state")`, and `policy_config.create_trained_policy(..., pytorch_device="cuda")`.
+- Produces: `ModelManager.get(model_id: str, checkpoint_path: str)`, `ModelManager.health_message`, `ModelUnavailableError`, and singleton `MODEL_MANAGER`.
+
+- [ ] **Step 1: Write failing cache and failure tests**
+
+```python
+import pytest
+from model_loader import ModelManager, ModelUnavailableError
+
+
+def test_manager_caches_same_key_and_replaces_changed_key():
+    calls = []
+    manager = ModelManager(loader=lambda model, path: calls.append((model, path)) or object())
+    first = manager.get("owner/model", "a")
+    assert manager.get("owner/model", "a") is first
+    second = manager.get("owner/model", "b")
+    assert second is not first
+    assert calls == [("owner/model", "a"), ("owner/model", "b")]
+
+
+def test_failed_load_is_reported_and_can_retry():
+    attempts = 0
+    def loader(*_):
+        nonlocal attempts
+        attempts += 1
+        raise RuntimeError("bad checkpoint")
+    manager = ModelManager(loader=loader)
+    with pytest.raises(ModelUnavailableError, match="bad checkpoint"):
+        manager.get("owner/model", "a")
+    with pytest.raises(ModelUnavailableError):
+        manager.get("owner/model", "a")
+    assert attempts == 2
+    assert "bad checkpoint" in manager.health_message
+```
+
+- [ ] **Step 2: Run tests to verify they fail**
+
+Run: `pytest tests/test_model_loader.py -v`
+
+Expected: FAIL with `ModuleNotFoundError: No module named 'model_loader'`.
+
+- [ ] **Step 3: Implement the model manager**
+
+Use a `threading.Lock`, one cached value, one active key, and an error string. The default loader must require `torch.cuda.is_available()`, call `download_checkpoint`, add `openpi_runtime` to the import path before importing OpenPI, resolve the exact UR config, and call:
+
+```python
+config = openpi_config.get_config("pi05_ur_demo_state")
+return policy_config.create_trained_policy(
+    config,
+    paths.checkpoint,
+    pytorch_device="cuda",
+)
+```
+
+On replacement or failure, drop references, run `gc.collect()`, and call `torch.cuda.empty_cache()` when CUDA is available.
+
+- [ ] **Step 4: Run manager tests**
+
+Run: `pytest tests/test_model_loader.py -v`
+
+Expected: PASS.
+
+- [ ] **Step 5: Commit**
+
+```bash
+git add model_loader.py tests/test_model_loader.py
+git commit -m "feat: add lazy pi05 policy lifecycle"
+```
+
+---
+
+### Task 5: Gradio application and end-to-end CPU verification
+
+**Files:**
+- Create: `app.py`
+- Create: `tests/test_app.py`
+- Modify: `README.md`
+
+**Interfaces:**
+- Consumes: `MODEL_MANAGER.get`, `run_prediction`, artifact defaults, thirteen Gradio input components, and three output components.
+- Produces: `predict_ui(...) -> tuple[pandas.DataFrame | None, str | None, str]` and queued Gradio `demo`.
+
+- [ ] **Step 1: Write failing UI-boundary tests**
+
+```python
+import app
+
+
+def test_predict_ui_returns_table_file_and_status(monkeypatch):
+    sentinel_policy = object()
+    monkeypatch.setattr(app.MODEL_MANAGER, "get", lambda *_: sentinel_policy)
+    monkeypatch.setattr(
+        app,
+        "run_prediction",
+        lambda *args: type("R", (), {"actions": "table", "json_path": "/tmp/result.json", "status": "done"})(),
+    )
+    result = app.predict_ui("owner/model", "checkpoint", object(), object(), "task", 1, 2, 3, 4, 5, 6, 0, 0)
+    assert result == ("table", "/tmp/result.json", "done")
+
+
+def test_predict_ui_turns_exceptions_into_status(monkeypatch):
+    monkeypatch.setattr(app.MODEL_MANAGER, "get", lambda *_: (_ for _ in ()).throw(RuntimeError("load failed")))
+    table, output_file, status = app.predict_ui("m", "c", object(), object(), "task", *([0] * 8))
+    assert table is None and output_file is None
+    assert status == "Error: load failed"
+
+
+def test_demo_exposes_prediction_controls():
+    config = app.demo.get_config_file()
+    labels = {component["props"].get("label") for component in config["components"]}
+    assert {"Fixed camera", "Wrist camera", "Task instruction", "Predict actions"} <= labels
+```
+
+- [ ] **Step 2: Run app tests to verify they fail**
+
+Run: `pytest tests/test_app.py -v`
+
+Expected: FAIL with `ModuleNotFoundError: No module named 'app'`.
+
+- [ ] **Step 3: Build the Gradio Space**
+
+Follow the `qwengr00t/app.py` layout. Add artifact textboxes, two `gr.Image(type="pil")` inputs, task textbox, seven labelled `gr.Number` state inputs, a non-negative integer trial index, primary predict button, status Markdown, non-editable action Dataframe, and JSON File output. Decorate `predict_ui` with `@spaces.GPU(duration=120)`, provide a no-op local fallback, and wire inputs in the exact signature order. Launch only under `if __name__ == "__main__"` using:
+
+```python
+demo.queue(default_concurrency_limit=1).launch()
+```
+
+Update README body with model-repository environment variables, checkpoint layout, input/state/action schemas, GPU requirement, Space deployment steps, and the optional local command `PYTHONPATH=openpi_runtime python app.py`.
+
+- [ ] **Step 4: Run focused and full CPU verification**
+
+Run: `pytest tests/test_app.py -v && pytest -v && python -m compileall -q app.py artifacts.py inference.py model_loader.py openpi_runtime/openpi openpi_runtime/openpi_client && git diff --check`
+
+Expected: all tests PASS, compilation exits zero, and `git diff --check` prints nothing.
+
+- [ ] **Step 5: Commit**
+
+```bash
+git add app.py README.md tests/test_app.py
+git commit -m "feat: deploy pi05 UR Gradio Space"
+```
+
+---
+
+### Task 6: Optional GPU checkpoint smoke test documentation
+
+**Files:**
+- Create: `tests/test_gpu_smoke.py`
+- Modify: `README.md`
+
+**Interfaces:**
+- Consumes: `PI05_GPU_SMOKE=1`, `PI05_MODEL_ID`, `PI05_CHECKPOINT_PATH`, CUDA, and external Hugging Face artifacts.
+- Produces: opt-in end-to-end proof that the real policy returns `(10, 7)` actions.
+
+- [ ] **Step 1: Add the skipped-by-default GPU test**
+
+```python
+import os
+import numpy as np
+from PIL import Image
+import pytest
+
+
+pytestmark = pytest.mark.skipif(os.getenv("PI05_GPU_SMOKE") != "1", reason="set PI05_GPU_SMOKE=1")
+
+
+def test_real_checkpoint_prediction():
+    from artifacts import resolve_checkpoint_path, resolve_model_id
+    from inference import run_prediction
+    from model_loader import MODEL_MANAGER
+    model_id = resolve_model_id()
+    checkpoint = resolve_checkpoint_path()
+    policy = MODEL_MANAGER.get(model_id, checkpoint)
+    image = Image.fromarray(np.zeros((224, 224, 3), dtype=np.uint8))
+    result = run_prediction(policy, image, image, "move safely", [0.0] * 7, 0, model_id, checkpoint)
+    assert result.actions.shape == (10, 7)
+```
+
+- [ ] **Step 2: Verify default test execution skips it safely**
+
+Run: `pytest tests/test_gpu_smoke.py -v`
+
+Expected: one SKIPPED test with reason `set PI05_GPU_SMOKE=1`.
+
+- [ ] **Step 3: Document the opt-in command**
+
+Add this command to README without embedding tokens:
+
+```bash
+PI05_GPU_SMOKE=1 PI05_MODEL_ID=owner/model PI05_CHECKPOINT_PATH=checkpoints/30000 pytest tests/test_gpu_smoke.py -v
+```
+
+- [ ] **Step 4: Run final verification**
+
+Run: `pytest -v && git diff --check && git status --short`
+
+Expected: CPU tests PASS with exactly the GPU test skipped; diff check is clean; status contains only Task 6 files before commit.
+
+- [ ] **Step 5: Commit**
+
+```bash
+git add README.md tests/test_gpu_smoke.py
+git commit -m "test: document pi05 GPU smoke check"
+```