mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-06 09:37:06 +00:00
refactor(runtime): make language runtime policy-agnostic; drop VQA viz
Set up the runtime so a second language-conditioned policy reuses the CLI/REPL/UI instead of copying pi052's. The tick loop, REPL, panel, and interactive CLI are now policy-independent in lerobot/runtime/; a policy plugs in only a LanguageConditionedPolicyAdapter. - Move repl.py, ui.py, and runtime_cli.py (-> cli.py) from pi052/inference/ into lerobot/runtime/. Generalize labels/titles (panel_label param, [runtime] prefixes). - lerobot.runtime.cli.run(argv, *, adapter_factory, panel_label, prog) is the shared entry; policy loading already dispatches generically via the factory on cfg.type. - lerobot-pi052-runtime is now a thin entry (scripts/lerobot_pi052_runtime.py) that passes PI052PolicyAdapter into run(). pi052/inference/ keeps only the adapter. - Drop PI052Runtime back-compat wrapper (no consumers). - Drop VQA visualization: delete inference/vqa.py + test_pi052_vqa_loc.py, remove answer_vqa/VQAResult from the Protocol + adapter, and the /question command + overlay paths from the CLI/REPL. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Pepijn
@@ -32,8 +32,8 @@ This is the dual-head co-training pattern from the paper:
|
||||
|
||||
with α = 10.0 per § IV.D of arxiv:2504.16054. The π0.5 model splits
|
||||
inference into a text-prediction step followed by an action-prediction
|
||||
step, which the multi-rate ``PI052Runtime`` (in
|
||||
``lerobot.policies.pi052.inference``) drives at separate rates.
|
||||
step, which the multi-rate runtime (``lerobot.runtime``, via the
|
||||
``lerobot-pi052-runtime`` CLI) drives at separate rates.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
@@ -12,31 +12,14 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""PI052 runtime adapter and CLI helpers."""
|
||||
"""PI052 bridge to the generic language-conditioned runtime.
|
||||
|
||||
from lerobot.runtime import (
|
||||
LanguageConditionedRuntime,
|
||||
RuntimeState,
|
||||
Tick,
|
||||
TickClock,
|
||||
VQAResult,
|
||||
)
|
||||
The runtime, REPL, and CLI are policy-agnostic and live in
|
||||
:mod:`lerobot.runtime`. PI052 supplies only :class:`PI052PolicyAdapter`;
|
||||
the ``lerobot-pi052-runtime`` entry point wires it into
|
||||
:func:`lerobot.runtime.cli.run`.
|
||||
"""
|
||||
|
||||
from .pi052_adapter import PI052PolicyAdapter
|
||||
from .repl import StdinReader
|
||||
from .runtime import PI052Runtime
|
||||
from .ui import make_state_panel, print_robot_lines, print_user_line
|
||||
|
||||
__all__ = [
|
||||
"LanguageConditionedRuntime",
|
||||
"PI052PolicyAdapter",
|
||||
"PI052Runtime",
|
||||
"RuntimeState",
|
||||
"StdinReader",
|
||||
"Tick",
|
||||
"TickClock",
|
||||
"VQAResult",
|
||||
"make_state_panel",
|
||||
"print_robot_lines",
|
||||
"print_user_line",
|
||||
]
|
||||
__all__ = ["PI052PolicyAdapter"]
|
||||
|
||||
@@ -21,7 +21,7 @@ import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from lerobot.runtime import RuntimeState, VQAResult
|
||||
from lerobot.runtime import RuntimeState
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -73,18 +73,6 @@ class PI052PolicyAdapter:
|
||||
plan, _speech = split_plan_and_say(text)
|
||||
return "" if looks_like_gibberish(plan) else plan
|
||||
|
||||
def answer_vqa(
|
||||
self,
|
||||
question: str,
|
||||
camera: str | None,
|
||||
observation: dict[str, Any] | None,
|
||||
state: RuntimeState,
|
||||
) -> VQAResult:
|
||||
answer = self.select_text("vqa", observation, state, user_text=question)
|
||||
from .vqa import parse_vqa_answer # noqa: PLC0415
|
||||
|
||||
return VQAResult(answer=answer, parsed=parse_vqa_answer(answer), camera=camera)
|
||||
|
||||
def update_language_state(self, observation: dict[str, Any] | None, state: RuntimeState) -> None:
|
||||
chunks_per_gen = max(1, int(state.extra.get("subtask_chunks_per_gen", 1) or 1))
|
||||
if "_hl_chunks_until_gen" not in state.extra:
|
||||
@@ -171,8 +159,6 @@ class PI052PolicyAdapter:
|
||||
return messages
|
||||
if kind == "plan":
|
||||
return [{"role": "user", "content": state.task or ""}]
|
||||
if kind == "vqa":
|
||||
return [{"role": "user", "content": user_text or ""}]
|
||||
raise ValueError(f"Unknown PI052 text kind: {kind}")
|
||||
|
||||
|
||||
|
||||
@@ -1,68 +0,0 @@
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""PI052 compatibility wrapper for the generic language-conditioned runtime."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from lerobot.runtime import (
|
||||
LanguageConditionedRuntime,
|
||||
RuntimeState,
|
||||
Tick,
|
||||
TickClock,
|
||||
VQAResult,
|
||||
)
|
||||
|
||||
from .pi052_adapter import PI052PolicyAdapter
|
||||
|
||||
|
||||
class PI052Runtime(LanguageConditionedRuntime):
|
||||
"""Backwards-compatible PI052 runtime constructor."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
policy: Any,
|
||||
*,
|
||||
observation_provider: Callable[[], dict | None] | None = None,
|
||||
robot_executor: Callable[[Any], None] | None = None,
|
||||
event_collector: Callable[[RuntimeState], None] | None = None,
|
||||
chunk_hz: float = 4.0,
|
||||
ctrl_hz: float = 50.0,
|
||||
high_level_hz: float = 1.0,
|
||||
max_rate_hz: float = 50.0,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
policy_adapter=policy if isinstance(policy, PI052PolicyAdapter) else PI052PolicyAdapter(policy),
|
||||
observation_provider=observation_provider,
|
||||
action_executor=robot_executor,
|
||||
event_collector=event_collector,
|
||||
chunk_hz=chunk_hz,
|
||||
ctrl_hz=ctrl_hz,
|
||||
high_level_hz=high_level_hz,
|
||||
max_rate_hz=max_rate_hz,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LanguageConditionedRuntime",
|
||||
"PI052PolicyAdapter",
|
||||
"PI052Runtime",
|
||||
"RuntimeState",
|
||||
"Tick",
|
||||
"TickClock",
|
||||
"VQAResult",
|
||||
]
|
||||
@@ -1,406 +0,0 @@
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Interactive VQA for the PI052 runtime.
|
||||
|
||||
In ``/vlm`` mode a typed line is treated as a VQA question. This module
|
||||
runs the full interactive flow:
|
||||
|
||||
1. pull the current observation and list available cameras,
|
||||
2. ask the operator which camera to ground the question on,
|
||||
3. generate the answer with the VLM conditioned on that one camera,
|
||||
4. parse the JSON answer; if it carries a bounding box (``bbox``) or a
|
||||
point (``keypoint``), draw the overlay on the camera frame, save a
|
||||
PNG to ``./vqa_overlays/`` and auto-open it.
|
||||
|
||||
VQA answer schemas mirror the annotation pipeline's ``VQA_ANSWER_SHAPES``
|
||||
(see ``lerobot.annotations.steerable_pipeline.validator``):
|
||||
|
||||
* ``bbox`` — ``{"detections": [{"label", "bbox_format": "xyxy",
|
||||
"bbox": [x1, y1, x2, y2]}, ...]}``
|
||||
* ``keypoint`` — ``{"label", "point_format": "xy", "point": [x, y]}``
|
||||
* ``count`` / ``attribute`` / ``spatial`` — text-only, no overlay.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import webbrowser
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_IMAGE_PREFIX = "observation.images."
|
||||
|
||||
# PaliGemma detection / pointing vocabulary. PI052 trains spatial VQA
|
||||
# answers in this native ``<locNNNN>`` format (index in [0, 1023],
|
||||
# normalized to the image axis) instead of pixel-coordinate JSON, so the
|
||||
# answer string the runtime parses can be e.g.
|
||||
# ``<loc0512><loc0301> blue cube`` (point) or
|
||||
# ``<loc0100><loc0080><loc0400><loc0360> blue cube`` (box).
|
||||
_LOC_RE = re.compile(r"<loc(\d{1,4})>")
|
||||
|
||||
# Iteration order for shape matching — most specific keys first so an
|
||||
# answer is classified deterministically.
|
||||
_SHAPE_ORDER = ("bbox", "keypoint", "count", "attribute", "spatial")
|
||||
|
||||
_BBOX_COLOR = (255, 64, 64)
|
||||
_POINT_COLOR = (64, 220, 64)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Camera selection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def available_cameras(observation: dict | None) -> list[str]:
|
||||
"""Return the sorted ``observation.images.*`` keys present in ``observation``."""
|
||||
if not observation:
|
||||
return []
|
||||
return sorted(k for k in observation if isinstance(k, str) and k.startswith(_IMAGE_PREFIX))
|
||||
|
||||
|
||||
def camera_short_name(camera_key: str) -> str:
|
||||
"""Strip the ``observation.images.`` prefix for display."""
|
||||
return camera_key[len(_IMAGE_PREFIX) :] if camera_key.startswith(_IMAGE_PREFIX) else camera_key
|
||||
|
||||
|
||||
def prompt_camera_choice(
|
||||
cameras: list[str],
|
||||
*,
|
||||
input_fn: Any = input,
|
||||
print_fn: Any = print,
|
||||
) -> str | None:
|
||||
"""Ask the operator which camera frame to draw a VQA overlay on.
|
||||
|
||||
Accepts either the menu number or the (short or full) camera name.
|
||||
A single-camera setup auto-selects without prompting. Returns the
|
||||
chosen ``observation.images.*`` key, or ``None`` if the operator
|
||||
cancels / gives an invalid answer.
|
||||
"""
|
||||
if not cameras:
|
||||
return None
|
||||
if len(cameras) == 1:
|
||||
return cameras[0]
|
||||
print_fn("Draw the result on which camera?")
|
||||
for i, cam in enumerate(cameras, 1):
|
||||
print_fn(f" [{i}] {camera_short_name(cam)}")
|
||||
try:
|
||||
raw = str(input_fn("camera> ")).strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
return None
|
||||
if not raw:
|
||||
return cameras[0]
|
||||
if raw.isdigit():
|
||||
idx = int(raw) - 1
|
||||
return cameras[idx] if 0 <= idx < len(cameras) else None
|
||||
for cam in cameras:
|
||||
if raw == cam or raw == camera_short_name(cam):
|
||||
return cam
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Answer parsing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _loc_to_norm(idx: int) -> float:
|
||||
"""PaliGemma ``<locNNNN>`` index → normalized [0, 1] axis coordinate."""
|
||||
return max(0.0, min(1023.0, float(idx))) / 1023.0
|
||||
|
||||
|
||||
def parse_loc_answer(answer: str) -> dict | None:
|
||||
"""Parse a PaliGemma ``<loc>``-format spatial VQA answer.
|
||||
|
||||
Point: ``<label> <locY><locX>``; box: ``<label> <locY0><locX0><locY1><locX1>``;
|
||||
multiple boxes joined by `` ; `` (label/loc order irrelevant). Returns
|
||||
``{"kind", "payload", "normalized": True}`` with [0, 1] coords mirroring the
|
||||
JSON shapes (shared overlay code), or ``None`` without ``<loc>`` tokens.
|
||||
"""
|
||||
if not answer or "<loc" not in answer:
|
||||
return None
|
||||
segments = [seg for seg in answer.split(";") if "<loc" in seg]
|
||||
points: list[tuple[float, float, str]] = []
|
||||
boxes: list[tuple[float, float, float, float, str]] = []
|
||||
for seg in segments:
|
||||
locs = [int(m) for m in _LOC_RE.findall(seg)]
|
||||
label = _LOC_RE.sub("", seg).strip()
|
||||
if len(locs) == 2:
|
||||
y, x = (_loc_to_norm(v) for v in locs[:2])
|
||||
points.append((x, y, label))
|
||||
elif len(locs) >= 4:
|
||||
y1, x1, y2, x2 = (_loc_to_norm(v) for v in locs[:4])
|
||||
boxes.append((x1, y1, x2, y2, label))
|
||||
if boxes:
|
||||
detections = [
|
||||
{"label": lbl, "bbox_format": "xyxy", "bbox": [x1, y1, x2, y2]} for (x1, y1, x2, y2, lbl) in boxes
|
||||
]
|
||||
return {"kind": "bbox", "payload": {"detections": detections}, "normalized": True}
|
||||
if len(points) == 1:
|
||||
x, y, lbl = points[0]
|
||||
return {
|
||||
"kind": "keypoint",
|
||||
"payload": {"label": lbl, "point_format": "xy", "point": [x, y]},
|
||||
"normalized": True,
|
||||
}
|
||||
if points: # several bare points → treat as detections-as-points
|
||||
detections = [{"label": lbl, "bbox_format": "xyxy", "bbox": [x, y, x, y]} for (x, y, lbl) in points]
|
||||
return {"kind": "bbox", "payload": {"detections": detections}, "normalized": True}
|
||||
return None
|
||||
|
||||
|
||||
def parse_vqa_answer(answer: str) -> dict | None:
|
||||
"""Parse a VQA answer (``<loc>`` text or JSON) into ``{"kind", "payload"}``.
|
||||
|
||||
``kind`` is a ``VQA_ANSWER_SHAPES`` name or ``"unknown"``; ``<loc>`` answers
|
||||
are tried first. Returns ``None`` when neither format parses.
|
||||
"""
|
||||
if not answer or not answer.strip():
|
||||
return None
|
||||
loc_parsed = parse_loc_answer(answer)
|
||||
if loc_parsed is not None:
|
||||
return loc_parsed
|
||||
try:
|
||||
payload = json.loads(answer)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
|
||||
try:
|
||||
from lerobot.annotations.steerable_pipeline.validator import ( # noqa: PLC0415
|
||||
VQA_ANSWER_SHAPES,
|
||||
)
|
||||
|
||||
shapes = VQA_ANSWER_SHAPES
|
||||
except ImportError: # pragma: no cover - annotation extra not installed
|
||||
shapes = {
|
||||
"bbox": {"detections"},
|
||||
"keypoint": {"label", "point_format", "point"},
|
||||
"count": {"label", "count"},
|
||||
"attribute": {"label", "attribute", "value"},
|
||||
"spatial": {"subject", "relation", "object"},
|
||||
}
|
||||
|
||||
keys = set(payload)
|
||||
for kind in _SHAPE_ORDER:
|
||||
required = shapes.get(kind)
|
||||
if required and required <= keys:
|
||||
return {"kind": kind, "payload": payload}
|
||||
return {"kind": "unknown", "payload": payload}
|
||||
|
||||
|
||||
def answer_has_overlay(parsed: dict | None) -> bool:
|
||||
"""True iff ``parsed`` carries drawable spatial coordinates."""
|
||||
return bool(parsed) and parsed.get("kind") in ("bbox", "keypoint")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Overlay drawing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def observation_image_to_pil(image_tensor: Any) -> Any:
|
||||
"""Convert an ``observation.images.*`` tensor to a PIL RGB image.
|
||||
|
||||
The runtime observation stores images as ``(1, C, H, W)`` (or
|
||||
``(C, H, W)``) float tensors in ``[0, 1]``. Reuses
|
||||
``image_array_to_pil_image`` which handles the CHW→HWC transpose and
|
||||
the float→uint8 scaling.
|
||||
"""
|
||||
from lerobot.datasets.image_writer import image_array_to_pil_image # noqa: PLC0415
|
||||
|
||||
arr = image_tensor
|
||||
if hasattr(arr, "detach"):
|
||||
arr = arr.detach().cpu()
|
||||
if hasattr(arr, "numpy"):
|
||||
arr = arr.numpy()
|
||||
while arr.ndim > 3: # drop leading batch dim(s)
|
||||
arr = arr[0]
|
||||
return image_array_to_pil_image(arr).convert("RGB")
|
||||
|
||||
|
||||
def draw_vqa_overlay(image: Any, parsed: dict) -> Any:
|
||||
"""Draw ``bbox`` / ``keypoint`` answers onto a copy of ``image``.
|
||||
|
||||
Non-spatial answers (``count`` / ``attribute`` / ``spatial`` /
|
||||
``unknown``) are returned as an unmodified copy. When ``parsed`` has
|
||||
``normalized=True`` (PaliGemma ``<loc>`` answers) the [0, 1]
|
||||
coordinates are scaled to the image's pixel size.
|
||||
"""
|
||||
from PIL import ImageDraw # noqa: PLC0415
|
||||
|
||||
img = image.convert("RGB").copy()
|
||||
kind = parsed.get("kind")
|
||||
payload = parsed.get("payload") or {}
|
||||
draw = ImageDraw.Draw(img)
|
||||
w, h = img.size
|
||||
sx, sy = (w, h) if parsed.get("normalized") else (1, 1)
|
||||
|
||||
if kind == "bbox":
|
||||
for det in payload.get("detections") or []:
|
||||
if not isinstance(det, dict):
|
||||
continue
|
||||
box = det.get("bbox")
|
||||
if not (isinstance(box, list | tuple) and len(box) == 4):
|
||||
continue
|
||||
try:
|
||||
x1, y1, x2, y2 = (float(v) for v in box)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
x1, x2 = x1 * sx, x2 * sx
|
||||
y1, y2 = y1 * sy, y2 * sy
|
||||
draw.rectangle([x1, y1, x2, y2], outline=_BBOX_COLOR, width=3)
|
||||
label = str(det.get("label", "")).strip()
|
||||
if label:
|
||||
draw.text((x1 + 3, max(0.0, y1 - 12)), label, fill=_BBOX_COLOR)
|
||||
elif kind == "keypoint":
|
||||
point = payload.get("point")
|
||||
if isinstance(point, list | tuple) and len(point) == 2:
|
||||
try:
|
||||
x, y = float(point[0]) * sx, float(point[1]) * sy
|
||||
except (TypeError, ValueError):
|
||||
return img
|
||||
r = 6
|
||||
draw.ellipse([x - r, y - r, x + r, y + r], outline=_POINT_COLOR, width=3)
|
||||
draw.line([x - 2 * r, y, x + 2 * r, y], fill=_POINT_COLOR, width=2)
|
||||
draw.line([x, y - 2 * r, x, y + 2 * r], fill=_POINT_COLOR, width=2)
|
||||
label = str(payload.get("label", "")).strip()
|
||||
if label:
|
||||
draw.text((x + r + 3, y - r), label, fill=_POINT_COLOR)
|
||||
return img
|
||||
|
||||
|
||||
def _open_file(path: Path) -> None:
|
||||
"""Best-effort open ``path`` in the OS default viewer."""
|
||||
try:
|
||||
if sys.platform == "darwin":
|
||||
subprocess.run(["open", str(path)], check=False) # nosec B607
|
||||
elif sys.platform.startswith("linux"):
|
||||
subprocess.run(["xdg-open", str(path)], check=False) # nosec B607
|
||||
elif os.name == "nt":
|
||||
os.startfile(str(path)) # type: ignore[attr-defined] # noqa: S606 # nosec B606
|
||||
else: # pragma: no cover - exotic platform
|
||||
webbrowser.open(path.resolve().as_uri())
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("could not auto-open %s: %s", path, exc)
|
||||
|
||||
|
||||
def save_and_open_overlay(image: Any, out_dir: str | Path = "./vqa_overlays") -> Path:
|
||||
"""Save ``image`` as a timestamped PNG under ``out_dir`` and auto-open it."""
|
||||
out = Path(out_dir)
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
path = out / f"vqa_{int(time.time() * 1000)}.png"
|
||||
image.save(path)
|
||||
_open_file(path)
|
||||
return path
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Orchestrator
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def handle_vqa_query(
|
||||
*,
|
||||
policy_adapter: Any | None = None,
|
||||
policy: Any | None = None,
|
||||
observation_provider: Any,
|
||||
question: str,
|
||||
state: Any,
|
||||
input_fn: Any = input,
|
||||
print_fn: Any = print,
|
||||
) -> None:
|
||||
"""Run one interactive VQA question end to end.
|
||||
|
||||
Called synchronously from the input layer while the runtime is in
|
||||
``/question`` mode (the action loop is gated off, so the policy is
|
||||
not in concurrent use). Progress is reported via both
|
||||
``state.log`` (REPL panel scrollback) and ``print_fn`` (direct stdout)
|
||||
— in autonomous question mode the panel redraw is suspended,
|
||||
so the direct print is what the operator actually sees.
|
||||
"""
|
||||
if policy_adapter is None and policy is not None:
|
||||
from .pi052_adapter import PI052PolicyAdapter # noqa: PLC0415
|
||||
|
||||
policy_adapter = PI052PolicyAdapter(policy)
|
||||
|
||||
def report(line: str) -> None:
|
||||
"""Surface a line both to the panel scrollback and to stdout."""
|
||||
if hasattr(state, "log"):
|
||||
state.log(line)
|
||||
else:
|
||||
state.setdefault("log_lines", []).append(line)
|
||||
with suppress(Exception):
|
||||
print_fn(line)
|
||||
|
||||
if policy_adapter is None:
|
||||
report(" [warn] vqa: no policy adapter — skipping")
|
||||
return
|
||||
|
||||
observation: dict | None = None
|
||||
if observation_provider is not None:
|
||||
try:
|
||||
observation = observation_provider()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("observation_provider raised %s", exc)
|
||||
|
||||
# Feed the FULL observation (every camera + state) to the VLM. The
|
||||
# ``ask_vqa_*`` recipes look single-camera, but the image *block* is
|
||||
# stripped before tokenization — the actual frames reach the model
|
||||
# via PI052's ``OBS_IMAGES_*`` channels, and ``embed_prefix``
|
||||
# consumes *all* ``config.image_features`` regardless of which
|
||||
# camera the sub-recipe was tagged for. So the model always sees
|
||||
# every camera; the operator never has to name one to ask.
|
||||
result = policy_adapter.answer_vqa(question, None, observation, state)
|
||||
answer = result.answer
|
||||
if not answer:
|
||||
report(" [info] vqa gen returned empty")
|
||||
return
|
||||
report(f" vqa: {answer}")
|
||||
|
||||
parsed = result.parsed if result.parsed is not None else parse_vqa_answer(answer)
|
||||
if not answer_has_overlay(parsed):
|
||||
if parsed is None:
|
||||
report(" [info] vqa answer is not JSON — no overlay")
|
||||
return
|
||||
|
||||
# The answer carries a bounding box / point. Its pixel coordinates
|
||||
# are camera-specific and the text answer doesn't say which camera,
|
||||
# so ask the operator *now* — only when there is actually something
|
||||
# to draw — which camera frame to render the overlay on.
|
||||
cameras = available_cameras(observation)
|
||||
if observation is None or not cameras:
|
||||
report(" [info] no camera image — cannot draw overlay")
|
||||
return
|
||||
chosen = prompt_camera_choice(cameras, input_fn=input_fn, print_fn=print_fn)
|
||||
if chosen is None:
|
||||
report(" [info] overlay skipped — no camera selected")
|
||||
return
|
||||
try:
|
||||
pil = observation_image_to_pil(observation[chosen])
|
||||
overlay = draw_vqa_overlay(pil, parsed)
|
||||
path = save_and_open_overlay(overlay)
|
||||
report(f" vqa overlay ({camera_short_name(chosen)}) saved: {path}")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("vqa overlay failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG))
|
||||
report(f" [warn] vqa overlay failed: {type(exc).__name__}: {exc}")
|
||||
@@ -1786,7 +1786,7 @@ class PI052Policy(PreTrainedPolicy):
|
||||
suppress_loc_tokens: bool = False,
|
||||
use_kv_cache: bool = True,
|
||||
) -> str:
|
||||
"""Generate text continuation from a multimodal prefix (used by PI052Runtime).
|
||||
"""Generate text continuation from a multimodal prefix (used by the runtime CLI).
|
||||
|
||||
``suppress_loc_tokens=True`` masks PaliGemma's reserved ``<locDDDD>`` ids
|
||||
([256000, 257024)) before sampling — the pretraining prior drifts back to
|
||||
|
||||
@@ -12,7 +12,12 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Policy-agnostic high/low-level runtime for language-conditioned policies."""
|
||||
"""Policy-agnostic high/low-level runtime for language-conditioned policies.
|
||||
|
||||
The tick loop, REPL, and interactive CLI here are policy-independent; a
|
||||
policy plugs in by implementing :class:`LanguageConditionedPolicyAdapter`
|
||||
and calling :func:`lerobot.runtime.cli.run` with an adapter factory.
|
||||
"""
|
||||
|
||||
from .language_runtime import (
|
||||
LanguageConditionedPolicyAdapter,
|
||||
@@ -20,7 +25,6 @@ from .language_runtime import (
|
||||
RuntimeState,
|
||||
Tick,
|
||||
TickClock,
|
||||
VQAResult,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
@@ -29,5 +33,4 @@ __all__ = [
|
||||
"RuntimeState",
|
||||
"Tick",
|
||||
"TickClock",
|
||||
"VQAResult",
|
||||
]
|
||||
|
||||
@@ -12,13 +12,16 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""``lerobot-pi052-runtime`` — interactive REPL for trained PI052.
|
||||
"""Interactive REPL for a language-conditioned robot policy.
|
||||
|
||||
Drives the multi-rate runtime defined in
|
||||
:mod:`lerobot.policies.pi052.inference`. Stdin becomes the user
|
||||
channel: type a task, then natural-language interjections / questions.
|
||||
The runtime prints state changes (plan / subtask / memory / vqa /
|
||||
speech) as they happen.
|
||||
Policy-agnostic CLI over :class:`lerobot.runtime.LanguageConditionedRuntime`.
|
||||
A policy wires it up with :func:`run`, passing an adapter factory
|
||||
(``policy -> LanguageConditionedPolicyAdapter``); see
|
||||
``lerobot.scripts.lerobot_pi052_runtime`` for the PI052 entry point.
|
||||
|
||||
Stdin is the user channel: type a task, then natural-language
|
||||
interjections. The runtime prints state changes (plan / subtask /
|
||||
memory) as they happen.
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -27,16 +30,16 @@ Dry run on a Hub checkpoint, no robot connected — useful for sanity-
|
||||
checking text generation::
|
||||
|
||||
uv run lerobot-pi052-runtime \\
|
||||
--policy.path=pepijn223/pi052_hirobot_super_poulain_tool2 \\
|
||||
--policy.path=<repo-or-dir> \\
|
||||
--no_robot \\
|
||||
--task="please clean the kitchen"
|
||||
|
||||
Same, but feed real frames from an annotated dataset so plan / subtask
|
||||
/ memory / VQA generation runs against actual video + state::
|
||||
/ memory generation runs against actual video + state::
|
||||
|
||||
uv run lerobot-pi052-runtime \\
|
||||
--policy.path=pepijn223/pi052_hirobot_super_poulain_tool2 \\
|
||||
--dataset.repo_id=pepijn223/super_poulain_annotated \\
|
||||
--policy.path=<repo-or-dir> \\
|
||||
--dataset.repo_id=<annotated-dataset> \\
|
||||
--dataset.episode=0 \\
|
||||
--no_robot \\
|
||||
--task="please clean the kitchen"
|
||||
@@ -45,8 +48,7 @@ With a real robot::
|
||||
|
||||
uv run lerobot-pi052-runtime \\
|
||||
--policy.path=... \\
|
||||
--robot.type=so101 --robot.port=/dev/tty.usbmodem... \\
|
||||
--tts.voice=alba
|
||||
--robot.type=so101 --robot.port=/dev/tty.usbmodem...
|
||||
|
||||
``--policy.path`` accepts either a local directory or a Hugging Face
|
||||
Hub repo id. ``--dataset.repo_id`` likewise.
|
||||
@@ -61,23 +63,23 @@ from collections.abc import Callable
|
||||
from contextlib import suppress
|
||||
from typing import Any
|
||||
|
||||
from .language_runtime import LanguageConditionedPolicyAdapter, LanguageConditionedRuntime
|
||||
from .repl import _emit
|
||||
|
||||
logger = logging.getLogger("lerobot.pi052.runtime")
|
||||
logger = logging.getLogger("lerobot.runtime")
|
||||
|
||||
|
||||
def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
||||
def _parse_args(argv: list[str] | None = None, *, prog: str | None = None) -> argparse.Namespace:
|
||||
p = argparse.ArgumentParser(
|
||||
description=("Interactive REPL runtime for a trained PI052 hierarchical VLA checkpoint."),
|
||||
prog=prog,
|
||||
description="Interactive REPL runtime for a language-conditioned robot policy.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--policy.path",
|
||||
dest="policy_path",
|
||||
type=str,
|
||||
required=True,
|
||||
help=(
|
||||
"Local directory or Hugging Face Hub repo id pointing at a trained PI052 ``pretrained_model``."
|
||||
),
|
||||
help="Local directory or Hugging Face Hub repo id pointing at a trained ``pretrained_model``.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--dataset.repo_id",
|
||||
@@ -88,8 +90,8 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
||||
"Optional dataset (local path or Hub repo id) used to drive "
|
||||
"observations during dry-run inference. When set, the runtime "
|
||||
"reads camera frames + state from the chosen episode and feeds "
|
||||
"them into all forward passes — so plan / subtask / memory / "
|
||||
"VQA generation see the same visual context the policy was "
|
||||
"them into all forward passes — so plan / subtask / memory "
|
||||
"generation see the same visual context the policy was "
|
||||
"trained on."
|
||||
),
|
||||
)
|
||||
@@ -168,7 +170,7 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
||||
# an action executor that postprocesses (denormalises) the policy's
|
||||
# output and calls ``robot.send_action(...)`` at ``--ctrl_hz``. The
|
||||
# high-level REPL-style stdin still works in a background thread
|
||||
# for interjections / VQA.
|
||||
# for interjections.
|
||||
p.add_argument(
|
||||
"--robot.type",
|
||||
dest="robot_type",
|
||||
@@ -306,7 +308,7 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
||||
|
||||
|
||||
# Columns the runtime supplies itself via its own message stream — strip
|
||||
# them so ``RenderMessagesStep`` / ``PI052TextTokenizerStep`` are no-ops.
|
||||
# them so the recipe render + text-tokenizer processor steps are no-ops.
|
||||
_RUNTIME_OWNED_LANGUAGE_COLS = ("language_persistent", "language_events")
|
||||
|
||||
|
||||
@@ -331,7 +333,7 @@ def _load_policy_and_preprocessor(
|
||||
policy_path: str,
|
||||
dataset_repo_id: str | None,
|
||||
) -> tuple[Any, Any, Any, Any]:
|
||||
"""Load a PI052 checkpoint (local path or Hub repo id).
|
||||
"""Load a policy checkpoint (local path or Hub repo id).
|
||||
|
||||
Returns ``(policy, preprocessor, postprocessor, ds_meta)``.
|
||||
``preprocessor`` / ``postprocessor`` / ``ds_meta`` are ``None``
|
||||
@@ -439,7 +441,7 @@ def _bootstrap_state_from_dataset(
|
||||
) -> dict[str, str]:
|
||||
"""Pull task / active plan / memory / subtask at ``start_frame``, so the
|
||||
runtime's first prompt matches the canonical training prompts (an OOD
|
||||
prompt makes the model fall back to its dominant mode, VQA JSON spam).
|
||||
prompt makes the model fall back to its dominant training mode).
|
||||
"""
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset # noqa: PLC0415
|
||||
|
||||
@@ -519,7 +521,7 @@ def _select_task_interactively(
|
||||
# bootstrap default (may be None — REPL handles that).
|
||||
return bootstrap_task
|
||||
|
||||
print("\n[pi052] Select startup task:", flush=True)
|
||||
print("\n[runtime] Select startup task:", flush=True)
|
||||
if options:
|
||||
for i, opt in enumerate(options, 1):
|
||||
marker = " (dataset default)" if opt == bootstrap_task else ""
|
||||
@@ -887,18 +889,13 @@ def _build_robot_action_executor(
|
||||
def _print_runtime_help() -> None:
|
||||
"""Print the slash-command reference."""
|
||||
print(
|
||||
"[pi052] commands (arguments need no quotes):\n"
|
||||
"[runtime] commands (arguments need no quotes):\n"
|
||||
" /action <task> run the robot; an argument switches to that task\n"
|
||||
" /action resume the robot on the current task\n"
|
||||
" /action <seconds> run the robot for N seconds, then auto-pause\n"
|
||||
" /pause pause the action loop — robot holds position\n"
|
||||
" /question <text> pause and answer one VQA question\n"
|
||||
" /help show this help\n"
|
||||
" stop | quit | exit end the session\n"
|
||||
"\n"
|
||||
" VQA examples:\n"
|
||||
" /question point to the yellow cube -> point overlay\n"
|
||||
" /question detect the blue cube -> bounding-box overlay",
|
||||
" stop | quit | exit end the session",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
@@ -935,7 +932,6 @@ def _handle_slash_command(runtime: Any, line: str) -> bool:
|
||||
(seconds), no argument resumes the current
|
||||
task.
|
||||
``/pause`` pause the action loop — the robot holds.
|
||||
``/question "text"`` pause and answer one VQA question.
|
||||
``/help`` print the command reference.
|
||||
|
||||
Returns ``True`` when ``line`` was a recognised command (consumed).
|
||||
@@ -955,7 +951,7 @@ def _handle_slash_command(runtime: Any, line: str) -> bool:
|
||||
secs = float(rest)
|
||||
runtime.state["action_deadline"] = _time.monotonic() + secs
|
||||
print(
|
||||
f"[pi052] action — running {secs:g}s, then auto-pause",
|
||||
f"[runtime] action — running {secs:g}s, then auto-pause",
|
||||
flush=True,
|
||||
)
|
||||
else:
|
||||
@@ -965,16 +961,16 @@ def _handle_slash_command(runtime: Any, line: str) -> bool:
|
||||
# New task → drop the stale subtask so the high-level
|
||||
# loop regenerates one for the new goal.
|
||||
runtime.state["current_subtask"] = None
|
||||
print(f"[pi052] action — task: {rest!r}", flush=True)
|
||||
print(f"[runtime] action — task: {rest!r}", flush=True)
|
||||
elif runtime.state.get("task"):
|
||||
print(
|
||||
f"[pi052] action — resuming: {runtime.state['task']!r}",
|
||||
f"[runtime] action — resuming: {runtime.state['task']!r}",
|
||||
flush=True,
|
||||
)
|
||||
else:
|
||||
runtime.state["mode"] = "paused"
|
||||
print(
|
||||
"[pi052] no task set — use /action <your task>",
|
||||
"[runtime] no task set — use /action <your task>",
|
||||
flush=True,
|
||||
)
|
||||
return True
|
||||
@@ -983,22 +979,7 @@ def _handle_slash_command(runtime: Any, line: str) -> bool:
|
||||
runtime.state["mode"] = "paused"
|
||||
runtime.state["action_deadline"] = None
|
||||
_clear_action_queue(runtime)
|
||||
print("[pi052] paused — robot holding position", flush=True)
|
||||
return True
|
||||
|
||||
if cmd in {"/question", "/q", "/ask", "/vqa", "/vlm"}:
|
||||
# A question always pauses the action loop first so the policy
|
||||
# is not used concurrently by the background runtime thread.
|
||||
runtime.state["mode"] = "paused"
|
||||
runtime.state["action_deadline"] = None
|
||||
_clear_action_queue(runtime)
|
||||
if not rest:
|
||||
print(
|
||||
"[pi052] usage: /question <your question> (e.g. /question point to the yellow cube)",
|
||||
flush=True,
|
||||
)
|
||||
return True
|
||||
_run_vqa_query(runtime, rest)
|
||||
print("[runtime] paused — robot holding position", flush=True)
|
||||
return True
|
||||
|
||||
if cmd in {"/help", "/?"}:
|
||||
@@ -1007,22 +988,6 @@ def _handle_slash_command(runtime: Any, line: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def _run_vqa_query(runtime: Any, question: str) -> None:
|
||||
"""Run one interactive VQA question against the runtime's policy.
|
||||
|
||||
Invoked by ``/question`` — the action loop is paused first so the
|
||||
policy is free for a synchronous VQA call.
|
||||
"""
|
||||
from lerobot.policies.pi052.inference.vqa import handle_vqa_query # noqa: PLC0415
|
||||
|
||||
handle_vqa_query(
|
||||
policy_adapter=runtime.policy_adapter,
|
||||
observation_provider=runtime.observation_provider,
|
||||
question=question,
|
||||
state=runtime.state,
|
||||
)
|
||||
|
||||
|
||||
def _run_autonomous(
|
||||
runtime: Any,
|
||||
*,
|
||||
@@ -1030,6 +995,7 @@ def _run_autonomous(
|
||||
auto_start: bool,
|
||||
initial_task: str | None,
|
||||
max_ticks: int | None,
|
||||
panel_label: str = "Runtime",
|
||||
) -> int:
|
||||
"""Drive the runtime continuously at ``ctrl_hz`` while accepting
|
||||
stdin events in the foreground.
|
||||
@@ -1049,10 +1015,10 @@ def _run_autonomous(
|
||||
if not auto_start and runtime.state.get("mode", "paused") == "action":
|
||||
try:
|
||||
input(
|
||||
"[pi052] Robot connected — starting in ACTION mode. Press ENTER to begin, Ctrl+C to abort. "
|
||||
"[runtime] Robot connected — starting in ACTION mode. Press ENTER to begin, Ctrl+C to abort. "
|
||||
)
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
print("\n[pi052] aborted before start", flush=True)
|
||||
print("\n[runtime] aborted before start", flush=True)
|
||||
return 130
|
||||
|
||||
if initial_task:
|
||||
@@ -1061,7 +1027,7 @@ def _run_autonomous(
|
||||
thread = threading.Thread(
|
||||
target=runtime.run,
|
||||
kwargs={"max_ticks": max_ticks},
|
||||
name="pi052-runtime-loop",
|
||||
name="runtime-loop",
|
||||
daemon=True,
|
||||
)
|
||||
thread.start()
|
||||
@@ -1085,7 +1051,9 @@ def _run_autonomous(
|
||||
|
||||
runtime._flush_logs = _flush_into_scrollback # type: ignore[method-assign]
|
||||
|
||||
redraw = _make_state_panel_renderer(runtime, mode_label="autonomous", scrollback=_scrollback)
|
||||
redraw = _make_state_panel_renderer(
|
||||
runtime, mode_label="autonomous", panel_label=panel_label, scrollback=_scrollback
|
||||
)
|
||||
redraw()
|
||||
print(
|
||||
" [autonomous] /action <task> to run · /pause to stop · "
|
||||
@@ -1119,7 +1087,7 @@ def _run_autonomous(
|
||||
if hasattr(queue, "clear"):
|
||||
queue.clear()
|
||||
print(
|
||||
"\n[pi052] timed action elapsed — paused",
|
||||
"\n[runtime] timed action elapsed — paused",
|
||||
flush=True,
|
||||
)
|
||||
else:
|
||||
@@ -1130,7 +1098,7 @@ def _run_autonomous(
|
||||
print("> ", end="", flush=True)
|
||||
_panel_stop.wait(0.7)
|
||||
|
||||
panel_thread = threading.Thread(target=_panel_loop, name="pi052-panel-redraw", daemon=True)
|
||||
panel_thread = threading.Thread(target=_panel_loop, name="runtime-panel-redraw", daemon=True)
|
||||
panel_thread.start()
|
||||
|
||||
try:
|
||||
@@ -1160,11 +1128,11 @@ def _run_autonomous(
|
||||
_emit(runtime.state, "user_interjection")
|
||||
else:
|
||||
print(
|
||||
"[pi052] no task yet — use /action <your task> to start",
|
||||
"[runtime] no task yet — use /action <your task> to start",
|
||||
flush=True,
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
print("\n[pi052] interrupt — stopping", flush=True)
|
||||
print("\n[runtime] interrupt — stopping", flush=True)
|
||||
finally:
|
||||
_panel_stop.set()
|
||||
runtime.stop()
|
||||
@@ -1175,9 +1143,9 @@ def _run_autonomous(
|
||||
time.sleep(0.1)
|
||||
try:
|
||||
robot.disconnect()
|
||||
print("[pi052] robot disconnected", flush=True)
|
||||
print("[runtime] robot disconnected", flush=True)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f"[pi052] WARNING: robot.disconnect raised {exc}", flush=True)
|
||||
print(f"[runtime] WARNING: robot.disconnect raised {exc}", flush=True)
|
||||
|
||||
return 0
|
||||
|
||||
@@ -1186,6 +1154,7 @@ def _make_state_panel_renderer(
|
||||
runtime: Any,
|
||||
*,
|
||||
mode_label: str,
|
||||
panel_label: str = "Runtime",
|
||||
scrollback: list[str] | None = None,
|
||||
) -> Callable[[list[str] | None], None]:
|
||||
"""Return a closure that prints the task/subtask/plan/memory panel.
|
||||
@@ -1204,24 +1173,15 @@ def _make_state_panel_renderer(
|
||||
st = runtime.state
|
||||
run_mode = st.get("mode", "action")
|
||||
mode_tag = "[green]mode: action[/]" if run_mode == "action" else "[yellow]mode: paused[/]"
|
||||
console.rule(f"[bold]PI052[/] · {mode_label} · {mode_tag}", style="cyan")
|
||||
console.rule(f"[bold]{panel_label}[/] · {mode_label} · {mode_tag}", style="cyan")
|
||||
# Always-visible command hint so the operator never has to
|
||||
# remember the slash commands.
|
||||
if run_mode == "action":
|
||||
console.print(
|
||||
" [dim]commands:[/] [bold]/pause[/] stop · "
|
||||
"[bold]/question[/] <text> ask · [bold]/help[/] · [bold]stop[/]"
|
||||
)
|
||||
console.print(" [dim]commands:[/] [bold]/pause[/] stop · [bold]/help[/] · [bold]stop[/]")
|
||||
else:
|
||||
console.print(
|
||||
" [dim]commands:[/] [bold]/action[/] <task> run · "
|
||||
"[bold]/question[/] <text> ask · [bold]/help[/] · [bold]stop[/]"
|
||||
" [dim]commands:[/] [bold]/action[/] <task> run · [bold]/help[/] · [bold]stop[/]"
|
||||
)
|
||||
# Reference VQA prompts — the two answer shapes that draw an
|
||||
# overlay (point + bounding box). No quotes needed.
|
||||
console.print(
|
||||
" [dim]vqa examples:[/] /question point to the yellow cube · /question detect the blue cube"
|
||||
)
|
||||
for key, label in (
|
||||
("task", "task"),
|
||||
("current_subtask", "subtask"),
|
||||
@@ -1238,13 +1198,8 @@ def _make_state_panel_renderer(
|
||||
if isinstance(st.get("action_queue"), (list, tuple)) or hasattr(st.get("action_queue"), "__len__")
|
||||
else 0
|
||||
)
|
||||
pending = len(st.get("tool_calls_pending") or [])
|
||||
dispatched = int(st.get("actions_dispatched") or 0)
|
||||
console.print(
|
||||
f" [dim]queued actions: {queue_len} "
|
||||
f"dispatched: {dispatched} "
|
||||
f"pending tool calls: {pending}[/]"
|
||||
)
|
||||
console.print(f" [dim]queued actions: {queue_len} dispatched: {dispatched}[/]")
|
||||
|
||||
# Overfit / memorisation diagnostics. The high-level steps
|
||||
# surface the raw generation each time they fire (even when
|
||||
@@ -1278,8 +1233,8 @@ def _make_state_panel_renderer(
|
||||
console.print(f" [dim]gen rejects memory:{mem_gib} plan:{plan_gib}[/]")
|
||||
console.rule(style="cyan")
|
||||
# Runtime scrollback — log lines pushed from generation steps
|
||||
# (warnings, gibberish rejections, plan/say speech, vqa
|
||||
# answers). Last N lines, oldest first.
|
||||
# (warnings, gibberish rejections, plan speech). Last N lines,
|
||||
# oldest first.
|
||||
if scrollback:
|
||||
for line in scrollback:
|
||||
console.print(f" [magenta]{line.rstrip()}[/]")
|
||||
@@ -1290,9 +1245,7 @@ def _make_state_panel_renderer(
|
||||
console.print()
|
||||
if not st.get("task"):
|
||||
console.print(
|
||||
" [dim]Type [bold]/action <your task>[/bold] to begin, "
|
||||
"[bold]/question <text>[/bold] to ask, /help for commands, "
|
||||
"stop to exit.[/]"
|
||||
" [dim]Type [bold]/action <your task>[/bold] to begin, /help for commands, stop to exit.[/]"
|
||||
)
|
||||
|
||||
return _redraw
|
||||
@@ -1338,8 +1291,21 @@ def _silence_noisy_loggers() -> None:
|
||||
logging.getLogger("lerobot.robots.utils").setLevel(logging.ERROR)
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
args = _parse_args(argv)
|
||||
def run(
|
||||
argv: list[str] | None = None,
|
||||
*,
|
||||
adapter_factory: Callable[[Any], LanguageConditionedPolicyAdapter],
|
||||
panel_label: str = "Runtime",
|
||||
prog: str | None = None,
|
||||
) -> int:
|
||||
"""Run the interactive language-conditioned runtime CLI.
|
||||
|
||||
A policy wires this up by passing ``adapter_factory`` — a callable
|
||||
that turns a loaded policy into a :class:`LanguageConditionedPolicyAdapter`
|
||||
(typically the adapter class itself). ``panel_label`` names the state
|
||||
panel; ``prog`` sets the argparse program name for ``--help``.
|
||||
"""
|
||||
args = _parse_args(argv, prog=prog)
|
||||
logging.basicConfig(
|
||||
level=logging.DEBUG if args.verbose else logging.INFO,
|
||||
format="%(asctime)s %(levelname)s %(message)s",
|
||||
@@ -1349,14 +1315,14 @@ def main(argv: list[str] | None = None) -> int:
|
||||
autonomous_mode = bool(args.robot_type) and not args.no_robot
|
||||
if autonomous_mode and not args.dataset_repo_id:
|
||||
print(
|
||||
"[pi052] ERROR: autonomous robot mode requires --dataset.repo_id "
|
||||
"[runtime] ERROR: autonomous robot mode requires --dataset.repo_id "
|
||||
"for action-denormalisation stats and feature shapes. Pass the "
|
||||
"same dataset the policy was trained on.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 2
|
||||
|
||||
print(f"[pi052] loading policy from {args.policy_path}", flush=True)
|
||||
print(f"[runtime] loading policy from {args.policy_path}", flush=True)
|
||||
policy, preprocessor, postprocessor, ds_meta = _load_policy_and_preprocessor(
|
||||
args.policy_path, args.dataset_repo_id
|
||||
)
|
||||
@@ -1387,7 +1353,7 @@ def main(argv: list[str] | None = None) -> int:
|
||||
)
|
||||
if chosen:
|
||||
args.task = chosen
|
||||
print(f"[pi052] task: {args.task!r}", flush=True)
|
||||
print(f"[runtime] task: {args.task!r}", flush=True)
|
||||
|
||||
# No startup prompts — the runtime is command-driven. It comes up at
|
||||
# the command line in ``paused`` mode (robot idle) unless ``--mode``
|
||||
@@ -1401,7 +1367,7 @@ def main(argv: list[str] | None = None) -> int:
|
||||
|
||||
if autonomous_mode:
|
||||
print(
|
||||
f"[pi052] connecting to robot.type={args.robot_type} port={args.robot_port}",
|
||||
f"[runtime] connecting to robot.type={args.robot_type} port={args.robot_port}",
|
||||
flush=True,
|
||||
)
|
||||
robot = _build_robot(
|
||||
@@ -1425,7 +1391,7 @@ def main(argv: list[str] | None = None) -> int:
|
||||
)
|
||||
elif args.dataset_repo_id is not None:
|
||||
print(
|
||||
f"[pi052] streaming observations from {args.dataset_repo_id} "
|
||||
f"[runtime] streaming observations from {args.dataset_repo_id} "
|
||||
f"episode={args.dataset_episode} "
|
||||
f"start_frame={args.dataset_start_frame}",
|
||||
flush=True,
|
||||
@@ -1440,13 +1406,8 @@ def main(argv: list[str] | None = None) -> int:
|
||||
augment=getattr(args, "dataset_augment_at_inference", False),
|
||||
)
|
||||
|
||||
from lerobot.policies.pi052.inference import ( # noqa: PLC0415
|
||||
LanguageConditionedRuntime,
|
||||
PI052PolicyAdapter,
|
||||
)
|
||||
|
||||
runtime = LanguageConditionedRuntime(
|
||||
policy_adapter=PI052PolicyAdapter(policy),
|
||||
policy_adapter=adapter_factory(policy),
|
||||
observation_provider=observation_provider,
|
||||
action_executor=robot_executor,
|
||||
# No background event collector — the REPL drives ticks
|
||||
@@ -1466,7 +1427,7 @@ def main(argv: list[str] | None = None) -> int:
|
||||
runtime.state["text_gen_min_new_tokens"] = int(getattr(args, "text_min_new_tokens", 0) or 0)
|
||||
runtime.state["text_gen_temperature"] = float(getattr(args, "text_temperature", 0.0) or 0.0)
|
||||
runtime.state["text_gen_top_p"] = float(getattr(args, "text_top_p", 1.0) or 1.0)
|
||||
# Subtask throttle: the PI052 adapter updates language state only once every N
|
||||
# Subtask throttle: the adapter updates language state only once every N
|
||||
# action-chunk boundaries. Lets you run N action chunks per LM-head
|
||||
# subtask gen (e.g. ``--subtask_chunks_per_gen=5`` ≈ 5 flow-matching
|
||||
# chunks per subtask refresh) so the subtask doesn't churn while
|
||||
@@ -1493,6 +1454,7 @@ def main(argv: list[str] | None = None) -> int:
|
||||
auto_start=args.auto_start,
|
||||
initial_task=args.task,
|
||||
max_ticks=args.max_ticks,
|
||||
panel_label=panel_label,
|
||||
)
|
||||
# Fire one full pipeline tick at startup so the obs diagnostic
|
||||
# *and* the subtask generation actually run before the REPL
|
||||
@@ -1508,11 +1470,13 @@ def main(argv: list[str] | None = None) -> int:
|
||||
logger.warning("startup tick failed: %s", exc)
|
||||
startup_logs = []
|
||||
for line in startup_logs or []:
|
||||
print(f"[pi052] {line}", flush=True)
|
||||
return _run_repl(runtime, initial_task=args.task, max_ticks=args.max_ticks)
|
||||
print(f"[runtime] {line}", flush=True)
|
||||
return _run_repl(runtime, initial_task=args.task, max_ticks=args.max_ticks, panel_label=panel_label)
|
||||
|
||||
|
||||
def _run_repl(runtime: Any, *, initial_task: str | None, max_ticks: int | None) -> int:
|
||||
def _run_repl(
|
||||
runtime: Any, *, initial_task: str | None, max_ticks: int | None, panel_label: str = "Runtime"
|
||||
) -> int:
|
||||
"""Claude-Code-style block REPL.
|
||||
|
||||
Each turn redraws a status block (task / subtask / plan / memory)
|
||||
@@ -1526,12 +1490,12 @@ def _run_repl(runtime: Any, *, initial_task: str | None, max_ticks: int | None)
|
||||
from rich.console import Console # noqa: PLC0415
|
||||
except ImportError:
|
||||
print(
|
||||
"[pi052] rich is required for the interactive REPL. `pip install rich` and re-run.",
|
||||
"[runtime] rich is required for the interactive REPL. `pip install rich` and re-run.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 2
|
||||
|
||||
_redraw = _make_state_panel_renderer(runtime, mode_label="dry-run")
|
||||
_redraw = _make_state_panel_renderer(runtime, mode_label="dry-run", panel_label=panel_label)
|
||||
# Keep a local ``console`` just for the styled input prompt; the
|
||||
# state panel is owned by the shared renderer.
|
||||
console = Console(highlight=False)
|
||||
@@ -1570,7 +1534,7 @@ def _run_repl(runtime: Any, *, initial_task: str | None, max_ticks: int | None)
|
||||
# task to be meaningful.
|
||||
if not runtime.state.get("task"):
|
||||
print(
|
||||
"[pi052] no task yet — use /action <your task>",
|
||||
"[runtime] no task yet — use /action <your task>",
|
||||
flush=True,
|
||||
)
|
||||
_redraw(last_logs)
|
||||
@@ -1588,7 +1552,3 @@ def _run_repl(runtime: Any, *, initial_task: str | None, max_ticks: int | None)
|
||||
console.print("\n[dim]interrupted[/]")
|
||||
console.print("[dim]runtime stopped[/]")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -26,15 +26,6 @@ from typing import Any, Protocol
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class VQAResult:
|
||||
"""Text answer plus optional parsed spatial payload."""
|
||||
|
||||
answer: str
|
||||
parsed: dict[str, Any] | None = None
|
||||
camera: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RuntimeState:
|
||||
"""Explicit state shared by the runtime and policy adapter."""
|
||||
@@ -138,14 +129,6 @@ class LanguageConditionedPolicyAdapter(Protocol):
|
||||
user_text: str | None = None,
|
||||
) -> str: ...
|
||||
|
||||
def answer_vqa(
|
||||
self,
|
||||
question: str,
|
||||
camera: str | None,
|
||||
observation: dict[str, Any] | None,
|
||||
state: RuntimeState,
|
||||
) -> VQAResult: ...
|
||||
|
||||
|
||||
@dataclass
|
||||
class Tick:
|
||||
@@ -290,8 +273,6 @@ class LanguageConditionedRuntime:
|
||||
def maybe_handle_user_events(self) -> None:
|
||||
if self.state.take_event("user_interjection"):
|
||||
self._handle_user_interjection()
|
||||
if self.state.take_event("user_vqa_query"):
|
||||
self._handle_vqa_query()
|
||||
|
||||
def _handle_user_interjection(self) -> None:
|
||||
text = str(self.state.extra.get("recent_interjection") or "")
|
||||
@@ -306,16 +287,6 @@ class LanguageConditionedRuntime:
|
||||
self.state.set_context("plan", plan, label="plan")
|
||||
self.state.extra["recent_interjection"] = None
|
||||
|
||||
def _handle_vqa_query(self) -> None:
|
||||
question = str(self.state.extra.get("recent_vqa_query") or "")
|
||||
if not question:
|
||||
return
|
||||
observation = self._current_observation()
|
||||
result = self.policy_adapter.answer_vqa(question, None, observation, self.state)
|
||||
if result.answer:
|
||||
self.state.log(f" vqa: {result.answer}")
|
||||
self.state.extra["recent_vqa_query"] = None
|
||||
|
||||
def maybe_enqueue_action_chunk(self, *, force: bool = False) -> None:
|
||||
if self.state.mode != "action" or not self.state.task:
|
||||
return
|
||||
|
||||
@@ -11,21 +11,20 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Stdin REPL event collector for the PI052 runtime.
|
||||
"""Stdin REPL event collector for the language-conditioned runtime.
|
||||
|
||||
Reads non-blocking stdin lines, classifies each one heuristically:
|
||||
|
||||
"stop" / "quit" / "exit" → state["stop"] = True
|
||||
"/action" / "/pause" → set state["mode"]
|
||||
ends with "?" → user_vqa_query event
|
||||
starts with "task:" or first line → set runtime task
|
||||
anything else → user_interjection event
|
||||
|
||||
Plugged into the runtime via ``event_collector=StdinReader().poll``.
|
||||
|
||||
Note: the shipped CLI (``lerobot-pi052-runtime``) drives stdin
|
||||
directly in its REPL / autonomous loops and does *not* wire this
|
||||
collector; it's kept as the documented embedding hook and for tests.
|
||||
Note: the shipped CLI drives stdin directly in its REPL / autonomous
|
||||
loops and does *not* wire this collector; it's kept as the documented
|
||||
embedding hook and for tests.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -92,17 +91,12 @@ class StdinReader:
|
||||
if not state.get("task"):
|
||||
task = line[5:].strip() if lower.startswith("task:") else line
|
||||
state["task"] = task
|
||||
print(f"[pi052] Task: {task}", flush=True)
|
||||
print(f"[runtime] Task: {task}", flush=True)
|
||||
self._seen_first_line = True
|
||||
return
|
||||
|
||||
# Question → VQA; statement → interjection.
|
||||
if lower.endswith("?"):
|
||||
state["recent_vqa_query"] = line
|
||||
_emit(state, "user_vqa_query")
|
||||
else:
|
||||
state["recent_interjection"] = line
|
||||
_emit(state, "user_interjection")
|
||||
state["recent_interjection"] = line
|
||||
_emit(state, "user_interjection")
|
||||
|
||||
|
||||
def _emit(state: Any, event_name: str) -> None:
|
||||
@@ -11,7 +11,7 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Rich-based REPL layout for the PI052 runtime.
|
||||
"""Rich-based REPL layout for the language-conditioned runtime.
|
||||
|
||||
Two-zone terminal layout:
|
||||
|
||||
@@ -56,7 +56,7 @@ _STATE_KEYS = (
|
||||
)
|
||||
|
||||
|
||||
def make_state_panel(state: dict[str, Any]) -> Any:
|
||||
def make_state_panel(state: dict[str, Any], *, title: str = "Runtime state") -> Any:
|
||||
"""Render the persistent state panel for the live region.
|
||||
|
||||
Returns a :class:`rich.panel.Panel`. Caller passes it to
|
||||
@@ -76,19 +76,13 @@ def make_state_panel(state: dict[str, Any]) -> Any:
|
||||
table.add_row(label, rendered)
|
||||
queue = state.get("action_queue")
|
||||
queue_len = len(queue) if hasattr(queue, "__len__") else 0
|
||||
pending = state.get("tool_calls_pending") or []
|
||||
footer = Text.assemble(
|
||||
("queued actions: ", "dim"),
|
||||
(str(queue_len), "bold cyan"),
|
||||
(" pending tool calls: ", "dim"),
|
||||
(str(len(pending)), "bold magenta"),
|
||||
)
|
||||
footer = Text.assemble(("queued actions: ", "dim"), (str(queue_len), "bold cyan"))
|
||||
table.add_row("", footer)
|
||||
run_mode = state.get("mode", "action")
|
||||
mode_tag = "[green]action[/]" if run_mode == "action" else "[yellow]paused[/]"
|
||||
return Panel(
|
||||
table,
|
||||
title=f"[bold]PI052 state[/] · mode: {mode_tag}",
|
||||
title=f"[bold]{title}[/] · mode: {mode_tag}",
|
||||
border_style="cyan",
|
||||
)
|
||||
|
||||
@@ -13,13 +13,29 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Entry point for ``lerobot-pi052-runtime``."""
|
||||
"""Entry point for ``lerobot-pi052-runtime``.
|
||||
|
||||
Wires PI052's adapter into the generic runtime CLI. A new
|
||||
language-conditioned policy adds its own such entry point with its
|
||||
adapter — no runtime/REPL code to copy.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
from lerobot.policies.pi052.inference.runtime_cli import main
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
from lerobot.policies.pi052.inference import PI052PolicyAdapter
|
||||
from lerobot.runtime.cli import run
|
||||
|
||||
return run(
|
||||
argv,
|
||||
adapter_factory=PI052PolicyAdapter,
|
||||
panel_label="PI052",
|
||||
prog="lerobot-pi052-runtime",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
||||
@@ -23,9 +23,6 @@ def test_pi052_adapter_builds_recipe_prompts_from_runtime_state():
|
||||
{"role": "assistant", "content": "Previous plan:\npick then place"},
|
||||
{"role": "user", "content": "wait"},
|
||||
]
|
||||
assert adapter.messages_for("vqa", state, user_text="where is the cup?") == [
|
||||
{"role": "user", "content": "where is the cup?"}
|
||||
]
|
||||
|
||||
|
||||
def test_pi052_adapter_strips_say_markers_from_plan_text():
|
||||
@@ -37,15 +34,19 @@ def test_pi052_adapter_strips_say_markers_from_plan_text():
|
||||
|
||||
|
||||
def test_pi052_runtime_cli_smoke_does_not_load_model(monkeypatch):
|
||||
from lerobot.policies.pi052.inference import runtime_cli
|
||||
"""The pi052 entry wires its adapter into the generic runtime CLI."""
|
||||
from lerobot.runtime import cli
|
||||
from lerobot.scripts import lerobot_pi052_runtime
|
||||
|
||||
fake_policy = SimpleNamespace(config=SimpleNamespace(device="cpu"))
|
||||
|
||||
monkeypatch.setattr(
|
||||
runtime_cli,
|
||||
cli,
|
||||
"_load_policy_and_preprocessor",
|
||||
lambda policy_path, dataset_repo_id: (fake_policy, None, None, None),
|
||||
)
|
||||
monkeypatch.setattr(runtime_cli, "_run_repl", lambda runtime, initial_task, max_ticks: 0)
|
||||
monkeypatch.setattr(cli, "_run_repl", lambda runtime, **kwargs: 0)
|
||||
|
||||
assert runtime_cli.main(["--policy.path=fake", "--no_robot", "--task=clean", "--max_ticks=0"]) == 0
|
||||
assert (
|
||||
lerobot_pi052_runtime.main(["--policy.path=fake", "--no_robot", "--task=clean", "--max_ticks=0"]) == 0
|
||||
)
|
||||
|
||||
@@ -1,187 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Training-side conversion of VQA answers to PaliGemma ``<loc>`` text.
|
||||
|
||||
PI052 trains spatial VQA answers (``bbox`` / ``keypoint``) in
|
||||
PaliGemma's native ``<locNNNN>`` detection vocabulary so the LM head
|
||||
reuses the detection prior instead of fighting it (the ``<loc>``-salad
|
||||
bug). The dataset stores Qwen2.5-VL's grounding output — **0–1000
|
||||
normalized** coordinates, *not* pixels. (Verified empirically on the
|
||||
published datasets: x and y both span 0..1000 with ~30% of values
|
||||
exceeding the camera's pixel dimensions.) The conversion is therefore
|
||||
camera-resolution-independent. The dataset stays backbone-agnostic
|
||||
JSON; the conversion lives in PI052's tokenizer. These tests pin the
|
||||
JSON → ``<loc>`` rewrite.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("transformers")
|
||||
|
||||
from lerobot.policies.pi052.text_processor_pi052 import ( # noqa: E402
|
||||
_loc_token,
|
||||
_messages_vqa_to_loc,
|
||||
_vqa_answer_to_loc,
|
||||
register_paligemma_loc_tokens,
|
||||
)
|
||||
|
||||
|
||||
class _FakeTokenizer:
|
||||
"""Tracks ``add_tokens`` calls; mimics the bits ``register_paligemma_loc_tokens`` reads."""
|
||||
|
||||
def __init__(self, prepopulate: bool = False):
|
||||
self.added_tokens_encoder: dict[str, int] = {}
|
||||
self.calls: list[list[str]] = []
|
||||
if prepopulate:
|
||||
self.added_tokens_encoder["<loc0000>"] = 256000
|
||||
|
||||
def add_tokens(self, tokens: list[str]) -> int:
|
||||
self.calls.append(list(tokens))
|
||||
for t in tokens:
|
||||
self.added_tokens_encoder.setdefault(t, len(self.added_tokens_encoder) + 256000)
|
||||
return len(tokens)
|
||||
|
||||
|
||||
def test_register_loc_tokens_adds_full_1024_range():
|
||||
tok = _FakeTokenizer()
|
||||
out = register_paligemma_loc_tokens(tok)
|
||||
assert out is tok # returns same instance
|
||||
assert len(tok.calls) == 1
|
||||
added = tok.calls[0]
|
||||
assert len(added) == 1024
|
||||
assert added[0] == "<loc0000>"
|
||||
assert added[-1] == "<loc1023>"
|
||||
# Spot check a few in the middle.
|
||||
assert added[162] == "<loc0162>"
|
||||
assert added[759] == "<loc0759>"
|
||||
|
||||
|
||||
def test_register_loc_tokens_is_idempotent():
|
||||
"""If the loc tokens are already present we skip re-adding them."""
|
||||
tok = _FakeTokenizer(prepopulate=True)
|
||||
register_paligemma_loc_tokens(tok)
|
||||
register_paligemma_loc_tokens(tok)
|
||||
assert tok.calls == [] # never called add_tokens
|
||||
|
||||
|
||||
def test_loc_token_normalizes_and_clamps():
|
||||
# Default scale is the 0–1000 Qwen convention.
|
||||
assert _loc_token(0) == "<loc0000>"
|
||||
assert _loc_token(1000) == "<loc1023>"
|
||||
assert _loc_token(500) == f"<loc{round(500 / 1000 * 1023):04d}>"
|
||||
# out-of-range coordinates clamp into [0, 1023]
|
||||
assert _loc_token(9999) == "<loc1023>"
|
||||
assert _loc_token(-5) == "<loc0000>"
|
||||
|
||||
|
||||
def test_vqa_answer_to_loc_keypoint_normalized():
|
||||
# Label-first: avoids the "Assistant: → <loc>" attractor at training.
|
||||
answer = {"label": "blue cube", "point_format": "xy", "point": [500, 500]}
|
||||
assert _vqa_answer_to_loc(answer) == "blue cube <loc0512><loc0512>"
|
||||
|
||||
|
||||
def test_vqa_answer_to_loc_bbox_normalized():
|
||||
answer = {
|
||||
"detections": [{"label": "cube", "bbox_format": "xyxy", "bbox": [0, 0, 1000, 1000]}]
|
||||
}
|
||||
assert _vqa_answer_to_loc(answer) == "cube <loc0000><loc0000><loc1023><loc1023>"
|
||||
|
||||
|
||||
def test_vqa_answer_to_loc_multiple_detections_separator():
|
||||
answer = {
|
||||
"detections": [
|
||||
{"label": "blue", "bbox_format": "xyxy", "bbox": [0, 0, 500, 500]},
|
||||
{"label": "yellow", "bbox_format": "xyxy", "bbox": [500, 500, 1000, 1000]},
|
||||
]
|
||||
}
|
||||
out = _vqa_answer_to_loc(answer)
|
||||
# Each segment is "label <locs>", joined by " ; "
|
||||
assert out == (
|
||||
"blue <loc0000><loc0000><loc0512><loc0512> ; "
|
||||
"yellow <loc0512><loc0512><loc1023><loc1023>"
|
||||
)
|
||||
|
||||
|
||||
def test_vqa_answer_to_loc_returns_none_for_non_spatial():
|
||||
assert _vqa_answer_to_loc({"label": "cubes", "count": 2}) is None
|
||||
assert _vqa_answer_to_loc({"weird": "payload"}) is None
|
||||
|
||||
|
||||
def test_messages_vqa_to_loc_rewrites_target_turn():
|
||||
messages = [
|
||||
{"role": "user", "content": [{"type": "text", "text": "where is the cube?"}]},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": '{"label": "cube", "point_format": "xy", "point": [500, 500]}',
|
||||
},
|
||||
]
|
||||
out = _messages_vqa_to_loc(messages, target_indices=[1])
|
||||
assert out[1]["content"] == "cube <loc0512><loc0512>"
|
||||
# input messages are not mutated
|
||||
assert messages[1]["content"].startswith("{")
|
||||
|
||||
|
||||
def test_messages_vqa_to_loc_leaves_plain_text_targets_untouched():
|
||||
messages = [
|
||||
{"role": "user", "content": "pick the cube"},
|
||||
{"role": "assistant", "content": "pick up the cube"},
|
||||
]
|
||||
out = _messages_vqa_to_loc(messages, target_indices=[1])
|
||||
assert out[1]["content"] == "pick up the cube"
|
||||
|
||||
|
||||
def test_messages_vqa_to_loc_noop_without_target_indices():
|
||||
messages = [
|
||||
{"role": "assistant", "content": '{"label": "c", "point_format": "xy", "point": [1, 2]}'}
|
||||
]
|
||||
assert _messages_vqa_to_loc(messages, []) is messages
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Round-trip: training-side JSON -> <loc> -> runtime-side parse back
|
||||
#
|
||||
# Pins that the conversion preserves coordinate *order* (JSON is x-first,
|
||||
# PaliGemma <loc> is y-first) and the 0–1000 → [0, 1023] scaling. The
|
||||
# only loss is quantization to the 1024-bucket <loc> grid, so a coord
|
||||
# survives within half a bucket (~1000/2046 ≈ 0.49 on the 0–1000 scale).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_loc_round_trip_keypoint_preserves_normalized_coords():
|
||||
from lerobot.policies.pi052.inference.vqa import parse_vqa_answer
|
||||
|
||||
answer = {"label": "blue cube", "point_format": "xy", "point": [640, 480]}
|
||||
loc = _vqa_answer_to_loc(answer)
|
||||
parsed = parse_vqa_answer(loc)
|
||||
nx, ny = parsed["payload"]["point"]
|
||||
# parse_vqa_answer returns [0, 1] normalized; rescale back to 0–1000.
|
||||
assert abs(nx * 1000.0 - 640) <= 1000.0 / 2046 + 1e-6
|
||||
assert abs(ny * 1000.0 - 480) <= 1000.0 / 2046 + 1e-6
|
||||
assert parsed["payload"]["label"] == "blue cube"
|
||||
|
||||
|
||||
def test_loc_round_trip_bbox_preserves_order_and_scale():
|
||||
from lerobot.policies.pi052.inference.vqa import parse_vqa_answer
|
||||
|
||||
answer = {
|
||||
"detections": [{"label": "cube", "bbox_format": "xyxy", "bbox": [100, 200, 800, 900]}]
|
||||
}
|
||||
loc = _vqa_answer_to_loc(answer)
|
||||
parsed = parse_vqa_answer(loc)
|
||||
x1, y1, x2, y2 = parsed["payload"]["detections"][0]["bbox"]
|
||||
for got, want in ((x1, 100), (y1, 200), (x2, 800), (y2, 900)):
|
||||
assert abs(got * 1000.0 - want) <= 1000.0 / 2046 + 1e-6
|
||||
@@ -1,7 +1,6 @@
|
||||
from lerobot.runtime import (
|
||||
LanguageConditionedRuntime,
|
||||
RuntimeState,
|
||||
VQAResult,
|
||||
)
|
||||
|
||||
|
||||
@@ -19,9 +18,6 @@ class FakeAdapter:
|
||||
self.text_calls.append((kind, user_text))
|
||||
return "new plan"
|
||||
|
||||
def answer_vqa(self, question, camera, observation, state):
|
||||
return VQAResult(answer=f"answer: {question}")
|
||||
|
||||
def update_language_state(self, observation, state):
|
||||
self.updated = True
|
||||
state.set_context("subtask", "pick cup", label="subtask")
|
||||
|
||||
Reference in New Issue
Block a user