adapters/pi05.py

"""pi05 adapter — serves a LeRobot pi05 checkpoint behind the frozen /act contract.

The intended payload is a NORI FINETUNE (lerobot/pi05_base has NO normalization
stats and cannot run zero-shot — Gate B). For adapter smoke-tests the only
public stats-bearing PyTorch checkpoint is the community DROID conversion
(jcoholich/pi05_droid_converted) — wrong embodiment for a robot, fine for
proving the serving path end-to-end.

Load + inference mirror the PROVEN local path in lelab/nori_rollout.py at our
pinned lerobot commit: PreTrainedConfig -> get_policy_class -> from_pretrained
-> make_pre_post_processors (with the device_processor override — fitted
processors bake in the training device), then a select_action DRAIN LOOP: the
policy's internal queue does ONE forward for a chunk of n_action_steps and pops
per call, so `horizon` calls = one chunk with per-action postprocessing —
identical semantics to local inference, no reliance on chunk-level
postprocessor broadcasting.

CHUNK SEMANTICS DIFFER FROM MOLMOACT2 — the client must read meta():
  - horizon = the checkpoint's n_action_steps (pi05 default 50, NOT 30)
  - chunk_hz = the TRAINING DATASET's control rate (Nori fleet: ~15) — not
    stored in the checkpoint config, so it MUST come from NORI_CHUNK_HZ.
  - cameras = the checkpoint's image feature keys, in the order /act images
    must be sent.
"""

from __future__ import annotations

import os
from typing import Optional

import numpy as np
from fastapi import HTTPException

from adapters.base import resolve_source

# Smoke default; production endpoints point MODEL_PATH/repository at the finetune.
FALLBACK_REPO = os.environ.get("NORI_PI05_CHECKPOINT", "jcoholich/pi05_droid_converted")
# Control rate of the chunk. Config carries no fps; default to the Nori fleet's
# achieved ~15 (raw-bundle finding). Endpoints for other data MUST set this.
CHUNK_HZ = float(os.environ.get("NORI_CHUNK_HZ", "15"))
MAX_IMAGES = 6


class Pi05Adapter:
    def __init__(self, model_path: str):
        self._probe_path = model_path
        self._source: Optional[str] = None
        self._policy = None
        self._pre = None
        self._post = None
        self._device = None
        self._image_keys: list[str] = []
        self._state_dim: Optional[int] = None
        self._action_dim: Optional[int] = None
        self._horizon: int = 50

    def load(self) -> None:
        import torch
        from lerobot.configs.policies import PreTrainedConfig
        from lerobot.policies.factory import get_policy_class, make_pre_post_processors

        self._source = resolve_source(self._probe_path, FALLBACK_REPO)
        print(f"[pi05] loading from {self._source}", flush=True)
        cfg = PreTrainedConfig.from_pretrained(self._source)
        if cfg.type != "pi05":
            raise RuntimeError(f"checkpoint is policy type {cfg.type!r}, expected pi05")
        policy_cls = get_policy_class(cfg.type)
        policy = policy_cls.from_pretrained(self._source, config=cfg)
        self._device = "cuda" if torch.cuda.is_available() else "cpu"
        policy.to(self._device)
        policy.eval()
        policy.reset()

        # Fitted processors carry the training-time normalization stats — the
        # whole reason a base checkpoint can't serve. Same device override as
        # lelab (fitted device_processor bakes in 'cuda'; harmless here, but a
        # CPU dev box would silently drop stats without it — the "driving but
        # not moving" incident class).
        dev = {"device_processor": {"device": self._device}}
        try:
            self._pre, self._post = make_pre_post_processors(
                cfg, pretrained_path=self._source,
                preprocessor_overrides=dev, postprocessor_overrides=dev)
        except Exception as e:
            raise RuntimeError(
                f"no fitted pre/post processors in {self._source} ({e}) — a pi05 "
                f"checkpoint without stats cannot serve (Gate B: use a finetune, "
                f"not pi05_base)") from e

        # Feature contract from the checkpoint itself, not assumptions.
        for key, feat in (cfg.input_features or {}).items():
            if "image" in key:
                self._image_keys.append(key)
            elif key == "observation.state":
                self._state_dim = feat.shape[0]
        for key, feat in (cfg.output_features or {}).items():
            if key == "action":
                self._action_dim = feat.shape[0]
        self._horizon = int(getattr(cfg, "n_action_steps", 50) or 50)
        self._policy = policy
        print(f"[pi05] cameras={self._image_keys} state_dim={self._state_dim} "
              f"action_dim={self._action_dim} horizon={self._horizon}", flush=True)

    def meta(self) -> dict:
        return {"kind": "pi05", "chunk_hz": CHUNK_HZ, "horizon": self._horizon,
                "dof": self._action_dim, "state_dim": self._state_dim,
                "cameras": list(self._image_keys), "max_images": MAX_IMAGES,
                "source": self._source, "supports_point": False,
                "supports_rtc": False}

    def act(self, *, images, state, instruction, num_steps, extras):
        import torch

        if len(images) != len(self._image_keys):
            raise HTTPException(
                status_code=422,
                detail=f"pi05 checkpoint expects exactly {len(self._image_keys)} "
                       f"images in this order: {self._image_keys} (got {len(images)})")
        rtc_note = ({"skipped": "unsupported for policy kind pi05"}
                    if extras.get("rtc") is not None else None)

        # State: robot scale in, processors normalize + pad (max_state_dim 32).
        st = np.zeros((self._state_dim or len(state),), dtype=np.float32)
        n = min(len(state), st.shape[0])
        st[:n] = state[:n]
        obs: dict = {
            "observation.state": torch.from_numpy(st).unsqueeze(0).to(self._device),
            "task": [instruction],
        }
        for key, img in zip(self._image_keys, images):
            # HxWx3 uint8 -> 1x3xHxW float in [0,1] (lerobot image convention).
            t = torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1)
            obs[key] = (t.float() / 255.0).unsqueeze(0).to(self._device)

        try:
            with torch.no_grad():
                self._policy.reset()          # fresh chunk per /act observation
                processed = self._pre(obs)
                chunk = []
                for _ in range(self._horizon):  # 1 forward + horizon-1 queue pops
                    a = self._policy.select_action(processed)
                    a = self._post(a)
                    chunk.append(a.squeeze(0).detach().float().cpu().numpy())
        except torch.cuda.OutOfMemoryError as e:
            torch.cuda.empty_cache()
            raise HTTPException(status_code=507, detail=f"CUDA OOM: {e}") from e
        acts = np.stack(chunk).astype(np.float32)
        if self._action_dim and acts.shape[-1] > self._action_dim:
            acts = acts[..., : self._action_dim]   # strip pad dims if post kept them
        return acts, {"rtc": rtc_note}