From 020dbab8f9780c02315c2777e5c89b2d2db18103 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Tue, 23 Jun 2026 12:00:25 +0200 Subject: [PATCH] refactor(pi052): introduce generic language runtime --- .../policies/language_conditioned/__init__.py | 35 + .../policies/language_conditioned/runtime.py | 432 +++++ .../policies/pi052/inference/__init__.py | 59 +- .../policies/pi052/inference/pi052_adapter.py | 311 +++ src/lerobot/policies/pi052/inference/repl.py | 11 +- .../policies/pi052/inference/runtime.py | 227 +-- .../policies/pi052/inference/runtime_cli.py | 1688 +++++++++++++++++ .../policies/pi052/inference/runtime_state.py | 95 - src/lerobot/policies/pi052/inference/steps.py | 941 +-------- .../policies/pi052/inference/triggers.py | 134 -- src/lerobot/policies/pi052/inference/ui.py | 9 +- src/lerobot/policies/pi052/inference/vqa.py | 54 +- src/lerobot/scripts/lerobot_pi052_runtime.py | 1667 +--------------- .../language_conditioned/test_runtime.py | 88 + .../pi052/test_pi052_runtime_adapter.py | 54 + 15 files changed, 2723 insertions(+), 3082 deletions(-) create mode 100644 src/lerobot/policies/language_conditioned/__init__.py create mode 100644 src/lerobot/policies/language_conditioned/runtime.py create mode 100644 src/lerobot/policies/pi052/inference/pi052_adapter.py create mode 100644 src/lerobot/policies/pi052/inference/runtime_cli.py delete mode 100644 src/lerobot/policies/pi052/inference/runtime_state.py delete mode 100644 src/lerobot/policies/pi052/inference/triggers.py create mode 100644 tests/policies/language_conditioned/test_runtime.py create mode 100644 tests/policies/pi052/test_pi052_runtime_adapter.py diff --git a/src/lerobot/policies/language_conditioned/__init__.py b/src/lerobot/policies/language_conditioned/__init__.py new file mode 100644 index 000000000..e81b0b64a --- /dev/null +++ b/src/lerobot/policies/language_conditioned/__init__.py @@ -0,0 +1,35 @@ +# 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. + +"""Generic runtime primitives for language-conditioned policies.""" + +from .runtime import ( + LanguageConditionedPolicyAdapter, + LanguageConditionedRuntime, + RuntimeState, + Tick, + TickClock, + ToolCall, + VQAResult, +) + +__all__ = [ + "LanguageConditionedPolicyAdapter", + "LanguageConditionedRuntime", + "RuntimeState", + "Tick", + "TickClock", + "ToolCall", + "VQAResult", +] diff --git a/src/lerobot/policies/language_conditioned/runtime.py b/src/lerobot/policies/language_conditioned/runtime.py new file mode 100644 index 000000000..f252bf037 --- /dev/null +++ b/src/lerobot/policies/language_conditioned/runtime.py @@ -0,0 +1,432 @@ +# 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. + +"""Small reusable runtime for language-conditioned robot policies.""" + +from __future__ import annotations + +import logging +import time +from collections import deque +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any, Protocol + +logger = logging.getLogger(__name__) + + +@dataclass +class ToolCall: + """A pending runtime tool invocation.""" + + name: str + arguments: dict[str, Any] = field(default_factory=dict) + + +@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.""" + + task: str = "" + language_context: dict[str, str] = field(default_factory=dict) + action_queue: deque[Any] = field(default_factory=deque) + events: set[str] = field(default_factory=set) + pending_tools: list[ToolCall] = field(default_factory=list) + log_lines: list[str] = field(default_factory=list) + mode: str = "action" + stop: bool = False + tick: Tick | None = None + actions_dispatched: int = 0 + action_deadline: float | None = None + extra: dict[str, Any] = field(default_factory=dict) + + _ALIASES = { + "current_plan": ("language_context", "plan"), + "current_subtask": ("language_context", "subtask"), + "current_memory": ("language_context", "memory"), + "tool_calls_pending": ("pending_tools", None), + "events_this_tick": ("events", None), + "_tick": ("tick", None), + } + + def emit(self, event_name: str) -> None: + self.events.add(event_name) + + def take_event(self, event_name: str) -> bool: + if event_name not in self.events: + return False + self.events.remove(event_name) + return True + + def log(self, line: str) -> None: + self.log_lines.append(line) + + def set_context(self, key: str, value: str | None, *, label: str | None = None) -> bool: + previous = self.language_context.get(key) + if previous == value: + return False + if value is None: + self.language_context.pop(key, None) + else: + self.language_context[key] = value + if label is not None and value: + self.log(f" {label}: {value}") + return True + + def get(self, key: str, default: Any = None) -> Any: + try: + return self[key] + except KeyError: + return default + + def setdefault(self, key: str, default: Any = None) -> Any: + current = self.get(key, None) + if current is not None: + return current + self[key] = default + return default + + def __getitem__(self, key: str) -> Any: + alias = self._ALIASES.get(key) + if alias is not None: + target, subkey = alias + value = getattr(self, target) + return value if subkey is None else value.get(subkey) + if hasattr(self, key): + return getattr(self, key) + if key in self.extra: + return self.extra[key] + raise KeyError(key) + + def __setitem__(self, key: str, value: Any) -> None: + alias = self._ALIASES.get(key) + if alias is not None: + target, subkey = alias + if subkey is None: + setattr(self, target, value) + elif value is None: + getattr(self, target).pop(subkey, None) + else: + getattr(self, target)[subkey] = value + return + if hasattr(self, key): + setattr(self, key, value) + else: + self.extra[key] = value + + +class LanguageConditionedPolicyAdapter(Protocol): + """Policy-specific bridge used by :class:`LanguageConditionedRuntime`.""" + + def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: ... + + def select_text( + self, + kind: str, + observation: dict[str, Any] | None, + state: RuntimeState, + user_text: str | None = None, + ) -> str: ... + + def parse_tool_calls(self, text: str) -> list[ToolCall]: ... + + def answer_vqa( + self, + question: str, + camera: str | None, + observation: dict[str, Any] | None, + state: RuntimeState, + ) -> VQAResult: ... + + +@dataclass +class Tick: + index: int + monotonic_seconds: float + + +@dataclass +class TickClock: + max_rate_hz: float = 50.0 + _index: int = field(default=0, init=False) + _last_seconds: float | None = field(default=None, init=False) + + def advance(self) -> Tick: + period = 1.0 / max(self.max_rate_hz, 0.1) + now = time.monotonic() + if self._last_seconds is not None: + sleep_for = (self._last_seconds + period) - now + if sleep_for > 0: + time.sleep(sleep_for) + now = time.monotonic() + self._last_seconds = now + self._index += 1 + return Tick(index=self._index, monotonic_seconds=now) + + +@dataclass +class _RateGate: + hz: float + _last_seconds: float | None = None + + def due(self, tick: Tick, *, force: bool = False) -> bool: + if force: + self._last_seconds = tick.monotonic_seconds + return True + period = 1.0 / max(self.hz, 1e-6) + if self._last_seconds is None or tick.monotonic_seconds - self._last_seconds >= period: + self._last_seconds = tick.monotonic_seconds + return True + return False + + def rearm(self) -> None: + self._last_seconds = None + + +@dataclass +class LanguageConditionedRuntime: + """Generic tick loop for language-conditioned robot policies.""" + + policy_adapter: LanguageConditionedPolicyAdapter + observation_provider: Callable[[], dict[str, Any] | None] | None = None + action_executor: Callable[[Any], None] | None = None + tools: dict[str, Any] = field(default_factory=dict) + 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 + + state: RuntimeState = field(default_factory=RuntimeState) + _chunk_gate: _RateGate = field(init=False) + _ctrl_gate: _RateGate = field(init=False) + _language_gate: _RateGate = field(init=False) + _stop: bool = field(default=False, init=False) + _last_dispatch_seconds: float | None = field(default=None, init=False) + + def __post_init__(self) -> None: + self._chunk_gate = _RateGate(self.chunk_hz) + self._ctrl_gate = _RateGate(self.ctrl_hz) + self._language_gate = _RateGate(self.high_level_hz) + + @property + def policy(self) -> Any: + return getattr(self.policy_adapter, "policy", self.policy_adapter) + + def set_task(self, task: str) -> None: + self.state.task = task + self.state.log(f"Task: {task}") + + def stop(self) -> None: + self._stop = True + self.state.stop = True + + def run(self, *, max_ticks: int | None = None) -> None: + clock = TickClock(max_rate_hz=self.max_rate_hz) + while not self._stop: + tick = clock.advance() + self._run_tick(tick) + self._flush_logs() + if self.state.stop: + self._stop = True + if max_ticks is not None and tick.index >= max_ticks: + break + self._on_shutdown() + + def step_once(self) -> list[str]: + previous = self.state.tick.index if self.state.tick is not None else 0 + tick = Tick(index=previous + 1, monotonic_seconds=time.monotonic()) + self._run_tick(tick, force_rates=True) + return list(self.state.log_lines) + + def _run_tick(self, tick: Tick, *, force_rates: bool = False) -> None: + self.state.tick = tick + self.state.log_lines = [] + if self.event_collector is not None: + self.event_collector(self.state) + self._handle_action_deadline() + if self.state.stop: + return + self.maybe_update_language_state(force=force_rates) + self.maybe_handle_user_events() + self.maybe_enqueue_action_chunk(force=force_rates) + self.dispatch_action(force=force_rates) + self.dispatch_tools() + self.state.events.clear() + + def _current_observation(self) -> dict[str, Any] | None: + if self.observation_provider is None: + return None + try: + return self.observation_provider() + except Exception as exc: # noqa: BLE001 + logger.debug("observation_provider failed: %s", exc) + return None + + def maybe_update_language_state(self, *, force: bool = False) -> None: + if self.state.mode != "action" or not self.state.task: + return + if self.state.action_queue: + self._language_gate.rearm() + return + if self.state.tick is None or not self._language_gate.due(self.state.tick, force=force): + return + update = getattr(self.policy_adapter, "update_language_state", None) + if update is None: + return + observation = self._current_observation() + try: + update(observation, self.state) + except Exception as exc: # noqa: BLE001 + logger.warning("language update failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG)) + self.state.log(f" [warn] language update failed: {type(exc).__name__}: {exc}") + + 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 "") + if not text: + return + observation = self._current_observation() + out = self.policy_adapter.select_text("interjection", observation, self.state, user_text=text) + if not out: + return + calls = self.policy_adapter.parse_tool_calls(out) + for call in calls: + self.state.pending_tools.append(call) + if calls: + self.state.emit("tool_call_pending") + for call in calls: + if call.name == "say" and call.arguments.get("text"): + self.state.log(f" speech: {call.arguments['text']}") + plan = getattr(self.policy_adapter, "plan_from_text", lambda value: value)(out) + if plan: + 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 + if self.state.action_queue: + return + if self.state.tick is None or not self._chunk_gate.due(self.state.tick, force=force): + return + observation = self._current_observation() + if observation is None: + return + try: + chunk = self.policy_adapter.select_action(observation, self.state) + except Exception as exc: # noqa: BLE001 + logger.warning("select_action failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG)) + self.state.log(f" [warn] select_action failed: {type(exc).__name__}: {exc}") + return + self._enqueue_chunk(chunk) + + def _enqueue_chunk(self, chunk: Any) -> None: + if chunk is None: + return + chunk_iter = chunk[0] if getattr(chunk, "ndim", None) == 3 else chunk + if getattr(chunk_iter, "ndim", None) == 1: + chunk_iter = chunk_iter.unsqueeze(0) + for step in chunk_iter: + self.state.action_queue.append(step.unsqueeze(0) if hasattr(step, "unsqueeze") else step) + try: + self.state.extra["last_chunk_size"] = int(chunk_iter.shape[0]) + except Exception: # noqa: BLE001 + self.state.extra["last_chunk_size"] = len(self.state.action_queue) + + def dispatch_action(self, *, force: bool = False) -> None: + if self.state.mode != "action": + self._last_dispatch_seconds = None + return + if self.state.tick is None or not self._ctrl_gate.due(self.state.tick, force=force): + return + queue = self.state.action_queue + if not queue: + self._last_dispatch_seconds = None + return + now = time.monotonic() + if self._last_dispatch_seconds is None or self.ctrl_hz <= 0: + n_to_pop = 1 + else: + n_to_pop = max(1, min(len(queue), int(round((now - self._last_dispatch_seconds) * self.ctrl_hz)))) + self._last_dispatch_seconds = now + latest = None + for _ in range(n_to_pop): + if not queue: + break + latest = queue.popleft() + self.state.actions_dispatched += 1 + if latest is not None and self.action_executor is not None: + self.action_executor(latest) + + def dispatch_tools(self) -> None: + if not (self.state.take_event("tool_call_pending") or self.state.pending_tools): + return + pending = list(self.state.pending_tools) + self.state.pending_tools = [] + for call in pending: + name = call.name if isinstance(call, ToolCall) else (call.get("function") or {}).get("name") + args = ( + call.arguments + if isinstance(call, ToolCall) + else (call.get("function") or {}).get("arguments", {}) + ) + tool = self.tools.get(name) + if tool is None: + self.state.log(f" [warn] tool {name!r} not registered — skipping call") + continue + try: + tool.call(args) + except Exception as exc: # noqa: BLE001 + self.state.log(f" [error] tool dispatch failed: {exc}") + + def _handle_action_deadline(self) -> None: + deadline = self.state.action_deadline + if self.state.mode == "action" and deadline is not None and time.monotonic() >= deadline: + self.state.mode = "paused" + self.state.action_deadline = None + self.state.action_queue.clear() + self.state.log("timed action elapsed — paused") + + def _flush_logs(self) -> None: + for line in self.state.log_lines: + print(f"[runtime] {line}", flush=True) + + def _on_shutdown(self) -> None: + self.state.action_queue.clear() + print("[runtime] stopped", flush=True) diff --git a/src/lerobot/policies/pi052/inference/__init__.py b/src/lerobot/policies/pi052/inference/__init__.py index 10f7f4726..012ff440c 100644 --- a/src/lerobot/policies/pi052/inference/__init__.py +++ b/src/lerobot/policies/pi052/inference/__init__.py @@ -11,62 +11,33 @@ # 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 inference / runtime orchestration. -Multi-rate runtime that mirrors the recipe-time training shape: +"""PI052 runtime adapter and CLI helpers.""" - low_level_execution → LowLevelForward + DispatchAction (high Hz) - high_level_subtask → HighLevelSubtaskFwd (~1 Hz) - memory_update → MemoryUpdateFwd (event: subtask_change) - user_interjection_response → UserInterjectionFwd (event: stdin) - ask_vqa_* → AskVQAFwd (event: stdin question) - speech tool calls → DispatchToolCalls (event: tool_call_pending) - -The CLI ``lerobot-pi052-runtime`` builds a ``PI052Runtime`` and calls -``run()``. -""" +from lerobot.policies.language_conditioned import ( + LanguageConditionedRuntime, + RuntimeState, + Tick, + TickClock, + ToolCall, + VQAResult, +) +from .pi052_adapter import PI052PolicyAdapter from .repl import StdinReader from .runtime import PI052Runtime -from .runtime_state import initial_runtime_state, push_log, set_if_changed, take_event -from .steps import ( - AskVQAFwd, - DispatchAction, - DispatchToolCalls, - HighLevelSubtaskFwd, - InferenceStep, - LowLevelForward, - MemoryUpdateFwd, - UserInterjectionFwd, -) -from .triggers import EventTrigger, HzTrigger, Tick, TickClock, Trigger from .ui import make_state_panel, print_robot_lines, print_user_line __all__ = [ - # runtime + "LanguageConditionedRuntime", + "PI052PolicyAdapter", "PI052Runtime", + "RuntimeState", "StdinReader", - # state helpers - "initial_runtime_state", - "push_log", - "set_if_changed", - "take_event", - # triggers - "Trigger", "Tick", "TickClock", - "HzTrigger", - "EventTrigger", - # steps - "InferenceStep", - "LowLevelForward", - "DispatchAction", - "HighLevelSubtaskFwd", - "MemoryUpdateFwd", - "UserInterjectionFwd", - "AskVQAFwd", - "DispatchToolCalls", - # UI + "ToolCall", + "VQAResult", "make_state_panel", "print_robot_lines", "print_user_line", diff --git a/src/lerobot/policies/pi052/inference/pi052_adapter.py b/src/lerobot/policies/pi052/inference/pi052_adapter.py new file mode 100644 index 000000000..8f4820374 --- /dev/null +++ b/src/lerobot/policies/pi052/inference/pi052_adapter.py @@ -0,0 +1,311 @@ +# 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 adapter for the generic language-conditioned runtime.""" + +from __future__ import annotations + +import logging +import re +from dataclasses import dataclass +from typing import Any + +from lerobot.policies.language_conditioned import RuntimeState, ToolCall, VQAResult + +logger = logging.getLogger(__name__) + +_LOC_TOKENIZER_CACHE: dict[str, Any] = {} +_SAY_RE = re.compile(r"<\s*say\s*>(.*?)<\s*/\s*say\s*>", re.IGNORECASE | re.DOTALL) + + +@dataclass +class PI052PolicyAdapter: + """Runtime bridge for PI052 policies.""" + + policy: Any + + def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: + subtask = state.language_context.get("subtask") or state.task or "" + text_batch = _build_text_batch( + self.policy, + [{"role": "user", "content": subtask}], + add_generation_prompt=False, + ) + from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS # noqa: PLC0415 + + batch = dict(observation) + batch[OBS_LANGUAGE_TOKENS] = text_batch["lang_tokens"] + batch[OBS_LANGUAGE_ATTENTION_MASK] = text_batch["lang_masks"] + return self.policy.predict_action_chunk(batch) + + def select_text( + self, + kind: str, + observation: dict[str, Any] | None, + state: RuntimeState, + user_text: str | None = None, + ) -> str: + messages = self.messages_for(kind, state, user_text=user_text) + return _generate_with_policy( + self.policy, + messages, + observation=observation, + state=state, + label=f"{kind} gen", + min_new_tokens=int(state.extra.get("text_gen_min_new_tokens") or 0), + temperature=float(state.extra.get("text_gen_temperature") or 0.0), + top_p=float(state.extra.get("text_top_p") or 1.0), + suppress_loc_tokens=kind in {"subtask", "memory", "interjection"}, + ) + + def parse_tool_calls(self, text: str) -> list[ToolCall]: + _plan, speech = split_plan_and_say(text) + return [ToolCall("say", {"text": speech})] if speech else [] + + def plan_from_text(self, text: str) -> str: + 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: + state.extra["_hl_chunks_until_gen"] = 0 + if state.extra["_hl_chunks_until_gen"] > 0: + state.extra["_hl_chunks_until_gen"] -= 1 + return + state.extra["_hl_chunks_until_gen"] = chunks_per_gen - 1 + + msg = self.select_text("subtask", observation, state) + state.extra["last_subtask_raw"] = msg or "" + if not msg: + empties = int(state.extra.get("subtask_empty_count") or 0) + 1 + state.extra["subtask_empty_count"] = empties + if empties == 1 or empties % 5 == 0: + debug = getattr(self.policy, "_last_select_message_debug", "") or "" + state.log( + f" [info] subtask gen empty (x{empties}); {debug}" + if debug + else f" [info] subtask gen returned empty (x{empties})" + ) + return + if looks_like_gibberish(msg): + count = int(state.extra.get("subtask_gibberish_count") or 0) + 1 + state.extra["subtask_gibberish_count"] = count + if count == 1 or count % 30 == 0: + state.log(f" [info] subtask gen rejected (gibberish x{count}): {msg[:60]!r}") + return + + previous = state.language_context.get("subtask") + changed = state.set_context("subtask", msg, label="subtask") + if not changed: + state.extra["subtask_repeat_count"] = int(state.extra.get("subtask_repeat_count") or 0) + 1 + return + + state.extra["subtask_repeat_count"] = 0 + if previous: + state.extra["prior_subtask"] = previous + self._update_memory(observation, state) + + def _update_memory(self, observation: dict[str, Any] | None, state: RuntimeState) -> None: + new_memory = self.select_text("memory", observation, state) + state.extra["last_memory_raw"] = new_memory or "" + if not new_memory: + return + if looks_like_gibberish(new_memory): + count = int(state.extra.get("memory_gibberish_count") or 0) + 1 + state.extra["memory_gibberish_count"] = count + state.log(f" [info] memory gen rejected (gibberish x{count}): {new_memory[:60]!r}") + return + state.set_context("memory", new_memory, label="memory") + + def messages_for( + self, + kind: str, + state: RuntimeState, + *, + user_text: str | None = None, + ) -> list[dict[str, Any]]: + if kind == "subtask": + return [{"role": "user", "content": state.task or ""}] + if kind == "memory": + messages = [{"role": "user", "content": state.task or ""}] + if state.language_context.get("memory"): + messages.append( + { + "role": "assistant", + "content": f"Previous memory: {state.language_context['memory']}", + } + ) + if state.extra.get("prior_subtask"): + messages.append( + {"role": "user", "content": f"Completed subtask: {state.extra['prior_subtask']}"} + ) + return messages + if kind == "interjection": + messages = [{"role": "user", "content": state.task or ""}] + if state.language_context.get("plan"): + messages.append( + {"role": "assistant", "content": f"Previous plan:\n{state.language_context['plan']}"} + ) + if user_text: + messages.append({"role": "user", "content": user_text}) + 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}") + + +def _get_loc_tokenizer(tok_name: str, auto_tokenizer_cls: Any, register_loc_fn: Any) -> Any: + tokenizer = _LOC_TOKENIZER_CACHE.get(tok_name) + if tokenizer is None: + tokenizer = register_loc_fn(auto_tokenizer_cls.from_pretrained(tok_name)) + _LOC_TOKENIZER_CACHE[tok_name] = tokenizer + return tokenizer + + +def _build_text_batch( + policy: Any, + prompt_messages: list[dict[str, Any]], + *, + add_generation_prompt: bool = True, +) -> dict[str, Any]: + import torch # noqa: PLC0415 + from transformers import AutoTokenizer # noqa: PLC0415 + + from lerobot.policies.pi052.text_processor_pi052 import ( # noqa: PLC0415 + _flatten_say_tool_calls, + _format_messages, + _strip_blocks, + register_paligemma_loc_tokens, + ) + + tok_name = getattr(policy.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224" + tokenizer = _get_loc_tokenizer(tok_name, AutoTokenizer, register_paligemma_loc_tokens) + + messages = [_strip_blocks(_flatten_say_tool_calls(m)) for m in prompt_messages] + prompt, _spans = _format_messages(messages) + if add_generation_prompt: + prompt = prompt + "Assistant: " + + encoded = tokenizer(prompt, return_tensors="pt") + ids = encoded["input_ids"] + attn = encoded.get("attention_mask") + if attn is None and tokenizer.pad_token_id is not None: + attn = ids != tokenizer.pad_token_id + if attn is not None and hasattr(attn, "dtype") and attn.dtype != torch.bool: + attn = attn.bool() + + device = getattr(getattr(policy, "config", None), "device", None) + if device is not None: + try: + ids = ids.to(device) + if attn is not None and hasattr(attn, "to"): + attn = attn.to(device) + except Exception as exc: # noqa: BLE001 + logger.debug("could not move pi052 lang tokens to %s: %s", device, exc) + return {"lang_tokens": ids, "lang_masks": attn, "tokenizer": tokenizer} + + +def _generate_with_policy( + policy: Any, + messages: list[dict[str, Any]], + *, + observation: dict[str, Any] | None = None, + state: RuntimeState | None = None, + label: str = "select_message", + min_new_tokens: int = 0, + temperature: float = 0.0, + top_p: float = 1.0, + suppress_loc_tokens: bool = False, +) -> str: + if not hasattr(policy, "select_message"): + if state is not None: + state.log(f" [warn] policy has no select_message — skipping {label}") + return "" + text_batch = _build_text_batch(policy, messages) + try: + from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS # noqa: PLC0415 + + batch: dict[str, Any] = { + OBS_LANGUAGE_TOKENS: text_batch["lang_tokens"], + OBS_LANGUAGE_ATTENTION_MASK: text_batch["lang_masks"], + } + if observation: + for k, v in observation.items(): + if isinstance(k, str) and k.startswith("observation.") and k not in batch: + batch[k] = v + return policy.select_message( + batch, + tokenizer=text_batch["tokenizer"], + min_new_tokens=min_new_tokens, + temperature=temperature, + top_p=top_p, + suppress_loc_tokens=suppress_loc_tokens, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("%s failed: %s", label, exc, exc_info=logger.isEnabledFor(logging.DEBUG)) + if state is not None: + state.log(f" [warn] {label} failed: {type(exc).__name__}: {exc}") + return "" + + +def looks_like_gibberish(text: str) -> bool: + if not text or not text.strip(): + return True + stripped = text.strip() + alpha = sum(1 for c in stripped if c.isalpha()) + if alpha < max(3, len(stripped) // 8): + return True + if stripped.startswith('":') and stripped.count('"') > stripped.count(" "): + return True + if len(set(stripped)) <= 2 and len(stripped) > 4: + return True + cleaned = stripped.replace("\n", " ").replace(":", " ") + for marker in ("Assistant", "User", "Ass "): + if marker in cleaned and len(cleaned.split()) < 4: + return True + tokens = [t for t in cleaned.split() if any(c.isalpha() for c in t)] + unique_alpha = {t.lower() for t in tokens} + if len(unique_alpha) < 3 and len(stripped) < 80: + return True + return len(tokens) >= 8 and len(unique_alpha) <= max(3, len(tokens) // 10) + + +def split_plan_and_say(text: str) -> tuple[str, str]: + if not text: + return "", "" + match = _SAY_RE.search(text) + if not match: + return text.strip(), "" + speech = match.group(1).strip().strip('"').strip("'") + plan = (text[: match.start()] + text[match.end() :]).strip() + return plan, speech + + +def messages_for_vqa(question: str) -> list[dict[str, Any]]: + return [{"role": "user", "content": question}] diff --git a/src/lerobot/policies/pi052/inference/repl.py b/src/lerobot/policies/pi052/inference/repl.py index 2b8813f58..6f025b15a 100644 --- a/src/lerobot/policies/pi052/inference/repl.py +++ b/src/lerobot/policies/pi052/inference/repl.py @@ -99,7 +99,14 @@ class StdinReader: # Question → VQA; statement → interjection. if lower.endswith("?"): state["recent_vqa_query"] = line - state.setdefault("events_this_tick", []).append("user_vqa_query") + _emit(state, "user_vqa_query") else: state["recent_interjection"] = line - state.setdefault("events_this_tick", []).append("user_interjection") + _emit(state, "user_interjection") + + +def _emit(state: Any, event_name: str) -> None: + if hasattr(state, "emit"): + state.emit(event_name) + else: + state.setdefault("events_this_tick", []).append(event_name) diff --git a/src/lerobot/policies/pi052/inference/runtime.py b/src/lerobot/policies/pi052/inference/runtime.py index 3ca3e0576..57ec3b1d7 100644 --- a/src/lerobot/policies/pi052/inference/runtime.py +++ b/src/lerobot/policies/pi052/inference/runtime.py @@ -11,195 +11,62 @@ # 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 runtime loop. -Threads the multi-rate inference pipeline together with a stdin REPL -event collector, drives ticks through :class:`TickClock`, and prints -state-change updates to the user. -""" +"""PI052 compatibility wrapper for the generic language-conditioned runtime.""" from __future__ import annotations -import logging -from collections import deque -from dataclasses import dataclass, field -from typing import Any, Callable +from collections.abc import Callable +from typing import Any -from .runtime_state import initial_runtime_state, push_log -from .steps import ( - AskVQAFwd, - DispatchAction, - DispatchToolCalls, - HighLevelSubtaskFwd, - InferenceStep, - LowLevelForward, - MemoryUpdateFwd, +from lerobot.policies.language_conditioned import ( + LanguageConditionedRuntime, + RuntimeState, + Tick, + TickClock, + ToolCall, + VQAResult, ) -from .triggers import EventTrigger, HzTrigger, TickClock -logger = logging.getLogger(__name__) +from .pi052_adapter import PI052PolicyAdapter -@dataclass -class PI052Runtime: - """Compose the inference pipeline and drive it tick-by-tick.""" +class PI052Runtime(LanguageConditionedRuntime): + """Backwards-compatible PI052 runtime constructor.""" - policy: Any - tools: dict[str, Any] = field(default_factory=dict) - """Name → tool-instance dict, e.g. ``{"say": SayTool(...)}``. Read - from :func:`lerobot.tools.get_tools(meta)` when wiring the - runtime.""" - observation_provider: Callable[[], dict | None] | None = None - """Closure returning the current preprocessed observation batch. - ``None`` for dry-run / language-only sessions.""" - robot_executor: Callable[[Any], None] | None = None - """Closure that takes one action chunk and forwards it to the - robot. ``None`` for dry-run.""" - event_collector: Callable[[dict], None] | None = None - """Per-tick hook that polls external sources (stdin, network) and - appends event names to ``state["events_this_tick"]``.""" - chunk_hz: float = 4.0 - ctrl_hz: float = 50.0 - high_level_hz: float = 1.0 - max_rate_hz: float = 50.0 + def __init__( + self, + policy: Any, + *, + tools: dict[str, Any] | None = None, + 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, + tools=tools or {}, + event_collector=event_collector, + chunk_hz=chunk_hz, + ctrl_hz=ctrl_hz, + high_level_hz=high_level_hz, + max_rate_hz=max_rate_hz, + ) - pipeline: list[InferenceStep] = field(init=False) - state: dict[str, Any] = field(init=False) - _stop: bool = field(default=False, init=False) - def __post_init__(self) -> None: - # Subtask + memory + VQA configuration. Pipeline: - # - # HighLevelSubtaskFwd → generate the next subtask via the LM - # head at ~``high_level_hz``; writes - # ``current_subtask`` and emits - # ``subtask_change`` on a transition. - # MemoryUpdateFwd → on ``subtask_change``, refresh - # ``current_memory`` from the - # ``memory_update`` head. - # AskVQAFwd → answer camera-grounded stdin questions. - # LowLevelForward → action chunk conditioned on the - # generated ``current_subtask``. - # DispatchAction → drain the chunk to the robot. - # DispatchToolCalls → fire any pending tool calls. - # - # Order matters: ``HighLevelSubtaskFwd`` must run before - # ``MemoryUpdateFwd`` so the event is visible the same tick, and - # both must run before ``LowLevelForward`` (which is gated on - # "action queue empty") so the chunk consumes the freshest - # subtask. ``UserInterjectionFwd`` is still importable but - # disabled until plan generation is wired in. - self.pipeline = [ - HighLevelSubtaskFwd( - trigger=HzTrigger(self.high_level_hz), - policy=self.policy, - observation_provider=self.observation_provider, - ), - # Listens for the ``subtask_change`` event raised by - # ``HighLevelSubtaskFwd`` and refreshes ``current_memory``. - MemoryUpdateFwd( - trigger=EventTrigger("subtask_change"), - policy=self.policy, - observation_provider=self.observation_provider, - ), - AskVQAFwd( - policy=self.policy, - observation_provider=self.observation_provider, - ), - LowLevelForward( - trigger=HzTrigger(self.chunk_hz), - policy=self.policy, - observation_provider=self.observation_provider, - ), - DispatchAction( - trigger=HzTrigger(self.ctrl_hz), - robot_executor=self.robot_executor, - ), - DispatchToolCalls(tools=self.tools), - ] - self.state = initial_runtime_state() - - # ------------------------------------------------------------------ - # Lifecycle - # ------------------------------------------------------------------ - - def set_task(self, task: str) -> None: - """Set or replace the active task. Logged for the REPL.""" - self.state["task"] = task - push_log(self.state, f"Task: {task}") - - def stop(self) -> None: - self._stop = True - - def run(self, *, max_ticks: int | None = None) -> None: - """Main loop. Returns when ``stop()`` is called or after - ``max_ticks`` ticks (useful for tests / dry-run).""" - clock = TickClock(max_rate_hz=self.max_rate_hz) - while not self._stop: - tick = clock.advance() - self.state["_tick"] = tick - self.state["events_this_tick"] = [] - self.state["log_lines"] = [] - - if self.event_collector is not None: - self.event_collector(self.state) - if self.state.get("stop"): - self._stop = True - break - - for step in self.pipeline: - self.state = step(self.state) - - self._flush_logs() - if max_ticks is not None and tick.index >= max_ticks: - break - - self._on_shutdown() - - # ------------------------------------------------------------------ - # REPL helper: drive one full pipeline pass and return its logs - # ------------------------------------------------------------------ - - def step_once(self) -> list[str]: - """Run one tick of the pipeline and return the log lines. - - Used by the interactive REPL: instead of a background thread, - the CLI drives ticks synchronously after each user input. Logs - are returned (not printed) so the caller can route them into - the rich-Live chat scrollback. - """ - from .triggers import Tick # noqa: PLC0415 - - # Synthesize a tick. We don't need the real wall-clock pacing - # here — the REPL drives the runtime, not vice versa — but - # ``HzTrigger`` uses ``tick.monotonic_seconds`` to gate, so we - # bump it generously so every Hz-triggered step considers - # itself due. - import time as _time # noqa: PLC0415 - - prev_index = self.state.get("_tick").index if isinstance(self.state.get("_tick"), Tick) else 0 - self.state["_tick"] = Tick(index=prev_index + 1, monotonic_seconds=_time.monotonic()) - self.state["log_lines"] = [] - # ``events_this_tick`` is set up by the caller before - # ``step_once`` (the REPL pushes user-driven events first). - self.state.setdefault("events_this_tick", []) - - for step in self.pipeline: - self.state = step(self.state) - - return list(self.state.get("log_lines") or []) - - # ------------------------------------------------------------------ - # I/O - # ------------------------------------------------------------------ - - def _flush_logs(self) -> None: - for line in self.state.get("log_lines") or []: - print(f"[pi052] {line}", flush=True) - - def _on_shutdown(self) -> None: - # Drain any queued action chunks safely. - queue = self.state.get("action_queue") - if isinstance(queue, deque): - queue.clear() - print("[pi052] runtime stopped", flush=True) +__all__ = [ + "LanguageConditionedRuntime", + "PI052PolicyAdapter", + "PI052Runtime", + "RuntimeState", + "Tick", + "TickClock", + "ToolCall", + "VQAResult", +] diff --git a/src/lerobot/policies/pi052/inference/runtime_cli.py b/src/lerobot/policies/pi052/inference/runtime_cli.py new file mode 100644 index 000000000..80b03cd9c --- /dev/null +++ b/src/lerobot/policies/pi052/inference/runtime_cli.py @@ -0,0 +1,1688 @@ +#!/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. +"""``lerobot-pi052-runtime`` — interactive REPL for trained PI052. + +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. + +Examples +-------- + +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 \\ + --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:: + + uv run lerobot-pi052-runtime \\ + --policy.path=pepijn223/pi052_hirobot_super_poulain_tool2 \\ + --dataset.repo_id=pepijn223/super_poulain_annotated \\ + --dataset.episode=0 \\ + --no_robot \\ + --task="please clean the kitchen" + +With a real robot:: + + uv run lerobot-pi052-runtime \\ + --policy.path=... \\ + --robot.type=so101 --robot.port=/dev/tty.usbmodem... \\ + --tts.voice=alba + +``--policy.path`` accepts either a local directory or a Hugging Face +Hub repo id. ``--dataset.repo_id`` likewise. + +Tool dispatch (TTS via ``SayTool``) is enabled by default when +``pocket-tts`` is installed; pass ``--no_tts`` to disable. +""" + +from __future__ import annotations + +import argparse +import logging +import sys +from collections.abc import Callable +from contextlib import suppress +from typing import Any + +logger = logging.getLogger("lerobot.pi052.runtime") + + +def _emit(state: Any, event_name: str) -> None: + if hasattr(state, "emit"): + state.emit(event_name) + else: + state.setdefault("events_this_tick", []).append(event_name) + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + p = argparse.ArgumentParser( + description=("Interactive REPL runtime for a trained PI052 hierarchical VLA checkpoint."), + ) + 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``." + ), + ) + p.add_argument( + "--dataset.repo_id", + dest="dataset_repo_id", + type=str, + default=None, + help=( + "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 " + "trained on." + ), + ) + p.add_argument( + "--dataset.episode", + dest="dataset_episode", + type=int, + default=0, + help="Episode index to walk through (default: 0).", + ) + p.add_argument( + "--dataset.start_frame", + dest="dataset_start_frame", + type=int, + default=0, + help="Frame index within the episode to start from (default: 0).", + ) + p.add_argument( + "--dataset.advance_per_tick", + dest="dataset_advance_per_tick", + type=int, + default=1, + help=( + "How many dataset frames to advance per runtime tick. The " + "default of 1 means the runtime walks the episode forward " + "frame by frame; set to 0 to freeze on ``start_frame``." + ), + ) + p.add_argument( + "--dataset.augment_at_inference", + dest="dataset_augment_at_inference", + action="store_true", + help=( + "Apply the same torchvision-v2 ColorJitter / SharpnessJitter " + "/ RandomAffine pipeline that training used to each dataset " + "frame fed to the policy. Use to test whether the LM head " + "generalises under the augmentation distribution it was " + "supervised on — if dry-run still produces coherent subtask " + "text with this flag on, the head has learned beyond exact " + "frames; if it collapses to '\\n' the head is hyper-specific " + "to the unperturbed training samples." + ), + ) + p.add_argument( + "--task", + dest="task", + type=str, + default=None, + help=( + "Initial task. When given, the startup task picker is skipped " + "and this task is used directly. If omitted, the picker is " + "shown (or the first stdin line is treated as the task)." + ), + ) + p.add_argument( + "--mode", + dest="mode", + type=str, + choices=["action", "paused"], + default=None, + help=( + "Start-up run mode. 'action' runs the robot immediately on " + "--task; 'paused' (the default) comes up at the command line " + "with the robot idle. Flip any time with /action and /pause." + ), + ) + p.add_argument( + "--no_robot", + action="store_true", + help="Skip robot connection — language-only / dry-run mode.", + ) + # --- Real-robot mode args ---------------------------------------- + # Setting ``--robot.type`` flips the runtime into autonomous mode: + # it connects to the robot, builds an observation provider that + # reads ``robot.get_observation()`` instead of dataset frames, and + # 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. + p.add_argument( + "--robot.type", + dest="robot_type", + type=str, + default=None, + help=( + "Robot config choice (e.g. ``so101``, ``so101_follower``). " + "When set, the runtime drives the actual robot at " + "``--ctrl_hz`` instead of running the dataset-driven dry-run " + "REPL. Implies ``--autonomous`` unless ``--no_robot`` is also " + "passed (in which case the flag is ignored). See " + "``lerobot.robots`` for available choices." + ), + ) + p.add_argument( + "--robot.port", + dest="robot_port", + type=str, + default=None, + help="Serial port for the robot (e.g. ``/dev/tty.usbmodem...``).", + ) + p.add_argument( + "--robot.id", + dest="robot_id", + type=str, + default=None, + help="Optional robot identifier (passed through to ``RobotConfig.id``).", + ) + p.add_argument( + "--robot.cameras", + dest="robot_cameras", + type=str, + default=None, + help=( + "Optional JSON dict describing camera configs to attach to " + 'the robot (e.g. ``\'{"top": {"type": "opencv", "index": 0}}\'``). ' + "Camera keys MUST match the ``observation.images.*`` features " + "the policy was trained on." + ), + ) + p.add_argument( + "--robot.max_relative_target", + dest="robot_max_relative_target", + type=str, + default=None, + help=( + "Safety clip on per-motor relative motion, passed through to " + "``RobotConfig.max_relative_target``. Accepts either a float " + "(applied to every motor — e.g. ``5.0`` degrees) or a JSON " + "object mapping motor names to caps " + '(e.g. ``\'{"shoulder_pan": 5, "gripper": 30}\'``). The ' + "robot driver clips each commanded position relative to the " + "current measured position before sending — same kill-switch " + "``lerobot-record`` uses. Default ``None`` = no clipping." + ), + ) + p.add_argument( + "--auto_start", + action="store_true", + help=( + "Skip the ``Press ENTER to start`` confirmation prompt before " + "the autonomous control loop begins. Off by default — having " + "to confirm catches a lot of stupid mistakes (wrong policy, " + "wrong robot, robot not at home pose)." + ), + ) + p.add_argument( + "--no_tts", + action="store_true", + help="Disable the ``say`` tool dispatch.", + ) + p.add_argument( + "--tts.voice", + dest="tts_voice", + type=str, + default="alba", + help="Pocket-tts voice name (or path to a .wav for cloning).", + ) + p.add_argument( + "--chunk_hz", + type=float, + default=1.0, + help=( + "Action-chunk generation rate (Hz). Default ``1.0`` — one " + "new chunk per second. Lower = less inference cost / " + "smoother behaviour but longer reaction time to changes. " + "Higher = fresher actions / more inference cost; cap at " + "~1/(forward-pass latency)." + ), + ) + p.add_argument("--ctrl_hz", type=float, default=50.0, help="Action dispatch rate.") + p.add_argument( + "--high_level_hz", + type=float, + default=1.0, + help="High-level subtask generation rate.", + ) + p.add_argument( + "--subtask_chunks_per_gen", + type=int, + default=1, + help=( + "Throttle subtask gen to once every N action-chunk boundaries. " + "Default 1 = regenerate the subtask on every chunk refresh. " + "Set to 5 to run ~5 flow-matching action chunks per LM-head " + "subtask gen — saves compute and avoids re-planning trajectories " + "mid-grasp when a subtask is still valid across multiple chunks." + ), + ) + p.add_argument( + "--max_ticks", + type=int, + default=None, + help="Stop after N ticks (debug / smoke-test).", + ) + p.add_argument( + "--text_min_new_tokens", + type=int, + default=0, + help=( + "Debug knob for under-trained checkpoints: force the LM head " + "to emit at least N non-EOS tokens before EOS is allowed. " + "Use when the head's prior at position 0 still favours EOS " + "(short training run on a chat-pretrained backbone). 3-5 " + "is usually enough to reveal whether the model has real " + "subtask-token mass under the EOS argmax." + ), + ) + p.add_argument( + "--text_temperature", + type=float, + default=0.0, + help=( + "Sampling temperature for high-level text gen. 0 = greedy " + "argmax (default, matches training). Set 0.3-0.7 with an " + "under-trained checkpoint to escape stuck-at-EOS argmax." + ), + ) + p.add_argument( + "--text_top_p", + type=float, + default=1.0, + help="Nucleus filtering for high-level text gen.", + ) + p.add_argument("-v", "--verbose", action="store_true", help="Enable DEBUG logging.") + return p.parse_args(argv) + + +# Columns the runtime supplies itself via its own message stream — strip +# them so ``RenderMessagesStep`` / ``PI052TextTokenizerStep`` are no-ops. +_RUNTIME_OWNED_LANGUAGE_COLS = ("language_persistent", "language_events") + + +def _strip_runtime_owned_language_cols(sample: dict) -> None: + """In-place drop of language columns the runtime owns at inference.""" + for k in _RUNTIME_OWNED_LANGUAGE_COLS: + sample.pop(k, None) + + +def _select_observation_to_device(sample: dict, device: Any) -> dict: + """Filter to ``observation.*`` keys and move tensors to ``device``.""" + import torch # noqa: PLC0415 + + return { + k: v.to(device) if isinstance(v, torch.Tensor) else v + for k, v in sample.items() + if isinstance(k, str) and k.startswith("observation.") + } + + +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). + + Returns ``(policy, preprocessor, postprocessor, ds_meta)``. + ``preprocessor`` / ``postprocessor`` / ``ds_meta`` are ``None`` + when no dataset is provided (rare — needed for autonomous robot + mode to have action-denormalisation stats). + """ + from lerobot.configs import PreTrainedConfig # noqa: PLC0415 + from lerobot.policies.factory import make_policy, make_pre_post_processors # noqa: PLC0415 + + cfg = PreTrainedConfig.from_pretrained(policy_path) + cfg.pretrained_path = policy_path + + ds_meta = None + preprocessor = None + postprocessor = None + if dataset_repo_id is not None: + from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata # noqa: PLC0415 + + ds_meta = LeRobotDatasetMetadata(dataset_repo_id) + policy = make_policy(cfg, ds_meta=ds_meta) + # ``pretrained_path=None`` rebuilds fresh — the saved + # ``policy_preprocessor.json`` doesn't round-trip + # ``RenderMessagesStep.recipe``. Stats come from the dataset + # the user is feeding through, so normalisation is consistent. + preprocessor, postprocessor = make_pre_post_processors( + cfg, + pretrained_path=None, + dataset_stats=ds_meta.stats, + ) + else: + from lerobot.policies.factory import get_policy_class # noqa: PLC0415 + + policy_cls = get_policy_class(cfg.type) + policy = policy_cls.from_pretrained(policy_path, config=cfg) + policy.to(cfg.device) + + policy.eval() + return policy, preprocessor, postprocessor, ds_meta + + +def _build_observation_provider( + *, + dataset_repo_id: str, + episode: int, + start_frame: int, + advance_per_tick: int, + preprocessor: Any, + device: str, + augment: bool = False, +) -> Callable[[], dict | None]: + """Build a closure that feeds dataset frames into the runtime. + + Each call returns a preprocessed observation batch (images + + state, batched, on the policy's device, normalized) suitable for + ``policy.select_action`` and ``policy.select_message``. The + closure walks the chosen episode forward by ``advance_per_tick`` + frames per call, looping back to the episode start when it falls + off the end. + + The dataset's ``language_persistent`` / ``language_events`` + columns are stripped before the sample reaches the preprocessor, + so ``RenderMessagesStep`` and ``PI052TextTokenizerStep`` are + no-ops; the runtime supplies its own messages from current state. + """ + from lerobot.datasets.lerobot_dataset import LeRobotDataset # noqa: PLC0415 + + ds = LeRobotDataset(dataset_repo_id, episodes=[episode]) + if len(ds) == 0: + raise ValueError(f"Dataset {dataset_repo_id!r} episode {episode} is empty.") + + # Optional: apply the same torchvision-v2 augmentation pipeline + # that training used, so dry-run sees frames from the augmented + # support region (not just the unperturbed dataset frames). When + # the LM head still generates coherent text under this, it has + # learned over the augmentation distribution — the *opposite* of + # the "memorised one specific frame per supervision" failure + # mode. When it collapses to ``\n`` here too, the head is hyper- + # specific to the unperturbed training samples and only the + # retrain can help. + inference_aug = None + if augment: + from lerobot.transforms import ( # noqa: PLC0415 + ImageTransforms, + ImageTransformsConfig, + ) + + aug_cfg = ImageTransformsConfig(enable=True) + inference_aug = ImageTransforms(aug_cfg) + ds.set_image_transforms(inference_aug) + logger.warning( + "dry-run augmentation ENABLED — frames will be jittered " + "(brightness/contrast/saturation/hue/sharpness/affine) " + "before going to the policy" + ) + + state = {"cursor": max(0, min(start_frame, len(ds) - 1))} + + def _provider() -> dict | None: + idx = state["cursor"] + if advance_per_tick > 0: + state["cursor"] = (idx + advance_per_tick) % len(ds) + + sample = ds[idx] + _strip_runtime_owned_language_cols(sample) + + if preprocessor is not None: + sample = preprocessor(sample) + + return _select_observation_to_device(sample, device) + + return _provider + + +def _bootstrap_state_from_dataset( + *, + dataset_repo_id: str, + episode: int, + start_frame: int, +) -> dict[str, str]: + """Pull task / active plan / active memory / active subtask at ``start_frame``. + + The model is heavily memorised on the exact training prompts the + recipe rendered from this dataset (canonical task wording, + persistent atoms emitted earlier in the episode). Reconstructing + that state at REPL startup lets the runtime's first prompt line + up with what training looked like — without it the model sees an + out-of-distribution prompt and falls back to its dominant + training mode (VQA JSON spam). + """ + from lerobot.datasets.lerobot_dataset import LeRobotDataset # noqa: PLC0415 + + ds = LeRobotDataset(dataset_repo_id, episodes=[episode]) + if len(ds) == 0: + return {} + idx = max(0, min(start_frame, len(ds) - 1)) + sample = ds[idx] + + out: dict[str, str] = {} + task = sample.get("task") + if isinstance(task, str) and task.strip(): + out["task"] = task + + persistent = sample.get("language_persistent") or [] + # ``persistent`` is the broadcast slice of the episode; pick the + # *latest* row of each style whose ``timestamp`` is ≤ the + # frame's timestamp (matches the renderer's ``active_at`` + # semantics). + try: + frame_ts = ( + float(sample["timestamp"]) + if not hasattr(sample["timestamp"], "item") + else sample["timestamp"].item() + ) + except Exception: # noqa: BLE001 + frame_ts = float("inf") + + by_style: dict[str, tuple[float, str]] = {} + for row in persistent: + style = row.get("style") + ts = row.get("timestamp") + content = row.get("content") + if not (style and content) or ts is None: + continue + try: + ts_f = float(ts) + except (TypeError, ValueError): + continue + if ts_f > frame_ts: + continue + prev = by_style.get(style) + if prev is None or ts_f >= prev[0]: + by_style[style] = (ts_f, content) + for style, (_, content) in by_style.items(): + if style in {"plan", "memory", "subtask"}: + out[style] = content + return out + + +def _select_task_interactively( + *, + ds_meta: Any, + bootstrap_task: str | None, +) -> str | None: + """Ask the operator which task to run at startup. + + Behaviour: + + * If a dataset is loaded, build a numbered menu of every unique task + string in ``ds_meta.tasks`` (canonical bootstrap task listed first + as the default). Add a ``[c] type a custom task`` option. + * If no dataset is loaded, show a plain ``Enter task:`` prompt. + * Non-TTY runs (scripts, pipes) skip the prompt and return the + bootstrap task so the existing "first stdin line becomes task" + flow in ``_run_repl`` / ``_run_autonomous`` still works. + + Returns the chosen task string, or ``None`` when the operator declines + to pick one (Ctrl-D / empty + no default). + """ + options: list[str] = [] + seen: set[str] = set() + if bootstrap_task: + options.append(bootstrap_task) + seen.add(bootstrap_task) + if ds_meta is not None and getattr(ds_meta, "tasks", None) is not None: + try: + for t in list(ds_meta.tasks.index): + if isinstance(t, str) and t and t not in seen: + options.append(t) + seen.add(t) + except Exception as exc: # noqa: BLE001 — defensive: tasks shape varies + logger.debug("could not enumerate dataset tasks: %s", exc) + + if not sys.stdin.isatty(): + # Scripted / piped run: no interactive prompt; fall back to the + # bootstrap default (may be None — REPL handles that). + return bootstrap_task + + print("\n[pi052] Select startup task:", flush=True) + if options: + for i, opt in enumerate(options, 1): + marker = " (dataset default)" if opt == bootstrap_task else "" + print(f" [{i}] {opt}{marker}", flush=True) + print(" [c] type a custom task", flush=True) + prompt = "Choice [1]: " if bootstrap_task else "Choice: " + else: + print(" (no tasks available from dataset)", flush=True) + prompt = "Enter task: " + + while True: + try: + choice = input(prompt).strip() + except EOFError: + print(flush=True) + return bootstrap_task + + # No dataset options at all: the entered line *is* the task. + if not options: + return choice or None + + # Empty input: take the default (item 1) when there is one. + if not choice: + return options[0] if bootstrap_task else None + + if choice.lower() in ("c", "custom"): + try: + free = input("Enter task: ").strip() + except EOFError: + print(flush=True) + return bootstrap_task + if free: + return free + # Empty free-form input → loop back to the menu. + continue + + if choice.isdigit(): + idx = int(choice) + if 1 <= idx <= len(options): + return options[idx - 1] + + print( + f" invalid choice {choice!r}; pick 1–{len(options)} or 'c'.", + flush=True, + ) + + +def _build_robot( + *, + robot_type: str, + robot_port: str | None, + robot_id: str | None, + robot_cameras_json: str | None, + robot_max_relative_target: str | None, +): + """Build and connect a robot from CLI args. + + Mirrors how ``lerobot-record`` builds a robot but takes the args + flat from argparse instead of through draccus, so the runtime + keeps its plain ``--key=value`` CLI surface. ``max_relative_target`` + is passed through to the RobotConfig — the driver itself clips each + commanded joint position relative to the current measured one + before issuing it on the bus. + """ + import importlib # noqa: PLC0415 + import json # noqa: PLC0415 + import pkgutil # noqa: PLC0415 + + import lerobot.robots as _robots_pkg # noqa: PLC0415 + from lerobot.robots import ( # noqa: PLC0415 + RobotConfig, + make_robot_from_config, + ) + + # ``RobotConfig._choice_registry`` is populated lazily — each robot's + # ``config_.py`` calls ``@RobotConfig.register_subclass`` at + # import time. ``lerobot.robots/__init__.py`` doesn't import the + # individual robot packages, so ``get_choice_class(robot_type)`` + # raises ``KeyError`` until at least one robot module has been + # imported. Mirror what ``make_robot_from_config`` does internally: + # walk the robots package's submodules and import each so the + # decorator side-effect runs. Slow only on the first call (~200ms + # for ~10 dataclass modules); negligible for an autonomous run that + # then loops at ctrl_hz for minutes. + for _modinfo in pkgutil.iter_modules(_robots_pkg.__path__): + if _modinfo.name.startswith("_"): + continue + try: + importlib.import_module(f"lerobot.robots.{_modinfo.name}") + except Exception as exc: # noqa: BLE001 + logger.debug("could not import lerobot.robots.%s: %s", _modinfo.name, exc) + + try: + cls = RobotConfig.get_choice_class(robot_type) + except KeyError as exc: + available = sorted(RobotConfig._choice_registry.keys()) + raise ValueError(f"Unknown robot type {robot_type!r}. Available choices: {available}") from exc + kwargs: dict[str, Any] = {} + if robot_port: + kwargs["port"] = robot_port + if robot_id: + kwargs["id"] = robot_id + if robot_cameras_json: + try: + cameras_raw = json.loads(robot_cameras_json) + except json.JSONDecodeError as exc: + raise ValueError( + f"--robot.cameras must be a JSON object, got {robot_cameras_json!r}: {exc}" + ) from exc + # ``RobotConfig`` expects ``cameras: dict[str, CameraConfig]`` — + # each inner value must be an actual ``CameraConfig`` subclass + # instance, not a raw dict. Look up the matching subclass via + # ``CameraConfig.get_choice_class()`` (registered by + # ``@CameraConfig.register_subclass`` decorators on each camera + # backend's config) and instantiate it. Mirror the lazy-import + # pattern from above so the registry is populated. + import lerobot.cameras as _cameras_pkg # noqa: PLC0415 + from lerobot.cameras import CameraConfig # noqa: PLC0415 + + for _modinfo in pkgutil.iter_modules(_cameras_pkg.__path__): + if _modinfo.name.startswith("_"): + continue + try: + importlib.import_module(f"lerobot.cameras.{_modinfo.name}") + except Exception as exc: # noqa: BLE001 + logger.debug("could not import lerobot.cameras.%s: %s", _modinfo.name, exc) + + cameras: dict[str, Any] = {} + for cam_name, cam_dict in cameras_raw.items(): + if not isinstance(cam_dict, dict): + raise ValueError(f"camera {cam_name!r} value must be a dict, got {cam_dict!r}") + cam_dict = dict(cam_dict) # don't mutate caller's parsed JSON + cam_type = cam_dict.pop("type", None) + if cam_type is None: + raise ValueError( + f"camera {cam_name!r} is missing a 'type' field (e.g. 'opencv', 'intelrealsense')" + ) + try: + cam_cls = CameraConfig.get_choice_class(cam_type) + except KeyError as exc: + available = sorted(CameraConfig._choice_registry.keys()) + raise ValueError( + f"camera {cam_name!r}: unknown type {cam_type!r}. Available choices: {available}" + ) from exc + cameras[cam_name] = cam_cls(**cam_dict) + kwargs["cameras"] = cameras + if robot_max_relative_target: + # Accept either a bare float (uniform cap) or a JSON object + # (per-motor cap). Matches ``RobotConfig.max_relative_target``'s + # ``float | dict[str, float] | None`` shape. + s = robot_max_relative_target.strip() + try: + if s.startswith("{"): + kwargs["max_relative_target"] = json.loads(s) + else: + kwargs["max_relative_target"] = float(s) + except (json.JSONDecodeError, ValueError) as exc: + raise ValueError( + f"--robot.max_relative_target must be a float or JSON dict, " + f"got {robot_max_relative_target!r}: {exc}" + ) from exc + cfg = cls(**kwargs) + robot = make_robot_from_config(cfg) + robot.connect() + return robot + + +def _build_robot_observation_provider( + *, + robot, + preprocessor: Any, + device: str, + task: str | None, + ds_features: dict[str, Any] | None, +) -> Callable[[], dict | None]: + """Closure that reads from the robot, runs the policy preprocessor. + + Each call: ``robot.get_observation()`` (raw per-joint + per-camera + dict, possibly with scalar floats) → ``build_inference_frame`` + (extract the keys the dataset declared, reshape per-joint floats + into a single ``observation.state`` vector, prefix camera keys + with ``observation.images.``, convert to tensors with batch dim + on device) → wrap in an ``EnvTransition`` (the preprocessor + pipeline is transition-shaped, keyed by ``TransitionKey``) → + preprocessor (rename, normalise) → unwrap and return the flat + observation batch ``policy.select_action`` / ``policy.select_message`` + consume. + """ + import torch # noqa: PLC0415 + + from lerobot.policies.utils import ( # noqa: PLC0415 + build_inference_frame, + prepare_observation_for_inference, + ) + + torch_device = torch.device(device) if isinstance(device, str) else device + robot_type = getattr(robot, "robot_type", None) or getattr(getattr(robot, "config", None), "type", None) + + # Pre-compute the camera-key → target (H, W) map from + # ``ds_features``. The training distribution sees frames at the + # recorded resolution (e.g. 480×640); a live Mac/USB camera will + # almost always hand us a different native size (720p / 1080p). + # PI052's internal ``resize_with_pad(512, 512)`` does pad the + # input to a fixed canvas, but the *geometry* of that pad differs + # by input aspect ratio — top/left padding varies, so the visual + # tokens at each tile carry different content than what the model + # saw at training. The action expert tolerates this (flow head + # rides broad geometry); the LM head, supervised much more + # tightly on visual features, goes out of distribution and the + # head's distribution at position 0 collapses to its dominant + # mode (a memorised ``\n``-only run in this checkpoint). + _resize_logged = {"done": False} + target_image_shapes: dict[str, tuple[int, int]] = {} + if ds_features: + for fkey, fmeta in ds_features.items(): + if not isinstance(fmeta, dict): + continue + dtype = fmeta.get("dtype") + if dtype not in ("image", "video"): + continue + shape = fmeta.get("shape") + if not shape or len(shape) != 3: + continue + names = fmeta.get("names") or [] + # Feature schema stores either (H, W, C) or (C, H, W); + # disambiguate by the ``names`` ordering when present. + if names and len(names) == 3 and names[0] == "channels": + _, h, w = shape + else: + h, w, _ = shape + cam_key = fkey.removeprefix("observation.images.") + target_image_shapes[cam_key] = (int(h), int(w)) + + def _provider() -> dict | None: + try: + raw = robot.get_observation() + except Exception as exc: # noqa: BLE001 + logger.warning("robot.get_observation failed: %s", exc) + return None + + # The runtime supplies messages itself; strip any language + # columns the robot stream may carry through. + _strip_runtime_owned_language_cols(raw) + + # Force-match the training-time visual distribution: + # every camera frame the model trained on came from the + # dataset at its recorded (H, W). Resize the live frame to + # that exact shape so the downstream resize_with_pad geometry + # matches training. Without this the LM head is OOD on every + # tick. + if target_image_shapes: + try: + import cv2 as _cv2 # noqa: PLC0415 + import numpy as _np # noqa: PLC0415 + + # Snapshot the gate state at the start of the call: the + # camera info and startup-state warnings are meant to fire + # exactly once (operator sanity check), so gate them on + # the *previous* value rather than the post-loop value. + first_call = not _resize_logged["done"] + for cam_key, (target_h, target_w) in target_image_shapes.items(): + img = raw.get(cam_key) + if img is None or not isinstance(img, _np.ndarray): + continue + if img.ndim != 3: + continue + cur_h, cur_w = img.shape[:2] + if first_call: + logger.warning( + "camera %s: live=%dx%d, training=%dx%d (resize=%s)", + cam_key, + cur_h, + cur_w, + target_h, + target_w, + "yes" if (cur_h, cur_w) != (target_h, target_w) else "no — already matched", + ) + if (cur_h, cur_w) == (target_h, target_w): + continue + raw[cam_key] = _cv2.resize(img, (target_w, target_h), interpolation=_cv2.INTER_AREA) + _resize_logged["done"] = True + # Print the state vector once so the operator can eyeball + # it against the dataset's stats. State OOD is a real + # failure mode for VLAs — the prefix carries state via + # the projection layer, and a neutral home pose can + # easily sit a couple σ off the supervised support + # region. Gated on ``first_call`` so this doesn't spam + # every observation tick. + if first_call and "observation.state" in (ds_features or {}): + state_names = ds_features["observation.state"].get("names") or [] + state_vals = [raw.get(n) for n in state_names] + logger.warning( + "robot state at startup: %s", + { + n: round(v, 2) if isinstance(v, float) else v + for n, v in zip(state_names, state_vals, strict=False) + }, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("camera resize to dataset shape failed: %s", exc) + + try: + if ds_features: + # Use the dataset's feature schema to pick the right + # raw keys and fold per-joint scalars into a single + # ``observation.state`` tensor. Then tensor-ise + + # device-place + add batch dim. + obs_tensors = build_inference_frame( + raw, + torch_device, + ds_features=ds_features, + task=task, + robot_type=robot_type, + ) + else: + # No dataset features available — fall back to the + # generic numpy-only path; only works when the robot + # already returns dataset-shaped keys. + obs_tensors = prepare_observation_for_inference( + raw, + torch_device, + task=task, + robot_type=robot_type, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("observation prep failed: %s", exc) + return None + + if preprocessor is not None: + # ``PolicyProcessorPipeline`` defaults its ``to_transition`` + # to ``batch_to_transition``, which expects a *flat batch + # dict* keyed by ``observation.*`` / ``action`` / etc., and + # wraps it into an ``EnvTransition`` itself. Pre-wrapping + # here would just have ``batch_to_transition`` look for + # ``observation.*`` keys at top level, find none (they'd + # be nested under ``TransitionKey.OBSERVATION``), and + # produce an empty observation → ``ObservationProcessorStep`` + # bails. Pass the flat dict straight in; ``to_output`` + # gives us a flat dict back. + try: + processed = preprocessor(obs_tensors) + except Exception as exc: # noqa: BLE001 + logger.warning("preprocessor failed on robot observation: %s", exc) + return None + obs_tensors = processed if isinstance(processed, dict) else {} + + return _select_observation_to_device(obs_tensors, torch_device) + + return _provider + + +def _build_robot_action_executor( + *, + robot, + postprocessor: Any, + ds_meta: Any, +) -> Callable[[Any], None]: + """Closure that postprocesses an action and dispatches to the robot. + + Mirrors ``lerobot-record``'s ``predict_action`` tail: postprocess + (denormalise) → ``make_robot_action`` (tensor → ``{joint: value}`` + dict) → ``robot.send_action(...)``. Safety clipping happens *inside* + ``robot.send_action`` via the driver's ``max_relative_target`` + cap (passed in at ``RobotConfig`` construction time) — same place + ``lerobot-record`` enforces it. + """ + import torch # noqa: PLC0415 + + from lerobot.policies.utils import make_robot_action # noqa: PLC0415 + + def _executor(action: Any) -> None: + try: + if postprocessor is not None: + action = postprocessor(action) + if isinstance(action, torch.Tensor): + if action.ndim > 1 and action.shape[0] == 1: + action = action.squeeze(0) + action_dict = make_robot_action(action, ds_meta.features) + elif isinstance(action, dict): + action_dict = action + else: + logger.warning("unsupported action type %r — skipping", type(action)) + return + robot.send_action(action_dict) + except Exception as exc: # noqa: BLE001 + logger.error("robot.send_action failed: %s", exc, exc_info=True) + + return _executor + + +def _print_runtime_help() -> None: + """Print the slash-command reference.""" + print( + "[pi052] commands (arguments need no quotes):\n" + " /action run the robot; an argument switches to that task\n" + " /action resume the robot on the current task\n" + " /action run the robot for N seconds, then auto-pause\n" + " /pause pause the action loop — robot holds position\n" + " /question 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", + flush=True, + ) + + +def _is_number(text: str) -> bool: + """True if ``text`` parses as a float (a ``/action`` duration arg).""" + try: + float(text) + return True + except ValueError: + return False + + +def _strip_quotes(text: str) -> str: + """Strip one pair of surrounding quotes from a command argument.""" + text = text.strip() + if len(text) >= 2 and text[0] == text[-1] and text[0] in {'"', "'"}: + return text[1:-1].strip() + return text + + +def _clear_action_queue(runtime: Any) -> None: + """Drop any queued action chunk so nothing fires while paused.""" + queue = runtime.state.get("action_queue") + if hasattr(queue, "clear"): + queue.clear() + + +def _handle_slash_command(runtime: Any, line: str) -> bool: + """Dispatch the runtime slash commands. + + ``/action ["task"]`` run the robot; a quoted/bare argument sets a + new task, a bare number is a timed burst + (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). + """ + stripped = line.strip() + if not stripped.startswith("/"): + return False + head, _, rest = stripped.partition(" ") + cmd = head.lower() + rest = _strip_quotes(rest) + + if cmd in {"/action", "/act", "/run"}: + runtime.state["mode"] = "action" + if rest and _is_number(rest): + import time as _time # noqa: PLC0415 + + secs = float(rest) + runtime.state["action_deadline"] = _time.monotonic() + secs + print( + f"[pi052] action — running {secs:g}s, then auto-pause", + flush=True, + ) + else: + runtime.state["action_deadline"] = None + if rest: + runtime.set_task(rest) + # 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) + elif runtime.state.get("task"): + print( + f"[pi052] action — resuming: {runtime.state['task']!r}", + flush=True, + ) + else: + runtime.state["mode"] = "paused" + print( + "[pi052] no task set — use /action ", + flush=True, + ) + return True + + if cmd in {"/pause", "/p"}: + 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 (e.g. /question point to the yellow cube)", + flush=True, + ) + return True + _run_vqa_query(runtime, rest) + return True + + if cmd in {"/help", "/?"}: + _print_runtime_help() + return True + 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, + *, + robot, + auto_start: bool, + initial_task: str | None, + max_ticks: int | None, +) -> int: + """Drive the runtime continuously at ``ctrl_hz`` while accepting + stdin events in the foreground. + + Different from ``_run_repl`` (dataset dry-run): the policy needs + to keep generating action chunks at ``chunk_hz`` and dispatching + them at ``ctrl_hz`` regardless of whether the user is typing, so + ``runtime.run()`` runs in a background thread and stdin handling + happens here in the main thread. + """ + import threading # noqa: PLC0415 + import time # noqa: PLC0415 + + # Only gate on ENTER when the robot will actually move at startup + # (``--mode=action``). The default is paused — the command line + # comes up immediately and nothing moves until ``/action``. + 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. " + ) + except (EOFError, KeyboardInterrupt): + print("\n[pi052] aborted before start", flush=True) + return 130 + + if initial_task: + runtime.set_task(initial_task) + + thread = threading.Thread( + target=runtime.run, + kwargs={"max_ticks": max_ticks}, + name="pi052-runtime-loop", + daemon=True, + ) + thread.start() + + # Capture log lines flushed by the runtime each tick into a + # bounded scrollback that the panel renderer prints inside the + # rule block. Without this, ``runtime._flush_logs`` just calls + # ``print(...)`` which the 2 Hz panel redraw clears immediately — + # so failure messages from generation (e.g. ``[warn] subtask gen + # failed: ...``) flash for ≤ 0.5 s and disappear, leaving the + # operator with no idea why ``last_raw`` stays empty. + _scrollback: list[str] = [] + _scrollback_max = 12 + + def _flush_into_scrollback() -> None: + for line in runtime.state.get("log_lines") or []: + _scrollback.append(line) + # Trim to the cap so the panel doesn't grow unbounded. + if len(_scrollback) > _scrollback_max: + del _scrollback[: len(_scrollback) - _scrollback_max] + + runtime._flush_logs = _flush_into_scrollback # type: ignore[method-assign] + + redraw = _make_state_panel_renderer(runtime, mode_label="autonomous", scrollback=_scrollback) + redraw() + print( + " [autonomous] /action to run · /pause to stop · " + "/question to ask · /help · stop", + flush=True, + ) + + # Background panel-redraw thread so state changes from the runtime + # loop (subtask refresh, plan update, etc.) are visible without the + # user typing anything. + # + # In ``/vlm`` mode the action loop is paused — nothing changes in the + # background — so the timer redraw is suspended entirely. That keeps + # the screen stable while the operator types a VQA question and the + # interactive camera prompt, instead of the panel clearing the + # prompt every tick. + _panel_stop = threading.Event() + + def _panel_loop() -> None: + while not _panel_stop.is_set(): + st = runtime.state + if st.get("mode", "action") == "action": + # Timed burst (``/action ``): once the deadline + # passes, auto-revert to question mode and clear the + # queue so the robot stops. + deadline = st.get("action_deadline") + if deadline is not None and time.monotonic() >= deadline: + st["mode"] = "paused" + st["action_deadline"] = None + queue = st.get("action_queue") + if hasattr(queue, "clear"): + queue.clear() + print( + "\n[pi052] timed action elapsed — paused", + flush=True, + ) + else: + with suppress(Exception): + redraw() + # Re-print the prompt the redraw just cleared so + # the operator always has a visible ``> ``. + print("> ", end="", flush=True) + _panel_stop.wait(0.7) + + panel_thread = threading.Thread(target=_panel_loop, name="pi052-panel-redraw", daemon=True) + panel_thread.start() + + try: + while thread.is_alive(): + try: + line = input("> ").strip() + except EOFError: + break + if not line: + continue + lower = line.lower() + if lower in {"stop", "quit", "exit"}: + break + # The runtime is command-driven: /action "task", /pause, + # /question "...", /help. ``_handle_slash_command`` runs the + # VQA query inline for /question (the action loop is paused + # first, so the policy isn't in concurrent use). + if _handle_slash_command(runtime, line): + with suppress(Exception): + redraw() + continue + # A bare (non-slash) line is treated as a user interjection + # — the trained ``user_interjection_response`` path. ``stop`` + # already handled above; everything else routes here. + if runtime.state.get("task"): + runtime.state["recent_interjection"] = line + _emit(runtime.state, "user_interjection") + else: + print( + "[pi052] no task yet — use /action to start", + flush=True, + ) + except KeyboardInterrupt: + print("\n[pi052] interrupt — stopping", flush=True) + finally: + _panel_stop.set() + runtime.stop() + # Give the loop a moment to drain. + for _ in range(10): + if not thread.is_alive(): + break + time.sleep(0.1) + try: + robot.disconnect() + print("[pi052] robot disconnected", flush=True) + except Exception as exc: # noqa: BLE001 + print(f"[pi052] WARNING: robot.disconnect raised {exc}", flush=True) + + return 0 + + +def _make_state_panel_renderer( + runtime: Any, + *, + mode_label: str, + scrollback: list[str] | None = None, +) -> Callable[[list[str] | None], None]: + """Return a closure that prints the task/subtask/plan/memory panel. + + Used by both ``_run_repl`` (dry-run, called per user input) and + ``_run_autonomous`` (real robot, called on a 2 Hz timer + + whenever the user types). Centralises the visual format so the + two modes look identical. + """ + from rich.console import Console # noqa: PLC0415 + + console = Console(highlight=False) + + def _redraw(robot_lines: list[str] | None = None) -> None: + console.clear() + 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") + # 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[/] ask · [bold]/help[/] · [bold]stop[/]" + ) + else: + console.print( + " [dim]commands:[/] [bold]/action[/] run · " + "[bold]/question[/] ask · [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"), + ("current_plan", "plan"), + ("current_memory", "memory"), + ): + value = st.get(key) + if value: + console.print(f" [bold cyan]{label:<8}[/] {value}") + else: + console.print(f" [dim]{label:<8} (not set)[/]") + queue_len = ( + len(st["action_queue"]) + 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}[/]" + ) + + # Overfit / memorisation diagnostics. The high-level steps + # surface the raw generation each time they fire (even when + # rejected as gibberish or unchanged), plus repeat/rejection + # counters. Rule of thumb: + # + # * subtask repeat ≥ ~5 and queue_len cycles fully → model + # can't move past current subtask (memorised one phase + # of the task — classic overfit signature) + # * subtask gibberish climbing → LM head collapsed to + # chat-template fragments / one-token salads + # * last raw differs from accepted → at least the LM is + # varying, the gibberish filter is doing its job + raw_subtask = st.get("last_subtask_raw") + sub_rep = int(st.get("subtask_repeat_count") or 0) + sub_gib = int(st.get("subtask_gibberish_count") or 0) + sub_empty = int(st.get("subtask_empty_count") or 0) + if raw_subtask is not None or sub_rep or sub_gib or sub_empty: + raw_display = (raw_subtask or "(empty)")[:80] + color = "yellow" if (sub_rep >= 3 or sub_gib >= 3 or sub_empty >= 3) else "dim" + console.print( + f" [{color}]subtask diag repeat:{sub_rep} " + f"gibberish:{sub_gib} empty:{sub_empty} " + f"last_raw: {raw_display!r}[/]" + ) + + # Same diagnostics for memory and plan when available. + mem_gib = int(st.get("memory_gibberish_count") or 0) + plan_gib = int(st.get("plan_gibberish_count") or 0) + if mem_gib or plan_gib: + 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. + if scrollback: + for line in scrollback: + console.print(f" [magenta]{line.rstrip()}[/]") + console.rule(style="cyan") + if robot_lines: + for line in robot_lines: + console.print(f" [magenta]{line.strip()}[/]") + console.print() + if not st.get("task"): + console.print( + " [dim]Type [bold]/action [/bold] to begin, " + "[bold]/question [/bold] to ask, /help for commands, " + "stop to exit.[/]" + ) + + return _redraw + + +def _build_tools(no_tts: bool, tts_voice: str) -> dict[str, Any]: + """Instantiate the tools declared on this dataset/policy.""" + if no_tts: + return {} + try: + from lerobot.tools import SayTool # noqa: PLC0415 + + return {"say": SayTool(voice=tts_voice)} + except Exception as exc: # noqa: BLE001 + logger.warning("Could not initialise SayTool (%s) — speech disabled.", exc) + return {} + + +def _silence_noisy_loggers() -> None: + """Drop chatty third-party loggers down to WARNING. + + HuggingFace / httpx / urllib3 emit one log line per HTTP request, + which the REPL has to print between the state block and the + prompt — completely unreadable. We never need that detail in the + REPL and the user can opt back into it via ``-v`` (verbose mode + keeps DEBUG on the lerobot loggers but still gates the noisy ones + here unless they explicitly want them). + """ + for name in ( + "httpcore", + "httpcore.connection", + "httpcore.http11", + "httpcore.proxy", + "httpx", + "urllib3", + "urllib3.connectionpool", + "huggingface_hub", + "huggingface_hub.repocard", + "huggingface_hub.file_download", + "transformers", + "transformers.modeling_utils", + "transformers.tokenization_utils_base", + "datasets", + "filelock", + ): + logging.getLogger(name).setLevel(logging.WARNING) + + # The robot's relative-goal-position clamp warning fires *every* + # dispatch tick on a memorised model — the LM is trying to jump + # the wrist far past where max_relative_target allows, so the + # warning floods the panel at ~30 Hz. Promote it from WARNING to + # DEBUG: the dispatch counter on the panel already tells the + # operator the loop is running, and the panel itself shows + # whether motion is happening. If anyone needs the per-action + # clamp detail, ``-v`` puts it back via DEBUG. + logging.getLogger("lerobot.robots.utils").setLevel(logging.ERROR) + + +def main(argv: list[str] | None = None) -> int: + args = _parse_args(argv) + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(asctime)s %(levelname)s %(message)s", + ) + _silence_noisy_loggers() + + 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 " + "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) + policy, preprocessor, postprocessor, ds_meta = _load_policy_and_preprocessor( + args.policy_path, args.dataset_repo_id + ) + + # Bootstrap the canonical task from the dataset whenever one is + # provided, so the interactive picker below can offer it as the + # default. The model is memorised on the exact training wording, so + # matching it is what gets recall to fire. + bootstrap_state: dict[str, str] = {} + if args.dataset_repo_id is not None: + bootstrap_state = _bootstrap_state_from_dataset( + dataset_repo_id=args.dataset_repo_id, + episode=args.dataset_episode, + start_frame=args.dataset_start_frame, + ) + + # Interactive task picker. Skipped when ``--task`` is already set on + # the CLI (scripted runs and explicit overrides win). When no task + # was passed, prompt the operator: pick from the dataset's tasks or + # type a custom one. Non-TTY runs fall back to the bootstrap task + # silently — the existing "first stdin line becomes task" flow in + # ``_run_repl`` / ``_run_autonomous`` still handles the no-default + # case. + if not args.task: + chosen = _select_task_interactively( + ds_meta=ds_meta, + bootstrap_task=bootstrap_state.get("task"), + ) + if chosen: + args.task = chosen + print(f"[pi052] 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`` + # forces a mode. The operator drives it with /action, /pause and + # /question. + startup_mode = args.mode or "paused" + + observation_provider: Callable[[], dict | None] | None = None + robot_executor: Callable[[Any], None] | None = None + robot = None + + if autonomous_mode: + print( + f"[pi052] connecting to robot.type={args.robot_type} port={args.robot_port}", + flush=True, + ) + robot = _build_robot( + robot_type=args.robot_type, + robot_port=args.robot_port, + robot_id=args.robot_id, + robot_cameras_json=args.robot_cameras, + robot_max_relative_target=args.robot_max_relative_target, + ) + observation_provider = _build_robot_observation_provider( + robot=robot, + preprocessor=preprocessor, + device=str(getattr(policy.config, "device", "cpu")), + task=args.task, + ds_features=ds_meta.features if ds_meta is not None else None, + ) + robot_executor = _build_robot_action_executor( + robot=robot, + postprocessor=postprocessor, + ds_meta=ds_meta, + ) + elif args.dataset_repo_id is not None: + print( + f"[pi052] streaming observations from {args.dataset_repo_id} " + f"episode={args.dataset_episode} " + f"start_frame={args.dataset_start_frame}", + flush=True, + ) + observation_provider = _build_observation_provider( + dataset_repo_id=args.dataset_repo_id, + episode=args.dataset_episode, + start_frame=args.dataset_start_frame, + advance_per_tick=args.dataset_advance_per_tick, + preprocessor=preprocessor, + device=str(getattr(policy.config, "device", "cpu")), + augment=getattr(args, "dataset_augment_at_inference", False), + ) + + tools = _build_tools(args.no_tts, args.tts_voice) + if tools: + print(f"[pi052] tools loaded: {list(tools)}", flush=True) + + from lerobot.policies.pi052.inference import ( # noqa: PLC0415 + LanguageConditionedRuntime, + PI052PolicyAdapter, + ) + + runtime = LanguageConditionedRuntime( + policy_adapter=PI052PolicyAdapter(policy), + tools=tools, + observation_provider=observation_provider, + action_executor=robot_executor, + # No background event collector — the REPL drives ticks + # synchronously after each user input (REPL mode). Autonomous + # mode runs ``runtime.run()`` in a thread; stdin events are + # injected from the foreground. + event_collector=None, + chunk_hz=args.chunk_hz, + ctrl_hz=args.ctrl_hz, + high_level_hz=args.high_level_hz, + ) + # Stash text-gen knobs on the state dict so the high-level steps + # (which read state) can pick them up and forward them to + # policy.select_message. Letting the operator try + # ``--text_min_new_tokens=5 --text_temperature=0.6`` on an + # under-trained checkpoint without recompiling. + 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 + # 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 + # the previous one is still being executed. + runtime.state["subtask_chunks_per_gen"] = max(1, int(getattr(args, "subtask_chunks_per_gen", 1) or 1)) + # Apply the startup mode chosen above the task picker. + runtime.state["mode"] = startup_mode + if args.task: + runtime.set_task(args.task) + # Seed the current subtask from the dataset so the first chunk — + # before the adapter has generated one — has a real subtask to + # condition the action expert on instead of falling back to the + # bare task. Plan and memory are NOT seeded: the current recipe + # trains neither, no inference step consumes them, and seeding + # them only put a stale plan in the status panel that does + # nothing. + if bootstrap_state.get("subtask"): + runtime.state["current_subtask"] = bootstrap_state["subtask"] + + if autonomous_mode: + return _run_autonomous( + runtime, + robot=robot, + auto_start=args.auto_start, + initial_task=args.task, + max_ticks=args.max_ticks, + ) + # Fire one full pipeline tick at startup so the obs diagnostic + # *and* the subtask generation actually run before the REPL + # blocks on stdin. The REPL otherwise only ticks on user input, + # which made the dry-run bisection test (does the LM head produce + # text at start_frame=0?) require typing something. Doing + # ``step_once`` here means the diag row populates without any + # manual interaction. + if observation_provider is not None: + try: + startup_logs = runtime.step_once() + except Exception as exc: # noqa: BLE001 + 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) + + +def _run_repl(runtime: Any, *, initial_task: str | None, max_ticks: int | None) -> int: + """Claude-Code-style block REPL. + + Each turn redraws a status block (task / subtask / plan / memory) + at the top, prints any robot log lines that came in since the last + turn, then asks for input on a clean ``> `` prompt at the bottom. + No live region, no panel re-renders, no rendering races with HTTP + log lines — just clear-screen + reprint each turn, the way a + chat-style REPL is meant to look. + """ + try: + from rich.console import Console # noqa: PLC0415 + except ImportError: + print( + "[pi052] 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") + # Keep a local ``console`` just for the styled input prompt; the + # state panel is owned by the shared renderer. + console = Console(highlight=False) + + last_logs: list[str] = [] + _redraw() + if initial_task is None: + # Already shown the help line in _redraw when task is None. + pass + ticks_done = 0 + try: + while True: + try: + line = console.input("[bold cyan]> [/]").strip() + except EOFError: + break + if not line: + _redraw(last_logs) + continue + lower = line.lower() + if lower in {"stop", "quit", "exit"}: + break + + # Command-driven: /action "task", /pause, /question "...", + # /help. ``_handle_slash_command`` runs the VQA query inline + # for /question (single-threaded REPL — no concurrency). + if _handle_slash_command(runtime, line): + last_logs = list(runtime.state.get("log_lines") or []) + _redraw(last_logs) + ticks_done += 1 + if max_ticks is not None and ticks_done >= max_ticks: + break + continue + + # A bare (non-slash) line is a user interjection — needs a + # task to be meaningful. + if not runtime.state.get("task"): + print( + "[pi052] no task yet — use /action ", + flush=True, + ) + _redraw(last_logs) + continue + runtime.state["recent_interjection"] = line + _emit(runtime.state, "user_interjection") + + last_logs = runtime.step_once() or [] + _redraw(last_logs) + + ticks_done += 1 + if max_ticks is not None and ticks_done >= max_ticks: + break + except KeyboardInterrupt: + console.print("\n[dim]interrupted[/]") + console.print("[dim]runtime stopped[/]") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lerobot/policies/pi052/inference/runtime_state.py b/src/lerobot/policies/pi052/inference/runtime_state.py deleted file mode 100644 index f2ca6f03a..000000000 --- a/src/lerobot/policies/pi052/inference/runtime_state.py +++ /dev/null @@ -1,95 +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. -"""Runtime state passed between inference steps each tick. - -The runtime threads a single dict through the pipeline; this module -documents the shape and provides factories. We use a plain ``dict`` -rather than a frozen dataclass because steps freely add and remove -keys (``events_this_tick``, ``messages_pending``, ``tool_calls_pending``, -…) and dataclass field churn would just get in the way. - -Stable keys (read by multiple steps): - - task str the current top-level task - current_plan str | None latest plan emitted by the planner - current_subtask str | None latest subtask the policy is executing - current_memory str | None latest compressed memory - recent_interjection str | None most recent user interjection text (consumed) - - action_queue collections.deque[Tensor] pending action chunks - tool_calls_pending list[dict] parsed but not-yet-dispatched tool calls - - events_this_tick list[str] triggers consumed this tick - _tick Tick current tick (set by the loop) - - mode str "action" (run the robot) | "paused" - (action loop stopped — robot holds) - - log_lines list[str] human-readable status lines printed each tick -""" - -from __future__ import annotations - -from collections import deque -from typing import Any - - -def initial_runtime_state(task: str | None = None) -> dict[str, Any]: - """Build a fresh runtime state dict with sensible defaults.""" - return { - "task": task, - "current_plan": None, - "current_subtask": None, - "current_memory": None, - "recent_interjection": None, - "action_queue": deque(), - "tool_calls_pending": [], - "events_this_tick": [], - "log_lines": [], - "mode": "action", - "stop": False, - } - - -def take_event(state: dict[str, Any], event_name: str) -> bool: - """Pop ``event_name`` from ``events_this_tick`` if present. - - Steps that consume an event call this so the same event doesn't - re-fire on a sibling step within the same tick. - """ - events: list[str] = state.get("events_this_tick") or [] - if event_name in events: - events.remove(event_name) - return True - return False - - -def push_log(state: dict[str, Any], line: str) -> None: - """Append ``line`` to the per-tick log buffer; the runtime prints - it at the end of the tick.""" - state.setdefault("log_lines", []).append(line) - - -def set_if_changed(state: dict[str, Any], key: str, value: Any, label: str | None = None) -> bool: - """Update ``state[key]`` and log a diff line if the value changed. - - Returns ``True`` if the value actually changed. - """ - prev = state.get(key) - if prev == value: - return False - state[key] = value - if label is not None: - push_log(state, f" {label}: {value}") - return True diff --git a/src/lerobot/policies/pi052/inference/steps.py b/src/lerobot/policies/pi052/inference/steps.py index a9af3b7a1..60ead4ac4 100644 --- a/src/lerobot/policies/pi052/inference/steps.py +++ b/src/lerobot/policies/pi052/inference/steps.py @@ -11,930 +11,19 @@ # 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. -"""Inference steps for the PI052 multi-rate runtime. -Each step is a tiny class with a ``trigger`` and an ``__call__(state)``; -the runtime applies them in order each tick. When a step's trigger -doesn't fire, the step is a no-op and the runtime moves on. - -Stream-to-step mapping mirrors the PI052 training recipe: - -* ``LowLevelForward`` — calls ``policy.select_action`` for the - action chunk; trained by - ``low_level_execution`` -* ``EnqueueChunk`` — pushes the chunk to ``action_queue`` -* ``DispatchAction`` — pops one action per control tick and - forwards to the robot -* ``HighLevelSubtaskFwd`` — calls ``policy.select_message`` for the - next subtask; trained by - ``high_level_subtask`` -* ``MemoryUpdateFwd`` — fires on subtask boundary; trained by - ``memory_update`` -* ``UserInterjectionFwd`` — fires on stdin interjection; trained by - ``user_interjection_response`` -* ``AskVQAFwd`` — fires on stdin question; trained by - ``ask_vqa_*`` -* ``DispatchToolCalls`` — pops ``tool_calls_pending`` and calls - the matching ``Tool`` instance -""" - -from __future__ import annotations - -import logging -import re -from dataclasses import dataclass, field -from typing import Any - -from .runtime_state import push_log, set_if_changed, take_event -from .triggers import EventTrigger, HzTrigger, Trigger - -logger = logging.getLogger(__name__) - - -# --------------------------------------------------------------------------- -# Step base + runner -# --------------------------------------------------------------------------- - - -@dataclass -class InferenceStep: - """A trigger-gated callable. Subclasses override :meth:`run`.""" - - trigger: Trigger - - def __call__(self, state: dict[str, Any]) -> dict[str, Any]: - if not self.trigger.should_fire(state["_tick"], state): - return state - return self.run(state) or state - - def run(self, state: dict[str, Any]) -> dict[str, Any] | None: # pragma: no cover - raise NotImplementedError - - -# --------------------------------------------------------------------------- -# Low-level (action) path -# --------------------------------------------------------------------------- - - -@dataclass -class LowLevelForward(InferenceStep): - """Run the policy's action head and produce one action chunk.""" - - policy: Any = None - observation_provider: Any = None - """Callable ``() -> dict``: returns the current observation batch - (already preprocessed). Typically wraps the robot's camera / - proprio reads. ``None`` in dry-run mode → step skips.""" - - trigger: Trigger = field(default_factory=lambda: HzTrigger(hz=4.0)) - - def run(self, state: dict[str, Any]) -> dict[str, Any] | None: - if self.policy is None or self.observation_provider is None: - return None - # ``/vlm`` mode pauses the whole action loop so the robot holds - # position while the operator probes the VLM with VQA. - if state.get("mode", "action") != "action": - return None - if not state.get("task"): - return None - - # PI052 produces *action chunks* (typically 50 steps via - # flow-matching). Every step gets dispatched to the robot; - # popping one per dispatch tick is essentially free. Only - # generate a new chunk once the previous one has fully - # drained — this is the canonical "sense → think → act" - # loop. Refreshing while a chunk is still queued causes the - # new chunk to "telescope" past the old one (planned from an - # observation that's already 25+ steps stale by the time it - # starts dispatching). - queue = state.setdefault("action_queue", []) - if len(queue) > 0: - return None - - observation = self.observation_provider() - if observation is None: - return None - - # The action expert is conditioned on the SUBTASK generated by - # the high-level loop (``HighLevelSubtaskFwd`` runs earlier in - # the pipeline and writes ``current_subtask``). Matches the - # training-time ``low_level_execution`` recipe — ``user(${subtask})``. - # Falls back to the task string only on the very first frame, - # before the high-level loop has produced a subtask. - subtask = state.get("current_subtask") or state.get("task") or "" - ctx = [{"role": "user", "content": subtask}] - # ``add_generation_prompt=False`` to match the training-time - # prefix shape: at training the action expert sees the rendered - # user turn ending at ``<|im_end|>`` (no trailing - # ``<|im_start|>assistant\n``). Passing True here would append - # extra role-marker tokens the action expert never saw during - # training. - text_batch = _build_text_batch(self.policy, ctx, add_generation_prompt=False) - from lerobot.utils.constants import ( # noqa: PLC0415 - OBS_LANGUAGE_ATTENTION_MASK, - OBS_LANGUAGE_TOKENS, - ) - - observation = dict(observation) - observation[OBS_LANGUAGE_TOKENS] = text_batch["lang_tokens"] - observation[OBS_LANGUAGE_ATTENTION_MASK] = text_batch["lang_masks"] - - try: - # ``predict_action_chunk`` returns the *full* chunk shape - # ``(batch, n_action_steps, action_dim)``. Enqueue every - # step so DispatchAction at ctrl_hz can drain them - # smoothly until the next refresh. - chunk = self.policy.predict_action_chunk(observation) - except Exception as exc: # noqa: BLE001 - logger.warning( - "predict_action_chunk failed: %s", - exc, - exc_info=logger.isEnabledFor(logging.DEBUG), - ) - push_log( - state, - f" [warn] predict_action_chunk failed: {type(exc).__name__}: {exc}", - ) - return None - - # ``chunk`` shape: ``(batch, n_action_steps, action_dim)``. Push - # each step as a ``(1, action_dim)`` tensor so the existing - # action executor's batch-squeeze logic works unchanged. - if chunk.ndim == 3: - chunk_iter = chunk[0] # ``(n_action_steps, action_dim)`` - elif chunk.ndim == 2: - chunk_iter = chunk - else: - chunk_iter = chunk.unsqueeze(0) - - for step in chunk_iter: - queue.append(step.unsqueeze(0)) - state["last_chunk_size"] = int(chunk_iter.shape[0]) - return None - - -@dataclass -class DispatchAction(InferenceStep): - """Pop one action per tick and hand it to the robot. - - In dry-run mode (``robot_executor=None``) the step still pops the - queue so it doesn't grow unbounded — the popped tensor is logged - instead of executed. - - Wall-clock catch-up: the action queue represents an open-loop - trajectory at a fixed step rate (``trigger.hz`` ≈ ``ctrl_hz``). - When the main loop stalls — e.g. an LLM call for the high-level - subtask blocks for ~2 s on MPS — the dispatch trigger fires only - once over that whole interval. Naively popping a single entry per - fire makes the robot lag further and further behind the planned - timeline, and a 50-step chunk would take ~125 s to drain instead - of ~1.7 s. Track real elapsed time between dispatches and pop - ``round(elapsed * hz)`` entries, sending the most recent one. The - skipped intermediate joint targets are stale anyway — the dynamixel - will smooth toward the latest goal position. - """ - - robot_executor: Any = None - trigger: Trigger = field(default_factory=lambda: HzTrigger(hz=50.0)) - _last_dispatch_t: float | None = field(default=None, init=False) - - def run(self, state: dict[str, Any]) -> dict[str, Any] | None: - import time as _time # noqa: PLC0415 - - # ``/vlm`` mode pauses dispatch — the robot holds its last - # commanded position while the operator runs VQA. - if state.get("mode", "action") != "action": - self._last_dispatch_t = None - return None - - queue = state.get("action_queue") - if not queue: - # Reset wall-clock anchor when the queue is empty so the - # next chunk doesn't see a huge fake "elapsed" window. - self._last_dispatch_t = None - return None - - now = _time.monotonic() - hz = getattr(self.trigger, "hz", 30.0) - if self._last_dispatch_t is None or hz <= 0: - n_to_pop = 1 - else: - elapsed = now - self._last_dispatch_t - # ``max(1, ...)`` so we always pop at least one when the - # trigger fires; ``min(len(queue), ...)`` so we don't run - # off the end of the chunk. - n_to_pop = max(1, min(len(queue), int(round(elapsed * hz)))) - self._last_dispatch_t = now - - # Drain ``n_to_pop`` stale entries, keep only the latest as the - # action actually sent. The intermediate joint targets would - # all be ~10–30 ms apart in chunk time — the robot can't track - # them individually anyway when the host loop is slow. - latest = None - for _ in range(n_to_pop): - if not queue: - break - latest = queue.popleft() if hasattr(queue, "popleft") else queue.pop(0) - state["actions_dispatched"] = state.get("actions_dispatched", 0) + 1 - - if latest is not None and self.robot_executor is not None: - self.robot_executor(latest) - return None - - -# --------------------------------------------------------------------------- -# High-level (text) paths — all use policy.select_message -# --------------------------------------------------------------------------- - - -_LOC_TOKENIZER_CACHE: dict[str, Any] = {} - - -def _get_loc_tokenizer(tok_name: str, auto_tokenizer_cls: Any, register_loc_fn: Any) -> Any: - """Return a loc-token-registered tokenizer, loading from disk only once. - - ``AutoTokenizer.from_pretrained`` + loc-token registration is expensive and - the result is immutable, so cache per ``tok_name``. - """ - tokenizer = _LOC_TOKENIZER_CACHE.get(tok_name) - if tokenizer is None: - tokenizer = register_loc_fn(auto_tokenizer_cls.from_pretrained(tok_name)) - _LOC_TOKENIZER_CACHE[tok_name] = tokenizer - return tokenizer - - -def _build_text_batch( - policy: Any, - prompt_messages: list[dict[str, Any]], - *, - add_generation_prompt: bool = True, -) -> dict[str, Any]: - """Tokenize chat messages into the batch ``select_message`` expects. - - PI052's backbone (PaliGemma) ships no chat template, so we train on - a plain role-prefixed concatenation built by - ``PI052TextTokenizerStep``. We reuse that exact formatter so the - inference prefix matches training; ``add_generation_prompt`` appends - the bare ``Assistant: `` header the LM head continues from. - """ - import torch # noqa: PLC0415 - from transformers import AutoTokenizer # noqa: PLC0415 - - from lerobot.policies.pi052.text_processor_pi052 import ( # noqa: PLC0415 - _flatten_say_tool_calls, - _format_messages, - _strip_blocks, - register_paligemma_loc_tokens, - ) - - tok_name = getattr(policy.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224" - # Register PaliGemma's tokens so inference encoding / - # decoding sees them as single vocab ids — must match training. - # The tokenizer is read-only after registration, so cache it: rebuilding it - # from disk on every call dominated eval runtime (this runs twice per env - # per replan — subtask gen + action prompt). - tokenizer = _get_loc_tokenizer(tok_name, AutoTokenizer, register_paligemma_loc_tokens) - - messages = [_strip_blocks(_flatten_say_tool_calls(m)) for m in prompt_messages] - prompt, _spans = _format_messages(messages) - if add_generation_prompt: - prompt = prompt + "Assistant: " - - encoded = tokenizer(prompt, return_tensors="pt") - ids = encoded["input_ids"] - attn = encoded.get("attention_mask") - if attn is None and tokenizer.pad_token_id is not None: - attn = ids != tokenizer.pad_token_id - if attn is not None and hasattr(attn, "dtype") and attn.dtype != torch.bool: - attn = attn.bool() - - # Move tokens onto the policy's device — otherwise prefix embedding - # raises a device-mismatch on every forward (CPU tensor vs MPS / CUDA - # model), which the caller's broad except would swallow silently. - device = getattr(getattr(policy, "config", None), "device", None) - if device is not None: - try: - ids = ids.to(device) - if attn is not None and hasattr(attn, "to"): - attn = attn.to(device) - except Exception as exc: # noqa: BLE001 - logger.debug("could not move pi052 lang tokens to %s: %s", device, exc) - return {"lang_tokens": ids, "lang_masks": attn, "tokenizer": tokenizer} - - -def _strip_recipe_keys(m: dict[str, Any]) -> dict[str, Any]: - new = dict(m) - new.pop("stream", None) - new.pop("target", None) - return new - - -@dataclass -class HighLevelSubtaskFwd(InferenceStep): - """At ~1 Hz, ask the policy for the next subtask. - - Mirrors the ``high_level_subtask`` recipe layout exactly: - - user: "${task}\\nPlan: ${plan}\\nMemory: ${memory}" - user: "Current subtask: ${subtask}" (if subtask present) - ↓ generate ↓ - assistant: - """ - - policy: Any = None - observation_provider: Any = None - """Same shape as ``LowLevelForward.observation_provider``. When - set, the resulting observation is merged into ``select_message``'s - batch so text generation runs against real video + state.""" - - trigger: Trigger = field(default_factory=lambda: HzTrigger(hz=1.0)) - - def run(self, state: dict[str, Any]) -> dict[str, Any] | None: - if self.policy is None or not state.get("task"): - return None - # ``/vlm`` mode pauses subtask generation along with the rest of - # the action loop. - if state.get("mode", "action") != "action": - return None - # Gate to chunk boundaries: only generate a fresh subtask when - # the action queue is empty (i.e. right before LowLevelForward - # refreshes the chunk). ``select_message`` takes ~2 s on MPS, - # and running it every loop iteration starves DispatchAction - # at ctrl_hz=30 — the queue drains at ~0.4 actions/sec instead - # of 30/sec and the robot barely moves. Tying it to the same - # "queue empty" condition as the chunk refresh produces a - # clean sense → think → act cycle. - # - # Rearm the trigger when skipping so a low-hz schedule - # (e.g. ``--high_level_hz=0.2`` = once per 5 s) doesn't lose - # the slot: the trigger fires once on the timer but the brief - # queue-empty window almost never coincides, so without rearm - # HL would effectively never run. - queue = state.get("action_queue") or [] - if len(queue) > 0: - if hasattr(self.trigger, "rearm"): - self.trigger.rearm() - return None - # Per-chunk-boundary throttle: at each "queue empty" moment we - # increment a counter; subtask gen only fires once the counter - # reaches ``subtask_chunks_per_gen``. Lets the operator run e.g. - # 5 action chunks per subtask-gen so the LM head doesn't churn - # every 1.7 s (a fresh subtask while the previous one is still - # being executed is wasted compute *and* causes the action - # expert's flow trajectory to be re-planned mid-grasp). - chunks_per_gen = max(1, int(state.get("subtask_chunks_per_gen", 1) or 1)) - # Initialise so the first chunk boundary fires immediately - # (counter starts at chunks_per_gen, decrements per skip, - # generates and resets when it hits 0). - if "_hl_chunks_until_gen" not in state: - state["_hl_chunks_until_gen"] = 0 - if state["_hl_chunks_until_gen"] > 0: - state["_hl_chunks_until_gen"] -= 1 - if hasattr(self.trigger, "rearm"): - self.trigger.rearm() - return None - state["_hl_chunks_until_gen"] = chunks_per_gen - 1 - ctx = _msgs_for_subtask(state) - observation = _maybe_observation(self.observation_provider) - # Default: greedy argmax, no min_new_tokens, no special-token - # suppression — matches training. Operator can override via - # ``--text_min_new_tokens=N --text_temperature=T --text_top_p=P`` - # on the CLI; useful for under-trained checkpoints whose LM - # head still favours EOS at position 0 (pre-trained chat - # backbone's short-turn prior hasn't been fully overridden - # by the fine-tuning supervision yet). - msg = _generate_with_policy( - self.policy, - ctx, - observation=observation, - state=state, - label="subtask gen", - min_new_tokens=int(state.get("text_gen_min_new_tokens") or 0), - temperature=float(state.get("text_gen_temperature") or 0.0), - top_p=float(state.get("text_gen_top_p") or 1.0), - # Subtasks never legitimately contain PaliGemma ```` - # tokens — suppress them so a checkpoint whose LM head - # has drifted toward the pretrained loc-prior falls back - # to its (still-correct) text mass. - suppress_loc_tokens=True, - ) - # Diagnostics: surface what the model is *actually* producing - # at chunk boundaries, even when the output gets rejected or - # repeats. Memorisation collapse looks like "same accepted - # subtask N times in a row" or "gibberish_count rising while - # current_subtask is stuck". The state panel renders these. - state["last_subtask_raw"] = msg or "" - # Persistent empty completion is its own failure mode (model - # immediately EOS-es from the chat-template generation - # prompt) — surface it once every N occurrences so the - # operator can distinguish "generation failing silently" - # from "generating fine but filter rejecting". - if not msg: - empties = state.get("subtask_empty_count", 0) + 1 - state["subtask_empty_count"] = empties - if empties == 1 or empties % 5 == 0: - debug = getattr(self.policy, "_last_select_message_debug", "") or "" - if debug: - push_log( - state, - f" [info] subtask gen empty (×{empties}); {debug}", - ) - else: - push_log( - state, - f" [info] subtask gen returned empty (×{empties}) — " - "no tokens generated (head EOS-ing before any " - "non-special token).", - ) - if msg and _looks_like_gibberish(msg): - # Bump a counter so the operator can see the model is - # struggling without spamming the log every tick. A first - # rejection still logs once so the failure is visible. - count = state.get("subtask_gibberish_count", 0) + 1 - state["subtask_gibberish_count"] = count - if count == 1 or count % 30 == 0: - push_log( - state, - f" [info] subtask gen rejected (gibberish ×{count}): {msg[:60]!r}", - ) - return None - if msg: - prev_subtask = state.get("current_subtask") - changed = set_if_changed(state, "current_subtask", msg, label="subtask") - if changed: - # Stash the just-completed subtask so ``MemoryUpdateFwd`` - # can drop it into its prompt as ``Completed subtask:`` - # — the recipe binds ``completed_subtask`` to - # ``nth_prev(style=subtask, offset=1)``, i.e. the subtask - # that was active *before* the change. - if prev_subtask: - state["prior_subtask"] = prev_subtask - # Subtask change is a downstream trigger. - state.setdefault("events_this_tick", []).append("subtask_change") - state["subtask_repeat_count"] = 0 - else: - # Same accepted string regenerated — memorisation tell. - # Once this counter climbs past a few, you're seeing - # the model unable to move past the current subtask - # despite the chunk having drained (visual scene may - # have changed but the LM is replaying training - # tokens). - state["subtask_repeat_count"] = state.get("subtask_repeat_count", 0) + 1 - # Silently skip empty completions — common when the model - # warms up or generates only EOS; logging it every tick at - # ctrl_hz is just noise. - return None - - -@dataclass -class MemoryUpdateFwd(InferenceStep): - """On subtask boundary, refresh the compressed memory. - - Mirrors the ``memory_update`` recipe layout exactly: - - user: "${task}" - assistant: "Previous memory: ${prior_memory}" (if prior memory) - user: "Completed subtask: ${completed_subtask}" (if subtask) - ↓ generate ↓ - assistant: - """ - - policy: Any = None - observation_provider: Any = None - trigger: Trigger = field(default_factory=lambda: EventTrigger("subtask_change")) - - def run(self, state: dict[str, Any]) -> dict[str, Any] | None: - # Don't consume the event — multiple steps may want to react. - if self.policy is None: - return None - ctx = _msgs_for_memory(state) - observation = _maybe_observation(self.observation_provider) - new_memory = _generate_with_policy( - self.policy, - ctx, - observation=observation, - state=state, - label="memory gen", - suppress_loc_tokens=True, - ) - state["last_memory_raw"] = new_memory or "" - if new_memory and _looks_like_gibberish(new_memory): - count = state.get("memory_gibberish_count", 0) + 1 - state["memory_gibberish_count"] = count - push_log( - state, - f" [info] memory gen rejected (gibberish ×{count}): {new_memory[:60]!r}", - ) - return None - if new_memory: - set_if_changed(state, "current_memory", new_memory, label="memory") - return None - - -@dataclass -class UserInterjectionFwd(InferenceStep): - """On stdin interjection, refresh the plan + emit a paired ``say``. - - Mirrors the ``user_interjection_response`` recipe layout exactly: - - user: "${task}" - assistant: "Previous plan:\\n${prior_plan}" (if prior plan) - user: "${interjection}" (the new utterance) - ↓ generate ↓ - assistant: ...> - """ - - policy: Any = None - observation_provider: Any = None - trigger: Trigger = field(default_factory=lambda: EventTrigger("user_interjection")) - - def run(self, state: dict[str, Any]) -> dict[str, Any] | None: - if self.policy is None or not take_event(state, "user_interjection"): - return None - ctx = _msgs_for_interjection(state) - observation = _maybe_observation(self.observation_provider) - out = _generate_with_policy( - self.policy, - ctx, - observation=observation, - state=state, - label="plan/say gen", - suppress_loc_tokens=True, - ) - if not out: - # Don't log every empty completion — happens repeatedly on - # MPS during warm-up and floods the panel. The user can - # re-trigger by typing again. - return None - if _looks_like_gibberish(out): - count = state.get("plan_gibberish_count", 0) + 1 - state["plan_gibberish_count"] = count - push_log( - state, - f" [info] plan/say gen rejected (gibberish ×{count}): {out[:60]!r}", - ) - return None - # Heuristic split: model is trained to emit one assistant turn - # carrying both plan text AND a `say` tool call. Look for a - # "..." or "say(...)" marker; fall back to whole - # text → plan, no speech. - plan_text, speech_text = _split_plan_and_say(out) - if plan_text and _looks_like_gibberish(plan_text): - plan_text = "" - if plan_text: - set_if_changed(state, "current_plan", plan_text, label="plan") - if speech_text: - push_log(state, f" speech: {speech_text}") - state.setdefault("tool_calls_pending", []).append( - { - "type": "function", - "function": {"name": "say", "arguments": {"text": speech_text}}, - } - ) - state.setdefault("events_this_tick", []).append("tool_call_pending") - # Mark interjection consumed. - state["recent_interjection"] = None - return None - - -@dataclass -class AskVQAFwd(InferenceStep): - """On stdin question, answer a frame-grounded VQA. - - Mirrors the ``ask_vqa_*`` recipe layout exactly: a single user - turn carrying just the VQA question, plus the camera image block - in training (we drop the image at inference because the dataset's - image preprocessing doesn't match SmolVLM's vision tower input). - - user: - ↓ generate ↓ - assistant: - """ - - policy: Any = None - observation_provider: Any = None - trigger: Trigger = field(default_factory=lambda: EventTrigger("user_vqa_query")) - - def run(self, state: dict[str, Any]) -> dict[str, Any] | None: - if self.policy is None or not take_event(state, "user_vqa_query"): - return None - question = state.get("recent_vqa_query") - if not question: - return None - ctx = _msgs_for_vqa(question) - observation = _maybe_observation(self.observation_provider) - answer = _generate_with_policy( - self.policy, - ctx, - observation=observation, - state=state, - label="vqa gen", - ) - # VQA answers are intentionally JSON-like during training, so - # ``_looks_like_gibberish`` would false-positive on them. Keep - # the answer as-is — the VQA panel line lets the user judge. - if answer: - push_log(state, f" vqa: {answer}") - state["recent_vqa_query"] = None - return None - - -# --------------------------------------------------------------------------- -# Tool dispatch -# --------------------------------------------------------------------------- - - -@dataclass -class DispatchToolCalls(InferenceStep): - """Pop ``tool_calls_pending`` and execute them via :data:`TOOL_REGISTRY`.""" - - tools: dict[str, Any] = field(default_factory=dict) - trigger: Trigger = field(default_factory=lambda: EventTrigger("tool_call_pending")) - - def run(self, state: dict[str, Any]) -> dict[str, Any] | None: - take_event(state, "tool_call_pending") - pending = state.get("tool_calls_pending") or [] - for call in pending: - try: - fn = (call or {}).get("function") or {} - name = fn.get("name") - args = fn.get("arguments") or {} - tool = self.tools.get(name) - if tool is None: - push_log(state, f" [warn] tool {name!r} not registered — skipping call") - continue - tool.call(args) - except Exception as exc: # noqa: BLE001 - push_log(state, f" [error] tool dispatch failed: {exc}") - state["tool_calls_pending"] = [] - return None - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _looks_like_gibberish(text: str) -> bool: - """Heuristically detect generation that's clearly off the rails. - - Memorised models can collapse to dominant-mode outputs when the - prompt drifts even slightly from training distribution. Reject: - - * empty / whitespace-only - * too few alphabetic characters (mostly punctuation) - * a single character repeated past the threshold - * starts with ``":"`` and contains no letters - * too few unique tokens — e.g. ``"the"``, ``"the the the"``, - ``"Ass\\n::\\nthe"`` (the collapse seen on real-robot frames - where the model emits one or two memorised tokens repeatedly) - * chat-template fragment leakage (``Assistant:``, ``User:``, - ``Ass\\n``) - - Real subtasks look like ``"close the gripper to grasp the blue - cube"`` — multiple unique alphabetic tokens, no role-marker - fragments. Anything materially shorter than that is rejected. - """ - if not text or not text.strip(): - return True - stripped = text.strip() - alpha = sum(1 for c in stripped if c.isalpha()) - if alpha < max(3, len(stripped) // 8): - return True - if stripped.startswith('":') and stripped.count('"') > stripped.count(" "): - return True - # Single repeating char: e.g. ``""""""``. - if len(set(stripped)) <= 2 and len(stripped) > 4: - return True - # Chat-template fragment leakage — the model emits ``Ass``, - # ``Assistant:``, ``User:``, often with extra newlines/colons. - # Reject if the cleaned text is mostly role-marker shards. - cleaned = stripped.replace("\n", " ").replace(":", " ") - for marker in ("Assistant", "User", "Ass "): - if marker in cleaned and len(cleaned.split()) < 4: - return True - tokens = [t for t in cleaned.split() if any(c.isalpha() for c in t)] - unique_alpha = {t.lower() for t in tokens} - # Short degenerate output — model stuck on ``the`` or a couple of - # memorised single-token continuations. - if len(unique_alpha) < 3 and len(stripped) < 80: - return True - # Long repetition collapse — the LM head loops an n-gram for the - # whole generation budget ("the arm the arm … the the the the"). - # Length-independent: many tokens but a tiny unique ratio. The - # earlier ``< 80`` check missed these because the looped string - # blows well past 80 chars. - return len(tokens) >= 8 and len(unique_alpha) <= max(3, len(tokens) // 10) - - -def _control_context_messages( - state: dict[str, Any], - *, - include_completed: bool = False, - extra_user: str | None = None, -) -> list[dict[str, Any]]: - """Build a chat-template-ready prompt from current runtime state. - - Mirrors what the recipe renders into ``${task}\nPlan: - ${plan}\nMemory: ${memory}`` for the high-level branches. - """ - # Always emit ``Plan: `` / ``Memory: `` labels — even with empty - # values — to mirror the training-time recipe substitution. - task = state.get("task") or "" - plan = state.get("current_plan") or "" - memory = state.get("current_memory") or "" - parts = [task, f"Plan: {plan}", f"Memory: {memory}"] - if include_completed and state.get("current_subtask"): - parts.append(f"Completed subtask: {state['current_subtask']}") - head = "\n".join(parts) - msgs: list[dict[str, Any]] = [{"role": "user", "content": head}] - if extra_user: - msgs.append({"role": "user", "content": extra_user}) - return msgs - - -# --------------------------------------------------------------------------- -# Per-recipe prompt builders. Each one mirrors a single sub-recipe's -# message layout in the recipe so the chat-templated -# prompt at inference matches what the model saw during training. -# Generic ``_control_context_messages`` is kept around as a fallback -# for ad-hoc callers but the four high-level steps now use these. -# --------------------------------------------------------------------------- - - -def _hirobot_user_head(state: dict[str, Any]) -> str: - """Build the ``task\\nPlan: …\\nMemory: …`` user content string. - - Mirrors what the recipe renders at training time, where - ``language_render._substitute`` substitutes empty strings for - missing ``${plan}`` / ``${memory}`` bindings — i.e. the - ``Plan: `` / ``Memory: `` prefix labels are *always* in the - user turn, even when their values aren't set yet. Skipping them - here (the previous behaviour) produced a different prompt shape - on early frames before plan / memory are populated and on - samples where the dataset has no plan / memory annotation. - """ - task = state.get("task") or "" - plan = state.get("current_plan") or "" - memory = state.get("current_memory") or "" - return f"{task}\nPlan: {plan}\nMemory: {memory}" - - -def _msgs_for_subtask(state: dict[str, Any]) -> list[dict[str, Any]]: - """``high_level_subtask`` recipe layout — predict the subtask from the - task. The v-current recipe's user turn is just ``${task}`` (plan and - memory are not trained), so the inference prompt is the bare task — - no ``Plan: `` / ``Memory: `` lines. - """ - return [{"role": "user", "content": state.get("task") or ""}] - - -def _msgs_for_memory(state: dict[str, Any]) -> list[dict[str, Any]]: - """Memory-update prompt — mirrors ``memory_update`` recipe layout. - - Recipe layout (``subtask_mem.yaml``): - - user: "${task}" - assistant: "Previous memory: ${prior_memory}" (if_present prior) - user: "Completed subtask: ${completed}" (if_present completed) - assistant: → predicts new memory - - Fired by ``MemoryUpdateFwd`` on a ``subtask_change`` event: - ``state['current_memory']`` is the memory the policy last emitted - (= the ``prior_memory`` binding at training), and - ``state['prior_subtask']`` is the subtask that just got replaced - (= the ``completed_subtask`` binding at training). - """ - msgs: list[dict[str, Any]] = [ - {"role": "user", "content": state.get("task") or ""}, - ] - prior_memory = state.get("current_memory") - if prior_memory: - msgs.append({"role": "assistant", "content": f"Previous memory: {prior_memory}"}) - completed_subtask = state.get("prior_subtask") - if completed_subtask: - msgs.append({"role": "user", "content": f"Completed subtask: {completed_subtask}"}) - return msgs - - -def _msgs_for_interjection(state: dict[str, Any]) -> list[dict[str, Any]]: - """``user_interjection_response`` recipe layout.""" - msgs: list[dict[str, Any]] = [{"role": "user", "content": state.get("task") or ""}] - if state.get("current_plan"): - msgs.append({"role": "assistant", "content": f"Previous plan:\n{state['current_plan']}"}) - interjection = state.get("recent_interjection") - if interjection: - msgs.append({"role": "user", "content": interjection}) - return msgs - - -def _msgs_for_plan(state: dict[str, Any]) -> list[dict[str, Any]]: - """``plan_generation`` recipe layout — bare task → plan. - - The assistant turn is the generation target, so we only render - the user turn at inference; the runtime appends the predicted - plan after sampling. - """ - return [{"role": "user", "content": state.get("task") or ""}] - - -def _msgs_for_vqa(question: str) -> list[dict[str, Any]]: - """``ask_vqa_*`` recipe layout (text-only at inference).""" - return [{"role": "user", "content": question}] - - -def _maybe_observation(provider: Any) -> dict | None: - """Pull one observation from ``provider`` if it's set, else ``None``. - - Errors from the provider are logged at debug level and swallowed — - text generation still runs (in text-only mode) so a flaky frame - source doesn't kill the REPL. - """ - if provider is None: - return None - try: - return provider() - except Exception as exc: # noqa: BLE001 - logger.debug("observation_provider raised %s — falling back to text-only", exc) - return None - - -def _generate_with_policy( - policy: Any, - messages: list[dict[str, Any]], - *, - observation: dict | None = None, - state: dict[str, Any] | None = None, - label: str = "select_message", - min_new_tokens: int = 0, - temperature: float = 0.0, - top_p: float = 1.0, - suppress_loc_tokens: bool = False, -) -> str: - """Drive ``policy.select_message`` with a chat batch (and optional obs). - - When ``observation`` carries ``observation.images.*`` and - ``observation.state``, those are merged into the batch so - ``select_message`` runs the same VLM prefix the policy was trained - on. Without an observation the runtime falls back to a text-only - prompt — the text head still runs, but generations may drift from - the training distribution. - - Failures are surfaced both to the module logger (``warning``) and, - when ``state`` is given, to the runtime's user-visible log via - :func:`push_log`, so the REPL no longer "looks dead" when - something goes wrong inside generation. - """ - if not hasattr(policy, "select_message"): - if state is not None: - push_log(state, f" [warn] policy has no select_message — skipping {label}") - return "" - text_batch = _build_text_batch(policy, messages) - try: - from lerobot.utils.constants import ( # noqa: PLC0415 - OBS_LANGUAGE_ATTENTION_MASK, - OBS_LANGUAGE_TOKENS, - ) - - batch: dict[str, Any] = { - OBS_LANGUAGE_TOKENS: text_batch["lang_tokens"], - OBS_LANGUAGE_ATTENTION_MASK: text_batch["lang_masks"], - } - if observation: - for k, v in observation.items(): - if isinstance(k, str) and k.startswith("observation.") and k not in batch: - batch[k] = v - kwargs: dict[str, Any] = { - "tokenizer": text_batch["tokenizer"], - "min_new_tokens": min_new_tokens, - "temperature": temperature, - "top_p": top_p, - } - kwargs["suppress_loc_tokens"] = suppress_loc_tokens - return policy.select_message(batch, **kwargs) - except Exception as exc: # noqa: BLE001 - logger.warning("%s failed: %s", label, exc, exc_info=logger.isEnabledFor(logging.DEBUG)) - if state is not None: - push_log(state, f" [warn] {label} failed: {type(exc).__name__}: {exc}") - return "" - - -_SAY_RE = re.compile(r"<\s*say\s*>(.*?)<\s*/\s*say\s*>", re.IGNORECASE | re.DOTALL) - - -def _split_plan_and_say(text: str) -> tuple[str, str]: - """Pull a ``...`` snippet out of ``text``; remainder is plan. - - The training-time tool-call serializer wraps ``say(text="…")`` in a - deterministic textual marker so prefix-LM-style training learns to - emit it. The runtime parses it back here. If no marker is present, - the entire text is treated as plan with no speech. - """ - if not text: - return "", "" - match = _SAY_RE.search(text) - if not match: - return text.strip(), "" - speech = match.group(1).strip().strip('"').strip("'") - plan = (text[: match.start()] + text[match.end() :]).strip() - return plan, speech +"""Compatibility exports for PI052 model helper imports.""" + +from .pi052_adapter import ( + _build_text_batch, + _generate_with_policy, + _get_loc_tokenizer, + looks_like_gibberish as _looks_like_gibberish, +) + +__all__ = [ + "_build_text_batch", + "_generate_with_policy", + "_get_loc_tokenizer", + "_looks_like_gibberish", +] diff --git a/src/lerobot/policies/pi052/inference/triggers.py b/src/lerobot/policies/pi052/inference/triggers.py deleted file mode 100644 index dd9537f66..000000000 --- a/src/lerobot/policies/pi052/inference/triggers.py +++ /dev/null @@ -1,134 +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. -"""Trigger primitives for PI052's multi-rate inference runtime. - -Mirrors the plan's Section "Runtime orchestration": each -``InferenceStep`` is gated by a :class:`Trigger` that decides per tick -whether the step fires. Two trigger flavours cover all the cadences -the canonical recipe needs: - -* :class:`HzTrigger` for periodic beats (action chunks at ~3-5 Hz, - high-level subtask generation at ~1 Hz, action dispatch at ~50 Hz) -* :class:`EventTrigger` for one-shot reactions (subtask boundary → - memory update; user interjection → plan refresh; user VQA query → - vqa answer; pending tool call → dispatcher) - -Triggers are stateless except for ``HzTrigger``'s last-fire timestamp. -The runtime stores the :class:`Tick` clock as ``state["_tick"]`` so -every step shares a single time source. -""" - -from __future__ import annotations - -import time -from dataclasses import dataclass, field -from typing import Any, Protocol - - -@dataclass -class Tick: - """Single tick from :class:`TickClock`. Carries time references the - runtime steps consume to gate themselves.""" - - index: int - """Monotonic counter — increments by one per tick.""" - - monotonic_seconds: float - """``time.monotonic()`` at the start of this tick.""" - - -@dataclass -class TickClock: - """Drives the runtime loop at up to ``max_rate_hz``. - - Sleeps just enough between :meth:`advance` calls to enforce the - rate. With ``max_rate_hz=50`` the loop wakes ~every 20ms; the - higher-level ``HzTrigger`` slices that timeline into sub-cadences. - """ - - max_rate_hz: float = 50.0 - _index: int = field(default=0, init=False) - _last_seconds: float | None = field(default=None, init=False) - - def advance(self) -> Tick: - period = 1.0 / max(self.max_rate_hz, 0.1) - now = time.monotonic() - if self._last_seconds is not None: - sleep_for = (self._last_seconds + period) - now - if sleep_for > 0: - time.sleep(sleep_for) - now = time.monotonic() - self._last_seconds = now - self._index += 1 - return Tick(index=self._index, monotonic_seconds=now) - - -class Trigger(Protocol): - """Decide whether the next ``InferenceStep`` should fire.""" - - def should_fire(self, tick: Tick, state: dict[str, Any]) -> bool: ... - - -@dataclass -class HzTrigger: - """Fire at most ``hz`` times per second. - - A step that gates further (e.g. ``HighLevelSubtaskFwd`` skipping - when the action queue is non-empty) and wants the trigger to - retry next tick instead of waiting a full period can call - :meth:`rearm` from inside ``run``. Without this, a low-hz trigger - (e.g. ``hz=0.2`` = once per 5 s) almost never coincides with the - brief queue-empty window and the step never fires at all. - """ - - hz: float - _last_seconds: float | None = field(default=None, init=False) - - def should_fire(self, tick: Tick, state: dict[str, Any]) -> bool: - period = 1.0 / max(self.hz, 1e-6) - if self._last_seconds is None or (tick.monotonic_seconds - self._last_seconds) >= period: - self._last_seconds = tick.monotonic_seconds - return True - return False - - def rearm(self) -> None: - """Mark the trigger as not having fired, so the next tick re-evaluates. - - Used by a step that decided to skip after ``should_fire`` already - committed the firing — keeps the cadence honest without losing - the slot. - """ - self._last_seconds = None - - -@dataclass -class EventTrigger: - """Fire when ``event_name`` is in ``state["events_this_tick"]``. - - The runtime fills ``events_this_tick`` once per tick from: - - * stdin / network input (``user_interjection``, ``user_vqa_query``, - ``stop``) - * internal state transitions (``subtask_change``, - ``tool_call_pending``) - - The list is consumed (cleared at the end of the tick) so events - fire at most once. - """ - - event_name: str - - def should_fire(self, tick: Tick, state: dict[str, Any]) -> bool: - events: list[str] = state.get("events_this_tick") or [] - return self.event_name in events diff --git a/src/lerobot/policies/pi052/inference/ui.py b/src/lerobot/policies/pi052/inference/ui.py index a2d90f076..28b15a8d7 100644 --- a/src/lerobot/policies/pi052/inference/ui.py +++ b/src/lerobot/policies/pi052/inference/ui.py @@ -76,10 +76,7 @@ def make_state_panel(state: dict[str, Any]) -> Any: table.add_column(justify="left") for key, label in _STATE_KEYS: value = state.get(key) - if value is None: - rendered = Text("(not set)", style="dim italic") - else: - rendered = Text(str(value), style="bold") + rendered = Text("(not set)", style="dim italic") if value is None else Text(str(value), style="bold") table.add_row(label, rendered) queue = state.get("action_queue") queue_len = len(queue) if hasattr(queue, "__len__") else 0 @@ -92,9 +89,7 @@ def make_state_panel(state: dict[str, Any]) -> Any: ) table.add_row("", footer) run_mode = state.get("mode", "action") - mode_tag = ( - "[green]action[/]" if run_mode == "action" else "[yellow]paused[/]" - ) + mode_tag = "[green]action[/]" if run_mode == "action" else "[yellow]paused[/]" return Panel( table, title=f"[bold]PI052 state[/] · mode: {mode_tag}", diff --git a/src/lerobot/policies/pi052/inference/vqa.py b/src/lerobot/policies/pi052/inference/vqa.py index 6561a25e2..f2e6f8a7b 100644 --- a/src/lerobot/policies/pi052/inference/vqa.py +++ b/src/lerobot/policies/pi052/inference/vqa.py @@ -42,11 +42,10 @@ import subprocess import sys import time import webbrowser +from contextlib import suppress from pathlib import Path from typing import Any -from .runtime_state import push_log - logger = logging.getLogger(__name__) _IMAGE_PREFIX = "observation.images." @@ -162,8 +161,7 @@ def parse_loc_answer(answer: str) -> dict | None: 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 + {"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: @@ -174,9 +172,7 @@ def parse_loc_answer(answer: str) -> dict | None: "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 - ] + 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 @@ -311,11 +307,11 @@ 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) + subprocess.run(["open", str(path)], check=False) # nosec B607 elif sys.platform.startswith("linux"): - subprocess.run(["xdg-open", str(path)], check=False) + subprocess.run(["xdg-open", str(path)], check=False) # nosec B607 elif os.name == "nt": - os.startfile(str(path)) # type: ignore[attr-defined] # noqa: S606 + 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 @@ -339,10 +335,11 @@ def save_and_open_overlay(image: Any, out_dir: str | Path = "./vqa_overlays") -> def handle_vqa_query( *, - policy: Any, + policy_adapter: Any | None = None, + policy: Any | None = None, observation_provider: Any, question: str, - state: dict[str, Any], + state: Any, input_fn: Any = input, print_fn: Any = print, ) -> None: @@ -351,22 +348,26 @@ def handle_vqa_query( 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 - :func:`push_log` (REPL panel scrollback) and ``print_fn`` (direct - stdout) — in autonomous question mode the panel redraw is suspended, + ``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. """ - from .steps import _generate_with_policy, _msgs_for_vqa # noqa: PLC0415 + 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.""" - push_log(state, line) - try: + if hasattr(state, "log"): + state.log(line) + else: + state.setdefault("log_lines", []).append(line) + with suppress(Exception): print_fn(line) - except Exception: # noqa: BLE001 - pass - if policy is None or not hasattr(policy, "select_message"): - report(" [warn] vqa: policy has no select_message — skipping") + if policy_adapter is None: + report(" [warn] vqa: no policy adapter — skipping") return observation: dict | None = None @@ -383,19 +384,14 @@ def handle_vqa_query( # 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. - answer = _generate_with_policy( - policy, - _msgs_for_vqa(question), - observation=observation, - state=state, - label="vqa gen", - ) + 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 = parse_vqa_answer(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") diff --git a/src/lerobot/scripts/lerobot_pi052_runtime.py b/src/lerobot/scripts/lerobot_pi052_runtime.py index 01fb8a66f..a6339b44f 100644 --- a/src/lerobot/scripts/lerobot_pi052_runtime.py +++ b/src/lerobot/scripts/lerobot_pi052_runtime.py @@ -12,1677 +12,14 @@ # 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. -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. - -Examples --------- - -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 \\ - --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:: - - uv run lerobot-pi052-runtime \\ - --policy.path=pepijn223/pi052_hirobot_super_poulain_tool2 \\ - --dataset.repo_id=pepijn223/super_poulain_annotated \\ - --dataset.episode=0 \\ - --no_robot \\ - --task="please clean the kitchen" - -With a real robot:: - - uv run lerobot-pi052-runtime \\ - --policy.path=... \\ - --robot.type=so101 --robot.port=/dev/tty.usbmodem... \\ - --tts.voice=alba - -``--policy.path`` accepts either a local directory or a Hugging Face -Hub repo id. ``--dataset.repo_id`` likewise. - -Tool dispatch (TTS via ``SayTool``) is enabled by default when -``pocket-tts`` is installed; pass ``--no_tts`` to disable. -""" +"""Entry point for ``lerobot-pi052-runtime``.""" from __future__ import annotations -import argparse -import logging import sys -from typing import Any, Callable - -logger = logging.getLogger("lerobot.pi052.runtime") - - -def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: - p = argparse.ArgumentParser( - description=( - "Interactive REPL runtime for a trained PI052 hierarchical " - "VLA checkpoint." - ), - ) - 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``." - ), - ) - p.add_argument( - "--dataset.repo_id", - dest="dataset_repo_id", - type=str, - default=None, - help=( - "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 " - "trained on." - ), - ) - p.add_argument( - "--dataset.episode", - dest="dataset_episode", - type=int, - default=0, - help="Episode index to walk through (default: 0).", - ) - p.add_argument( - "--dataset.start_frame", - dest="dataset_start_frame", - type=int, - default=0, - help="Frame index within the episode to start from (default: 0).", - ) - p.add_argument( - "--dataset.advance_per_tick", - dest="dataset_advance_per_tick", - type=int, - default=1, - help=( - "How many dataset frames to advance per runtime tick. The " - "default of 1 means the runtime walks the episode forward " - "frame by frame; set to 0 to freeze on ``start_frame``." - ), - ) - p.add_argument( - "--dataset.augment_at_inference", - dest="dataset_augment_at_inference", - action="store_true", - help=( - "Apply the same torchvision-v2 ColorJitter / SharpnessJitter " - "/ RandomAffine pipeline that training used to each dataset " - "frame fed to the policy. Use to test whether the LM head " - "generalises under the augmentation distribution it was " - "supervised on — if dry-run still produces coherent subtask " - "text with this flag on, the head has learned beyond exact " - "frames; if it collapses to '\\n' the head is hyper-specific " - "to the unperturbed training samples." - ), - ) - p.add_argument( - "--task", - dest="task", - type=str, - default=None, - help=( - "Initial task. When given, the startup task picker is skipped " - "and this task is used directly. If omitted, the picker is " - "shown (or the first stdin line is treated as the task)." - ), - ) - p.add_argument( - "--mode", - dest="mode", - type=str, - choices=["action", "paused"], - default=None, - help=( - "Start-up run mode. 'action' runs the robot immediately on " - "--task; 'paused' (the default) comes up at the command line " - "with the robot idle. Flip any time with /action and /pause." - ), - ) - p.add_argument( - "--no_robot", - action="store_true", - help="Skip robot connection — language-only / dry-run mode.", - ) - # --- Real-robot mode args ---------------------------------------- - # Setting ``--robot.type`` flips the runtime into autonomous mode: - # it connects to the robot, builds an observation provider that - # reads ``robot.get_observation()`` instead of dataset frames, and - # 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. - p.add_argument( - "--robot.type", - dest="robot_type", - type=str, - default=None, - help=( - "Robot config choice (e.g. ``so101``, ``so101_follower``). " - "When set, the runtime drives the actual robot at " - "``--ctrl_hz`` instead of running the dataset-driven dry-run " - "REPL. Implies ``--autonomous`` unless ``--no_robot`` is also " - "passed (in which case the flag is ignored). See " - "``lerobot.robots`` for available choices." - ), - ) - p.add_argument( - "--robot.port", - dest="robot_port", - type=str, - default=None, - help="Serial port for the robot (e.g. ``/dev/tty.usbmodem...``).", - ) - p.add_argument( - "--robot.id", - dest="robot_id", - type=str, - default=None, - help="Optional robot identifier (passed through to ``RobotConfig.id``).", - ) - p.add_argument( - "--robot.cameras", - dest="robot_cameras", - type=str, - default=None, - help=( - "Optional JSON dict describing camera configs to attach to " - 'the robot (e.g. ``\'{"top": {"type": "opencv", "index": 0}}\'``). ' - "Camera keys MUST match the ``observation.images.*`` features " - "the policy was trained on." - ), - ) - p.add_argument( - "--robot.max_relative_target", - dest="robot_max_relative_target", - type=str, - default=None, - help=( - "Safety clip on per-motor relative motion, passed through to " - "``RobotConfig.max_relative_target``. Accepts either a float " - "(applied to every motor — e.g. ``5.0`` degrees) or a JSON " - "object mapping motor names to caps " - '(e.g. ``\'{"shoulder_pan": 5, "gripper": 30}\'``). The ' - "robot driver clips each commanded position relative to the " - "current measured position before sending — same kill-switch " - "``lerobot-record`` uses. Default ``None`` = no clipping." - ), - ) - p.add_argument( - "--auto_start", - action="store_true", - help=( - "Skip the ``Press ENTER to start`` confirmation prompt before " - "the autonomous control loop begins. Off by default — having " - "to confirm catches a lot of stupid mistakes (wrong policy, " - "wrong robot, robot not at home pose)." - ), - ) - p.add_argument( - "--no_tts", - action="store_true", - help="Disable the ``say`` tool dispatch.", - ) - p.add_argument( - "--tts.voice", - dest="tts_voice", - type=str, - default="alba", - help="Pocket-tts voice name (or path to a .wav for cloning).", - ) - p.add_argument( - "--chunk_hz", - type=float, - default=1.0, - help=( - "Action-chunk generation rate (Hz). Default ``1.0`` — one " - "new chunk per second. Lower = less inference cost / " - "smoother behaviour but longer reaction time to changes. " - "Higher = fresher actions / more inference cost; cap at " - "~1/(forward-pass latency)." - ), - ) - p.add_argument("--ctrl_hz", type=float, default=50.0, help="Action dispatch rate.") - p.add_argument( - "--high_level_hz", - type=float, - default=1.0, - help="High-level subtask generation rate.", - ) - p.add_argument( - "--subtask_chunks_per_gen", - type=int, - default=1, - help=( - "Throttle subtask gen to once every N action-chunk boundaries. " - "Default 1 = regenerate the subtask on every chunk refresh. " - "Set to 5 to run ~5 flow-matching action chunks per LM-head " - "subtask gen — saves compute and avoids re-planning trajectories " - "mid-grasp when a subtask is still valid across multiple chunks." - ), - ) - p.add_argument( - "--max_ticks", - type=int, - default=None, - help="Stop after N ticks (debug / smoke-test).", - ) - p.add_argument( - "--text_min_new_tokens", - type=int, - default=0, - help=( - "Debug knob for under-trained checkpoints: force the LM head " - "to emit at least N non-EOS tokens before EOS is allowed. " - "Use when the head's prior at position 0 still favours EOS " - "(short training run on a chat-pretrained backbone). 3-5 " - "is usually enough to reveal whether the model has real " - "subtask-token mass under the EOS argmax." - ), - ) - p.add_argument( - "--text_temperature", - type=float, - default=0.0, - help=( - "Sampling temperature for high-level text gen. 0 = greedy " - "argmax (default, matches training). Set 0.3-0.7 with an " - "under-trained checkpoint to escape stuck-at-EOS argmax." - ), - ) - p.add_argument( - "--text_top_p", - type=float, - default=1.0, - help="Nucleus filtering for high-level text gen.", - ) - p.add_argument("-v", "--verbose", action="store_true", help="Enable DEBUG logging.") - return p.parse_args(argv) - - -# Columns the runtime supplies itself via its own message stream — strip -# them so ``RenderMessagesStep`` / ``PI052TextTokenizerStep`` are no-ops. -_RUNTIME_OWNED_LANGUAGE_COLS = ("language_persistent", "language_events") - - -def _strip_runtime_owned_language_cols(sample: dict) -> None: - """In-place drop of language columns the runtime owns at inference.""" - for k in _RUNTIME_OWNED_LANGUAGE_COLS: - sample.pop(k, None) - - -def _select_observation_to_device(sample: dict, device: Any) -> dict: - """Filter to ``observation.*`` keys and move tensors to ``device``.""" - import torch # noqa: PLC0415 - - return { - k: v.to(device) if isinstance(v, torch.Tensor) else v - for k, v in sample.items() - if isinstance(k, str) and k.startswith("observation.") - } - - -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). - - Returns ``(policy, preprocessor, postprocessor, ds_meta)``. - ``preprocessor`` / ``postprocessor`` / ``ds_meta`` are ``None`` - when no dataset is provided (rare — needed for autonomous robot - mode to have action-denormalisation stats). - """ - from lerobot.configs import PreTrainedConfig # noqa: PLC0415 - from lerobot.policies.factory import make_policy, make_pre_post_processors # noqa: PLC0415 - - cfg = PreTrainedConfig.from_pretrained(policy_path) - cfg.pretrained_path = policy_path - - ds_meta = None - preprocessor = None - postprocessor = None - if dataset_repo_id is not None: - from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata # noqa: PLC0415 - - ds_meta = LeRobotDatasetMetadata(dataset_repo_id) - policy = make_policy(cfg, ds_meta=ds_meta) - # ``pretrained_path=None`` rebuilds fresh — the saved - # ``policy_preprocessor.json`` doesn't round-trip - # ``RenderMessagesStep.recipe``. Stats come from the dataset - # the user is feeding through, so normalisation is consistent. - preprocessor, postprocessor = make_pre_post_processors( - cfg, - pretrained_path=None, - dataset_stats=ds_meta.stats, - ) - else: - from lerobot.policies.factory import get_policy_class # noqa: PLC0415 - - policy_cls = get_policy_class(cfg.type) - policy = policy_cls.from_pretrained(policy_path, config=cfg) - policy.to(cfg.device) - - policy.eval() - return policy, preprocessor, postprocessor, ds_meta - - -def _build_observation_provider( - *, - dataset_repo_id: str, - episode: int, - start_frame: int, - advance_per_tick: int, - preprocessor: Any, - device: str, - augment: bool = False, -) -> Callable[[], dict | None]: - """Build a closure that feeds dataset frames into the runtime. - - Each call returns a preprocessed observation batch (images + - state, batched, on the policy's device, normalized) suitable for - ``policy.select_action`` and ``policy.select_message``. The - closure walks the chosen episode forward by ``advance_per_tick`` - frames per call, looping back to the episode start when it falls - off the end. - - The dataset's ``language_persistent`` / ``language_events`` - columns are stripped before the sample reaches the preprocessor, - so ``RenderMessagesStep`` and ``PI052TextTokenizerStep`` are - no-ops; the runtime supplies its own messages from current state. - """ - from lerobot.datasets.lerobot_dataset import LeRobotDataset # noqa: PLC0415 - - ds = LeRobotDataset(dataset_repo_id, episodes=[episode]) - if len(ds) == 0: - raise ValueError(f"Dataset {dataset_repo_id!r} episode {episode} is empty.") - - # Optional: apply the same torchvision-v2 augmentation pipeline - # that training used, so dry-run sees frames from the augmented - # support region (not just the unperturbed dataset frames). When - # the LM head still generates coherent text under this, it has - # learned over the augmentation distribution — the *opposite* of - # the "memorised one specific frame per supervision" failure - # mode. When it collapses to ``\n`` here too, the head is hyper- - # specific to the unperturbed training samples and only the - # retrain can help. - inference_aug = None - if augment: - from lerobot.transforms import ( # noqa: PLC0415 - ImageTransforms, - ImageTransformsConfig, - ) - - aug_cfg = ImageTransformsConfig(enable=True) - inference_aug = ImageTransforms(aug_cfg) - ds.set_image_transforms(inference_aug) - logger.warning( - "dry-run augmentation ENABLED — frames will be jittered " - "(brightness/contrast/saturation/hue/sharpness/affine) " - "before going to the policy" - ) - - state = {"cursor": max(0, min(start_frame, len(ds) - 1))} - - def _provider() -> dict | None: - idx = state["cursor"] - if advance_per_tick > 0: - state["cursor"] = (idx + advance_per_tick) % len(ds) - - sample = ds[idx] - _strip_runtime_owned_language_cols(sample) - - if preprocessor is not None: - sample = preprocessor(sample) - - return _select_observation_to_device(sample, device) - - return _provider - - -def _bootstrap_state_from_dataset( - *, - dataset_repo_id: str, - episode: int, - start_frame: int, -) -> dict[str, str]: - """Pull task / active plan / active memory / active subtask at ``start_frame``. - - The model is heavily memorised on the exact training prompts the - recipe rendered from this dataset (canonical task wording, - persistent atoms emitted earlier in the episode). Reconstructing - that state at REPL startup lets the runtime's first prompt line - up with what training looked like — without it the model sees an - out-of-distribution prompt and falls back to its dominant - training mode (VQA JSON spam). - """ - from lerobot.datasets.lerobot_dataset import LeRobotDataset # noqa: PLC0415 - - ds = LeRobotDataset(dataset_repo_id, episodes=[episode]) - if len(ds) == 0: - return {} - idx = max(0, min(start_frame, len(ds) - 1)) - sample = ds[idx] - - out: dict[str, str] = {} - task = sample.get("task") - if isinstance(task, str) and task.strip(): - out["task"] = task - - persistent = sample.get("language_persistent") or [] - # ``persistent`` is the broadcast slice of the episode; pick the - # *latest* row of each style whose ``timestamp`` is ≤ the - # frame's timestamp (matches the renderer's ``active_at`` - # semantics). - try: - frame_ts = ( - float(sample["timestamp"]) - if not hasattr(sample["timestamp"], "item") - else sample["timestamp"].item() - ) - except Exception: # noqa: BLE001 - frame_ts = float("inf") - - by_style: dict[str, tuple[float, str]] = {} - for row in persistent: - style = row.get("style") - ts = row.get("timestamp") - content = row.get("content") - if not (style and content) or ts is None: - continue - try: - ts_f = float(ts) - except (TypeError, ValueError): - continue - if ts_f > frame_ts: - continue - prev = by_style.get(style) - if prev is None or ts_f >= prev[0]: - by_style[style] = (ts_f, content) - for style, (_, content) in by_style.items(): - if style in {"plan", "memory", "subtask"}: - out[style] = content - return out - - -def _select_task_interactively( - *, - ds_meta: Any, - bootstrap_task: str | None, -) -> str | None: - """Ask the operator which task to run at startup. - - Behaviour: - - * If a dataset is loaded, build a numbered menu of every unique task - string in ``ds_meta.tasks`` (canonical bootstrap task listed first - as the default). Add a ``[c] type a custom task`` option. - * If no dataset is loaded, show a plain ``Enter task:`` prompt. - * Non-TTY runs (scripts, pipes) skip the prompt and return the - bootstrap task so the existing "first stdin line becomes task" - flow in ``_run_repl`` / ``_run_autonomous`` still works. - - Returns the chosen task string, or ``None`` when the operator declines - to pick one (Ctrl-D / empty + no default). - """ - options: list[str] = [] - seen: set[str] = set() - if bootstrap_task: - options.append(bootstrap_task) - seen.add(bootstrap_task) - if ds_meta is not None and getattr(ds_meta, "tasks", None) is not None: - try: - for t in list(ds_meta.tasks.index): - if isinstance(t, str) and t and t not in seen: - options.append(t) - seen.add(t) - except Exception: # noqa: BLE001 — defensive: tasks shape varies - pass - - if not sys.stdin.isatty(): - # Scripted / piped run: no interactive prompt; fall back to the - # bootstrap default (may be None — REPL handles that). - return bootstrap_task - - print("\n[pi052] Select startup task:", flush=True) - if options: - for i, opt in enumerate(options, 1): - marker = " (dataset default)" if opt == bootstrap_task else "" - print(f" [{i}] {opt}{marker}", flush=True) - print(" [c] type a custom task", flush=True) - prompt = "Choice [1]: " if bootstrap_task else "Choice: " - else: - print(" (no tasks available from dataset)", flush=True) - prompt = "Enter task: " - - while True: - try: - choice = input(prompt).strip() - except EOFError: - print(flush=True) - return bootstrap_task - - # No dataset options at all: the entered line *is* the task. - if not options: - return choice or None - - # Empty input: take the default (item 1) when there is one. - if not choice: - return options[0] if bootstrap_task else None - - if choice.lower() in ("c", "custom"): - try: - free = input("Enter task: ").strip() - except EOFError: - print(flush=True) - return bootstrap_task - if free: - return free - # Empty free-form input → loop back to the menu. - continue - - if choice.isdigit(): - idx = int(choice) - if 1 <= idx <= len(options): - return options[idx - 1] - - print( - f" invalid choice {choice!r}; pick 1–{len(options)} or 'c'.", - flush=True, - ) - - -def _build_robot( - *, - robot_type: str, - robot_port: str | None, - robot_id: str | None, - robot_cameras_json: str | None, - robot_max_relative_target: str | None, -): - """Build and connect a robot from CLI args. - - Mirrors how ``lerobot-record`` builds a robot but takes the args - flat from argparse instead of through draccus, so the runtime - keeps its plain ``--key=value`` CLI surface. ``max_relative_target`` - is passed through to the RobotConfig — the driver itself clips each - commanded joint position relative to the current measured one - before issuing it on the bus. - """ - import importlib # noqa: PLC0415 - import json # noqa: PLC0415 - import pkgutil # noqa: PLC0415 - - import lerobot.robots as _robots_pkg # noqa: PLC0415 - from lerobot.robots import ( # noqa: PLC0415 - RobotConfig, - make_robot_from_config, - ) - - # ``RobotConfig._choice_registry`` is populated lazily — each robot's - # ``config_.py`` calls ``@RobotConfig.register_subclass`` at - # import time. ``lerobot.robots/__init__.py`` doesn't import the - # individual robot packages, so ``get_choice_class(robot_type)`` - # raises ``KeyError`` until at least one robot module has been - # imported. Mirror what ``make_robot_from_config`` does internally: - # walk the robots package's submodules and import each so the - # decorator side-effect runs. Slow only on the first call (~200ms - # for ~10 dataclass modules); negligible for an autonomous run that - # then loops at ctrl_hz for minutes. - for _modinfo in pkgutil.iter_modules(_robots_pkg.__path__): - if _modinfo.name.startswith("_"): - continue - try: - importlib.import_module(f"lerobot.robots.{_modinfo.name}") - except Exception as exc: # noqa: BLE001 - logger.debug("could not import lerobot.robots.%s: %s", _modinfo.name, exc) - - try: - cls = RobotConfig.get_choice_class(robot_type) - except KeyError as exc: - available = sorted(RobotConfig._choice_registry.keys()) - raise ValueError(f"Unknown robot type {robot_type!r}. Available choices: {available}") from exc - kwargs: dict[str, Any] = {} - if robot_port: - kwargs["port"] = robot_port - if robot_id: - kwargs["id"] = robot_id - if robot_cameras_json: - try: - cameras_raw = json.loads(robot_cameras_json) - except json.JSONDecodeError as exc: - raise ValueError( - f"--robot.cameras must be a JSON object, got {robot_cameras_json!r}: {exc}" - ) from exc - # ``RobotConfig`` expects ``cameras: dict[str, CameraConfig]`` — - # each inner value must be an actual ``CameraConfig`` subclass - # instance, not a raw dict. Look up the matching subclass via - # ``CameraConfig.get_choice_class()`` (registered by - # ``@CameraConfig.register_subclass`` decorators on each camera - # backend's config) and instantiate it. Mirror the lazy-import - # pattern from above so the registry is populated. - import lerobot.cameras as _cameras_pkg # noqa: PLC0415 - from lerobot.cameras import CameraConfig # noqa: PLC0415 - - for _modinfo in pkgutil.iter_modules(_cameras_pkg.__path__): - if _modinfo.name.startswith("_"): - continue - try: - importlib.import_module(f"lerobot.cameras.{_modinfo.name}") - except Exception as exc: # noqa: BLE001 - logger.debug("could not import lerobot.cameras.%s: %s", _modinfo.name, exc) - - cameras: dict[str, Any] = {} - for cam_name, cam_dict in cameras_raw.items(): - if not isinstance(cam_dict, dict): - raise ValueError(f"camera {cam_name!r} value must be a dict, got {cam_dict!r}") - cam_dict = dict(cam_dict) # don't mutate caller's parsed JSON - cam_type = cam_dict.pop("type", None) - if cam_type is None: - raise ValueError( - f"camera {cam_name!r} is missing a 'type' field (e.g. 'opencv', 'intelrealsense')" - ) - try: - cam_cls = CameraConfig.get_choice_class(cam_type) - except KeyError as exc: - available = sorted(CameraConfig._choice_registry.keys()) - raise ValueError( - f"camera {cam_name!r}: unknown type {cam_type!r}. Available choices: {available}" - ) from exc - cameras[cam_name] = cam_cls(**cam_dict) - kwargs["cameras"] = cameras - if robot_max_relative_target: - # Accept either a bare float (uniform cap) or a JSON object - # (per-motor cap). Matches ``RobotConfig.max_relative_target``'s - # ``float | dict[str, float] | None`` shape. - s = robot_max_relative_target.strip() - try: - if s.startswith("{"): - kwargs["max_relative_target"] = json.loads(s) - else: - kwargs["max_relative_target"] = float(s) - except (json.JSONDecodeError, ValueError) as exc: - raise ValueError( - f"--robot.max_relative_target must be a float or JSON dict, " - f"got {robot_max_relative_target!r}: {exc}" - ) from exc - cfg = cls(**kwargs) - robot = make_robot_from_config(cfg) - robot.connect() - return robot - - -def _build_robot_observation_provider( - *, - robot, - preprocessor: Any, - device: str, - task: str | None, - ds_features: dict[str, Any] | None, -) -> Callable[[], dict | None]: - """Closure that reads from the robot, runs the policy preprocessor. - - Each call: ``robot.get_observation()`` (raw per-joint + per-camera - dict, possibly with scalar floats) → ``build_inference_frame`` - (extract the keys the dataset declared, reshape per-joint floats - into a single ``observation.state`` vector, prefix camera keys - with ``observation.images.``, convert to tensors with batch dim - on device) → wrap in an ``EnvTransition`` (the preprocessor - pipeline is transition-shaped, keyed by ``TransitionKey``) → - preprocessor (rename, normalise) → unwrap and return the flat - observation batch ``policy.select_action`` / ``policy.select_message`` - consume. - """ - import torch # noqa: PLC0415 - - from lerobot.policies.utils import ( # noqa: PLC0415 - build_inference_frame, - prepare_observation_for_inference, - ) - - torch_device = torch.device(device) if isinstance(device, str) else device - robot_type = getattr(robot, "robot_type", None) or getattr(getattr(robot, "config", None), "type", None) - - # Pre-compute the camera-key → target (H, W) map from - # ``ds_features``. The training distribution sees frames at the - # recorded resolution (e.g. 480×640); a live Mac/USB camera will - # almost always hand us a different native size (720p / 1080p). - # PI052's internal ``resize_with_pad(512, 512)`` does pad the - # input to a fixed canvas, but the *geometry* of that pad differs - # by input aspect ratio — top/left padding varies, so the visual - # tokens at each tile carry different content than what the model - # saw at training. The action expert tolerates this (flow head - # rides broad geometry); the LM head, supervised much more - # tightly on visual features, goes out of distribution and the - # head's distribution at position 0 collapses to its dominant - # mode (a memorised ``\n``-only run in this checkpoint). - _resize_logged = {"done": False} - target_image_shapes: dict[str, tuple[int, int]] = {} - if ds_features: - for fkey, fmeta in ds_features.items(): - if not isinstance(fmeta, dict): - continue - dtype = fmeta.get("dtype") - if dtype not in ("image", "video"): - continue - shape = fmeta.get("shape") - if not shape or len(shape) != 3: - continue - names = fmeta.get("names") or [] - # Feature schema stores either (H, W, C) or (C, H, W); - # disambiguate by the ``names`` ordering when present. - if names and len(names) == 3 and names[0] == "channels": - _, h, w = shape - else: - h, w, _ = shape - cam_key = fkey.removeprefix("observation.images.") - target_image_shapes[cam_key] = (int(h), int(w)) - - def _provider() -> dict | None: - try: - raw = robot.get_observation() - except Exception as exc: # noqa: BLE001 - logger.warning("robot.get_observation failed: %s", exc) - return None - - # The runtime supplies messages itself; strip any language - # columns the robot stream may carry through. - _strip_runtime_owned_language_cols(raw) - - # Force-match the training-time visual distribution: - # every camera frame the model trained on came from the - # dataset at its recorded (H, W). Resize the live frame to - # that exact shape so the downstream resize_with_pad geometry - # matches training. Without this the LM head is OOD on every - # tick. - if target_image_shapes: - try: - import cv2 as _cv2 # noqa: PLC0415 - import numpy as _np # noqa: PLC0415 - - # Snapshot the gate state at the start of the call: the - # camera info and startup-state warnings are meant to fire - # exactly once (operator sanity check), so gate them on - # the *previous* value rather than the post-loop value. - first_call = not _resize_logged["done"] - for cam_key, (target_h, target_w) in target_image_shapes.items(): - img = raw.get(cam_key) - if img is None or not isinstance(img, _np.ndarray): - continue - if img.ndim != 3: - continue - cur_h, cur_w = img.shape[:2] - if first_call: - logger.warning( - "camera %s: live=%dx%d, training=%dx%d (resize=%s)", - cam_key, - cur_h, - cur_w, - target_h, - target_w, - "yes" if (cur_h, cur_w) != (target_h, target_w) else "no — already matched", - ) - if (cur_h, cur_w) == (target_h, target_w): - continue - raw[cam_key] = _cv2.resize(img, (target_w, target_h), interpolation=_cv2.INTER_AREA) - _resize_logged["done"] = True - # Print the state vector once so the operator can eyeball - # it against the dataset's stats. State OOD is a real - # failure mode for VLAs — the prefix carries state via - # the projection layer, and a neutral home pose can - # easily sit a couple σ off the supervised support - # region. Gated on ``first_call`` so this doesn't spam - # every observation tick. - if first_call and "observation.state" in (ds_features or {}): - state_names = ds_features["observation.state"].get("names") or [] - state_vals = [raw.get(n) for n in state_names] - logger.warning( - "robot state at startup: %s", - { - n: round(v, 2) if isinstance(v, float) else v - for n, v in zip(state_names, state_vals, strict=False) - }, - ) - except Exception as exc: # noqa: BLE001 - logger.warning("camera resize to dataset shape failed: %s", exc) - - try: - if ds_features: - # Use the dataset's feature schema to pick the right - # raw keys and fold per-joint scalars into a single - # ``observation.state`` tensor. Then tensor-ise + - # device-place + add batch dim. - obs_tensors = build_inference_frame( - raw, - torch_device, - ds_features=ds_features, - task=task, - robot_type=robot_type, - ) - else: - # No dataset features available — fall back to the - # generic numpy-only path; only works when the robot - # already returns dataset-shaped keys. - obs_tensors = prepare_observation_for_inference( - raw, - torch_device, - task=task, - robot_type=robot_type, - ) - except Exception as exc: # noqa: BLE001 - logger.warning("observation prep failed: %s", exc) - return None - - if preprocessor is not None: - # ``PolicyProcessorPipeline`` defaults its ``to_transition`` - # to ``batch_to_transition``, which expects a *flat batch - # dict* keyed by ``observation.*`` / ``action`` / etc., and - # wraps it into an ``EnvTransition`` itself. Pre-wrapping - # here would just have ``batch_to_transition`` look for - # ``observation.*`` keys at top level, find none (they'd - # be nested under ``TransitionKey.OBSERVATION``), and - # produce an empty observation → ``ObservationProcessorStep`` - # bails. Pass the flat dict straight in; ``to_output`` - # gives us a flat dict back. - try: - processed = preprocessor(obs_tensors) - except Exception as exc: # noqa: BLE001 - logger.warning("preprocessor failed on robot observation: %s", exc) - return None - obs_tensors = processed if isinstance(processed, dict) else {} - - return _select_observation_to_device(obs_tensors, torch_device) - - return _provider - - -def _build_robot_action_executor( - *, - robot, - postprocessor: Any, - ds_meta: Any, -) -> Callable[[Any], None]: - """Closure that postprocesses an action and dispatches to the robot. - - Mirrors ``lerobot-record``'s ``predict_action`` tail: postprocess - (denormalise) → ``make_robot_action`` (tensor → ``{joint: value}`` - dict) → ``robot.send_action(...)``. Safety clipping happens *inside* - ``robot.send_action`` via the driver's ``max_relative_target`` - cap (passed in at ``RobotConfig`` construction time) — same place - ``lerobot-record`` enforces it. - """ - import torch # noqa: PLC0415 - - from lerobot.policies.utils import make_robot_action # noqa: PLC0415 - - def _executor(action: Any) -> None: - try: - if postprocessor is not None: - action = postprocessor(action) - if isinstance(action, torch.Tensor): - if action.ndim > 1 and action.shape[0] == 1: - action = action.squeeze(0) - action_dict = make_robot_action(action, ds_meta.features) - elif isinstance(action, dict): - action_dict = action - else: - logger.warning("unsupported action type %r — skipping", type(action)) - return - robot.send_action(action_dict) - except Exception as exc: # noqa: BLE001 - logger.error("robot.send_action failed: %s", exc, exc_info=True) - - return _executor - - -def _print_runtime_help() -> None: - """Print the slash-command reference.""" - print( - "[pi052] commands (arguments need no quotes):\n" - " /action run the robot; an argument switches to that task\n" - " /action resume the robot on the current task\n" - " /action run the robot for N seconds, then auto-pause\n" - " /pause pause the action loop — robot holds position\n" - " /question 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", - flush=True, - ) - - -def _is_number(text: str) -> bool: - """True if ``text`` parses as a float (a ``/action`` duration arg).""" - try: - float(text) - return True - except ValueError: - return False - - -def _strip_quotes(text: str) -> str: - """Strip one pair of surrounding quotes from a command argument.""" - text = text.strip() - if len(text) >= 2 and text[0] == text[-1] and text[0] in {'"', "'"}: - return text[1:-1].strip() - return text - - -def _clear_action_queue(runtime: Any) -> None: - """Drop any queued action chunk so nothing fires while paused.""" - queue = runtime.state.get("action_queue") - if hasattr(queue, "clear"): - queue.clear() - - -def _handle_slash_command(runtime: Any, line: str) -> bool: - """Dispatch the runtime slash commands. - - ``/action ["task"]`` run the robot; a quoted/bare argument sets a - new task, a bare number is a timed burst - (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). - """ - stripped = line.strip() - if not stripped.startswith("/"): - return False - head, _, rest = stripped.partition(" ") - cmd = head.lower() - rest = _strip_quotes(rest) - - if cmd in {"/action", "/act", "/run"}: - runtime.state["mode"] = "action" - if rest and _is_number(rest): - import time as _time # noqa: PLC0415 - - secs = float(rest) - runtime.state["action_deadline"] = _time.monotonic() + secs - print( - f"[pi052] action — running {secs:g}s, then auto-pause", - flush=True, - ) - else: - runtime.state["action_deadline"] = None - if rest: - runtime.set_task(rest) - # 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) - elif runtime.state.get("task"): - print( - f"[pi052] action — resuming: {runtime.state['task']!r}", - flush=True, - ) - else: - runtime.state["mode"] = "paused" - print( - "[pi052] no task set — use /action ", - flush=True, - ) - return True - - if cmd in {"/pause", "/p"}: - 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 " - "(e.g. /question point to the yellow cube)", - flush=True, - ) - return True - _run_vqa_query(runtime, rest) - return True - - if cmd in {"/help", "/?"}: - _print_runtime_help() - return True - 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=runtime.policy, - observation_provider=runtime.observation_provider, - question=question, - state=runtime.state, - ) - - -def _run_autonomous( - runtime: Any, - *, - robot, - auto_start: bool, - initial_task: str | None, - max_ticks: int | None, -) -> int: - """Drive the runtime continuously at ``ctrl_hz`` while accepting - stdin events in the foreground. - - Different from ``_run_repl`` (dataset dry-run): the policy needs - to keep generating action chunks at ``chunk_hz`` and dispatching - them at ``ctrl_hz`` regardless of whether the user is typing, so - ``runtime.run()`` runs in a background thread and stdin handling - happens here in the main thread. - """ - import threading # noqa: PLC0415 - import time # noqa: PLC0415 - - # Only gate on ENTER when the robot will actually move at startup - # (``--mode=action``). The default is paused — the command line - # comes up immediately and nothing moves until ``/action``. - 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. " - ) - except (EOFError, KeyboardInterrupt): - print("\n[pi052] aborted before start", flush=True) - return 130 - - if initial_task: - runtime.set_task(initial_task) - - thread = threading.Thread( - target=runtime.run, - kwargs={"max_ticks": max_ticks}, - name="pi052-runtime-loop", - daemon=True, - ) - thread.start() - - # Capture log lines flushed by the runtime each tick into a - # bounded scrollback that the panel renderer prints inside the - # rule block. Without this, ``runtime._flush_logs`` just calls - # ``print(...)`` which the 2 Hz panel redraw clears immediately — - # so failure messages from generation (e.g. ``[warn] subtask gen - # failed: ...``) flash for ≤ 0.5 s and disappear, leaving the - # operator with no idea why ``last_raw`` stays empty. - _scrollback: list[str] = [] - _scrollback_max = 12 - - def _flush_into_scrollback() -> None: - for line in runtime.state.get("log_lines") or []: - _scrollback.append(line) - # Trim to the cap so the panel doesn't grow unbounded. - if len(_scrollback) > _scrollback_max: - del _scrollback[: len(_scrollback) - _scrollback_max] - - runtime._flush_logs = _flush_into_scrollback # type: ignore[method-assign] - - redraw = _make_state_panel_renderer(runtime, mode_label="autonomous", scrollback=_scrollback) - redraw() - print( - " [autonomous] /action to run · /pause to stop · " - "/question to ask · /help · stop", - flush=True, - ) - - # Background panel-redraw thread so state changes from the runtime - # loop (subtask refresh, plan update, etc.) are visible without the - # user typing anything. - # - # In ``/vlm`` mode the action loop is paused — nothing changes in the - # background — so the timer redraw is suspended entirely. That keeps - # the screen stable while the operator types a VQA question and the - # interactive camera prompt, instead of the panel clearing the - # prompt every tick. - _panel_stop = threading.Event() - - def _panel_loop() -> None: - while not _panel_stop.is_set(): - st = runtime.state - if st.get("mode", "action") == "action": - # Timed burst (``/action ``): once the deadline - # passes, auto-revert to question mode and clear the - # queue so the robot stops. - deadline = st.get("action_deadline") - if deadline is not None and time.monotonic() >= deadline: - st["mode"] = "paused" - st["action_deadline"] = None - queue = st.get("action_queue") - if hasattr(queue, "clear"): - queue.clear() - print( - "\n[pi052] timed action elapsed — paused", - flush=True, - ) - else: - try: - redraw() - # Re-print the prompt the redraw just cleared so - # the operator always has a visible ``> ``. - print("> ", end="", flush=True) - except Exception: # noqa: BLE001 - pass - _panel_stop.wait(0.7) - - panel_thread = threading.Thread(target=_panel_loop, name="pi052-panel-redraw", daemon=True) - panel_thread.start() - - try: - while thread.is_alive(): - try: - line = input("> ").strip() - except EOFError: - break - if not line: - continue - lower = line.lower() - if lower in {"stop", "quit", "exit"}: - break - # The runtime is command-driven: /action "task", /pause, - # /question "...", /help. ``_handle_slash_command`` runs the - # VQA query inline for /question (the action loop is paused - # first, so the policy isn't in concurrent use). - if _handle_slash_command(runtime, line): - try: - redraw() - except Exception: # noqa: BLE001 - pass - continue - # A bare (non-slash) line is treated as a user interjection - # — the trained ``user_interjection_response`` path. ``stop`` - # already handled above; everything else routes here. - if runtime.state.get("task"): - runtime.state["recent_interjection"] = line - runtime.state.setdefault("events_this_tick", []).append("user_interjection") - else: - print( - "[pi052] no task yet — use /action to start", - flush=True, - ) - except KeyboardInterrupt: - print("\n[pi052] interrupt — stopping", flush=True) - finally: - _panel_stop.set() - runtime.stop() - # Give the loop a moment to drain. - for _ in range(10): - if not thread.is_alive(): - break - time.sleep(0.1) - try: - robot.disconnect() - print("[pi052] robot disconnected", flush=True) - except Exception as exc: # noqa: BLE001 - print(f"[pi052] WARNING: robot.disconnect raised {exc}", flush=True) - - return 0 - - -def _make_state_panel_renderer( - runtime: Any, - *, - mode_label: str, - scrollback: list[str] | None = None, -) -> Callable[[list[str] | None], None]: - """Return a closure that prints the task/subtask/plan/memory panel. - - Used by both ``_run_repl`` (dry-run, called per user input) and - ``_run_autonomous`` (real robot, called on a 2 Hz timer + - whenever the user types). Centralises the visual format so the - two modes look identical. - """ - from rich.console import Console # noqa: PLC0415 - - console = Console(highlight=False) - - def _redraw(robot_lines: list[str] | None = None) -> None: - console.clear() - 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") - # 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[/] ask · [bold]/help[/] · [bold]stop[/]" - ) - else: - console.print( - " [dim]commands:[/] [bold]/action[/] run · " - "[bold]/question[/] ask · [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"), - ("current_plan", "plan"), - ("current_memory", "memory"), - ): - value = st.get(key) - if value: - console.print(f" [bold cyan]{label:<8}[/] {value}") - else: - console.print(f" [dim]{label:<8} (not set)[/]") - queue_len = ( - len(st["action_queue"]) - 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}[/]" - ) - - # Overfit / memorisation diagnostics. The high-level steps - # surface the raw generation each time they fire (even when - # rejected as gibberish or unchanged), plus repeat/rejection - # counters. Rule of thumb: - # - # * subtask repeat ≥ ~5 and queue_len cycles fully → model - # can't move past current subtask (memorised one phase - # of the task — classic overfit signature) - # * subtask gibberish climbing → LM head collapsed to - # chat-template fragments / one-token salads - # * last raw differs from accepted → at least the LM is - # varying, the gibberish filter is doing its job - raw_subtask = st.get("last_subtask_raw") - sub_rep = int(st.get("subtask_repeat_count") or 0) - sub_gib = int(st.get("subtask_gibberish_count") or 0) - sub_empty = int(st.get("subtask_empty_count") or 0) - if raw_subtask is not None or sub_rep or sub_gib or sub_empty: - raw_display = (raw_subtask or "(empty)")[:80] - color = "yellow" if (sub_rep >= 3 or sub_gib >= 3 or sub_empty >= 3) else "dim" - console.print( - f" [{color}]subtask diag repeat:{sub_rep} " - f"gibberish:{sub_gib} empty:{sub_empty} " - f"last_raw: {raw_display!r}[/]" - ) - - # Same diagnostics for memory and plan when available. - mem_gib = int(st.get("memory_gibberish_count") or 0) - plan_gib = int(st.get("plan_gibberish_count") or 0) - if mem_gib or plan_gib: - 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. - if scrollback: - for line in scrollback: - console.print(f" [magenta]{line.rstrip()}[/]") - console.rule(style="cyan") - if robot_lines: - for line in robot_lines: - console.print(f" [magenta]{line.strip()}[/]") - console.print() - if not st.get("task"): - console.print( - " [dim]Type [bold]/action [/bold] to begin, " - "[bold]/question [/bold] to ask, /help for commands, " - "stop to exit.[/]" - ) - - return _redraw - - -def _build_tools(no_tts: bool, tts_voice: str) -> dict[str, Any]: - """Instantiate the tools declared on this dataset/policy.""" - if no_tts: - return {} - try: - from lerobot.tools import SayTool # noqa: PLC0415 - - return {"say": SayTool(voice=tts_voice)} - except Exception as exc: # noqa: BLE001 - logger.warning("Could not initialise SayTool (%s) — speech disabled.", exc) - return {} - - -def _silence_noisy_loggers() -> None: - """Drop chatty third-party loggers down to WARNING. - - HuggingFace / httpx / urllib3 emit one log line per HTTP request, - which the REPL has to print between the state block and the - prompt — completely unreadable. We never need that detail in the - REPL and the user can opt back into it via ``-v`` (verbose mode - keeps DEBUG on the lerobot loggers but still gates the noisy ones - here unless they explicitly want them). - """ - for name in ( - "httpcore", - "httpcore.connection", - "httpcore.http11", - "httpcore.proxy", - "httpx", - "urllib3", - "urllib3.connectionpool", - "huggingface_hub", - "huggingface_hub.repocard", - "huggingface_hub.file_download", - "transformers", - "transformers.modeling_utils", - "transformers.tokenization_utils_base", - "datasets", - "filelock", - ): - logging.getLogger(name).setLevel(logging.WARNING) - - # The robot's relative-goal-position clamp warning fires *every* - # dispatch tick on a memorised model — the LM is trying to jump - # the wrist far past where max_relative_target allows, so the - # warning floods the panel at ~30 Hz. Promote it from WARNING to - # DEBUG: the dispatch counter on the panel already tells the - # operator the loop is running, and the panel itself shows - # whether motion is happening. If anyone needs the per-action - # clamp detail, ``-v`` puts it back via DEBUG. - logging.getLogger("lerobot.robots.utils").setLevel(logging.ERROR) - - -def main(argv: list[str] | None = None) -> int: - args = _parse_args(argv) - logging.basicConfig( - level=logging.DEBUG if args.verbose else logging.INFO, - format="%(asctime)s %(levelname)s %(message)s", - ) - _silence_noisy_loggers() - - 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 " - "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) - policy, preprocessor, postprocessor, ds_meta = _load_policy_and_preprocessor( - args.policy_path, args.dataset_repo_id - ) - - # Bootstrap the canonical task from the dataset whenever one is - # provided, so the interactive picker below can offer it as the - # default. The model is memorised on the exact training wording, so - # matching it is what gets recall to fire. - bootstrap_state: dict[str, str] = {} - if args.dataset_repo_id is not None: - bootstrap_state = _bootstrap_state_from_dataset( - dataset_repo_id=args.dataset_repo_id, - episode=args.dataset_episode, - start_frame=args.dataset_start_frame, - ) - - # Interactive task picker. Skipped when ``--task`` is already set on - # the CLI (scripted runs and explicit overrides win). When no task - # was passed, prompt the operator: pick from the dataset's tasks or - # type a custom one. Non-TTY runs fall back to the bootstrap task - # silently — the existing "first stdin line becomes task" flow in - # ``_run_repl`` / ``_run_autonomous`` still handles the no-default - # case. - if not args.task: - chosen = _select_task_interactively( - ds_meta=ds_meta, - bootstrap_task=bootstrap_state.get("task"), - ) - if chosen: - args.task = chosen - print(f"[pi052] 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`` - # forces a mode. The operator drives it with /action, /pause and - # /question. - startup_mode = args.mode or "paused" - - observation_provider: Callable[[], dict | None] | None = None - robot_executor: Callable[[Any], None] | None = None - robot = None - - if autonomous_mode: - print( - f"[pi052] connecting to robot.type={args.robot_type} port={args.robot_port}", - flush=True, - ) - robot = _build_robot( - robot_type=args.robot_type, - robot_port=args.robot_port, - robot_id=args.robot_id, - robot_cameras_json=args.robot_cameras, - robot_max_relative_target=args.robot_max_relative_target, - ) - observation_provider = _build_robot_observation_provider( - robot=robot, - preprocessor=preprocessor, - device=str(getattr(policy.config, "device", "cpu")), - task=args.task, - ds_features=ds_meta.features if ds_meta is not None else None, - ) - robot_executor = _build_robot_action_executor( - robot=robot, - postprocessor=postprocessor, - ds_meta=ds_meta, - ) - elif args.dataset_repo_id is not None: - print( - f"[pi052] streaming observations from {args.dataset_repo_id} " - f"episode={args.dataset_episode} " - f"start_frame={args.dataset_start_frame}", - flush=True, - ) - observation_provider = _build_observation_provider( - dataset_repo_id=args.dataset_repo_id, - episode=args.dataset_episode, - start_frame=args.dataset_start_frame, - advance_per_tick=args.dataset_advance_per_tick, - preprocessor=preprocessor, - device=str(getattr(policy.config, "device", "cpu")), - augment=getattr(args, "dataset_augment_at_inference", False), - ) - - tools = _build_tools(args.no_tts, args.tts_voice) - if tools: - print(f"[pi052] tools loaded: {list(tools)}", flush=True) - - from lerobot.policies.pi052.inference import PI052Runtime # noqa: PLC0415 - - runtime = PI052Runtime( - policy=policy, - tools=tools, - observation_provider=observation_provider, - robot_executor=robot_executor, - # No background event collector — the REPL drives ticks - # synchronously after each user input (REPL mode). Autonomous - # mode runs ``runtime.run()`` in a thread; stdin events are - # injected from the foreground. - event_collector=None, - chunk_hz=args.chunk_hz, - ctrl_hz=args.ctrl_hz, - high_level_hz=args.high_level_hz, - ) - # Stash text-gen knobs on the state dict so the high-level steps - # (which read state) can pick them up and forward them to - # policy.select_message. Letting the operator try - # ``--text_min_new_tokens=5 --text_temperature=0.6`` on an - # under-trained checkpoint without recompiling. - 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: HighLevelSubtaskFwd fires 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 - # the previous one is still being executed. - runtime.state["subtask_chunks_per_gen"] = max( - 1, int(getattr(args, "subtask_chunks_per_gen", 1) or 1) - ) - # Apply the startup mode chosen above the task picker. - runtime.state["mode"] = startup_mode - if args.task: - runtime.set_task(args.task) - # Seed the current subtask from the dataset so the first chunk — - # before ``HighLevelSubtaskFwd`` has run — has a real subtask to - # condition the action expert on instead of falling back to the - # bare task. Plan and memory are NOT seeded: the current recipe - # trains neither, no inference step consumes them, and seeding - # them only put a stale plan in the status panel that does - # nothing. - if bootstrap_state.get("subtask"): - runtime.state["current_subtask"] = bootstrap_state["subtask"] - - if autonomous_mode: - return _run_autonomous( - runtime, - robot=robot, - auto_start=args.auto_start, - initial_task=args.task, - max_ticks=args.max_ticks, - ) - # Fire one full pipeline tick at startup so the obs diagnostic - # *and* the subtask generation actually run before the REPL - # blocks on stdin. The REPL otherwise only ticks on user input, - # which made the dry-run bisection test (does the LM head produce - # text at start_frame=0?) require typing something. Doing - # ``step_once`` here means the diag row populates without any - # manual interaction. - if observation_provider is not None: - try: - startup_logs = runtime.step_once() - except Exception as exc: # noqa: BLE001 - 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) - - -def _run_repl(runtime: Any, *, initial_task: str | None, max_ticks: int | None) -> int: - """Claude-Code-style block REPL. - - Each turn redraws a status block (task / subtask / plan / memory) - at the top, prints any robot log lines that came in since the last - turn, then asks for input on a clean ``> `` prompt at the bottom. - No live region, no panel re-renders, no rendering races with HTTP - log lines — just clear-screen + reprint each turn, the way a - chat-style REPL is meant to look. - """ - try: - from rich.console import Console # noqa: PLC0415 - except ImportError: - print( - "[pi052] 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") - # Keep a local ``console`` just for the styled input prompt; the - # state panel is owned by the shared renderer. - console = Console(highlight=False) - - last_logs: list[str] = [] - _redraw() - if initial_task is None: - # Already shown the help line in _redraw when task is None. - pass - ticks_done = 0 - try: - while True: - try: - line = console.input("[bold cyan]> [/]").strip() - except EOFError: - break - if not line: - _redraw(last_logs) - continue - lower = line.lower() - if lower in {"stop", "quit", "exit"}: - break - - # Command-driven: /action "task", /pause, /question "...", - # /help. ``_handle_slash_command`` runs the VQA query inline - # for /question (single-threaded REPL — no concurrency). - if _handle_slash_command(runtime, line): - last_logs = list(runtime.state.get("log_lines") or []) - _redraw(last_logs) - ticks_done += 1 - if max_ticks is not None and ticks_done >= max_ticks: - break - continue - - # A bare (non-slash) line is a user interjection — needs a - # task to be meaningful. - if not runtime.state.get("task"): - print( - "[pi052] no task yet — use /action ", - flush=True, - ) - _redraw(last_logs) - continue - runtime.state["recent_interjection"] = line - runtime.state.setdefault("events_this_tick", []).append("user_interjection") - - last_logs = runtime.step_once() or [] - _redraw(last_logs) - - ticks_done += 1 - if max_ticks is not None and ticks_done >= max_ticks: - break - except KeyboardInterrupt: - console.print("\n[dim]interrupted[/]") - console.print("[dim]runtime stopped[/]") - return 0 +from lerobot.policies.pi052.inference.runtime_cli import main if __name__ == "__main__": sys.exit(main()) diff --git a/tests/policies/language_conditioned/test_runtime.py b/tests/policies/language_conditioned/test_runtime.py new file mode 100644 index 000000000..ef1a9137b --- /dev/null +++ b/tests/policies/language_conditioned/test_runtime.py @@ -0,0 +1,88 @@ +from lerobot.policies.language_conditioned import ( + LanguageConditionedRuntime, + RuntimeState, + ToolCall, + VQAResult, +) + + +class FakeAdapter: + def __init__(self): + self.updated = False + self.text_calls = [] + + def select_action(self, observation, state): + assert observation == {"observation.state": 1} + assert state.task == "clean" + return ["a0", "a1"] + + def select_text(self, kind, observation, state, user_text=None): + self.text_calls.append((kind, user_text)) + return "new plan ok" + + def parse_tool_calls(self, text): + assert text == "new plan ok" + return [ToolCall("say", {"text": "ok"})] + + 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") + + +class FakeTool: + def __init__(self): + self.calls = [] + + def call(self, args): + self.calls.append(args) + + +def test_runtime_tick_updates_language_enqueues_and_dispatches_action(): + adapter = FakeAdapter() + executed = [] + runtime = LanguageConditionedRuntime( + policy_adapter=adapter, + observation_provider=lambda: {"observation.state": 1}, + action_executor=executed.append, + ) + runtime.set_task("clean") + + logs = runtime.step_once() + + assert adapter.updated + assert runtime.state.language_context["subtask"] == "pick cup" + assert executed == ["a0"] + assert list(runtime.state.action_queue) == ["a1"] + assert " subtask: pick cup" in logs + + +def test_runtime_handles_user_interjection_and_dispatches_tools(): + adapter = FakeAdapter() + tool = FakeTool() + runtime = LanguageConditionedRuntime( + policy_adapter=adapter, + observation_provider=lambda: {"observation.state": 1}, + tools={"say": tool}, + ) + runtime.set_task("clean") + runtime.state.extra["recent_interjection"] = "please say ok" + runtime.state.emit("user_interjection") + + logs = runtime.step_once() + + assert ("interjection", "please say ok") in adapter.text_calls + assert runtime.state.language_context["plan"] == "new plan ok" + assert tool.calls == [{"text": "ok"}] + assert " speech: ok" in logs + + +def test_runtime_state_aliases_legacy_keys_to_language_context(): + state = RuntimeState() + state["current_subtask"] = "open drawer" + state["current_memory"] = "drawer open" + + assert state.get("current_subtask") == "open drawer" + assert state.language_context == {"subtask": "open drawer", "memory": "drawer open"} diff --git a/tests/policies/pi052/test_pi052_runtime_adapter.py b/tests/policies/pi052/test_pi052_runtime_adapter.py new file mode 100644 index 000000000..a3ef2eab7 --- /dev/null +++ b/tests/policies/pi052/test_pi052_runtime_adapter.py @@ -0,0 +1,54 @@ +from types import SimpleNamespace + +from lerobot.policies.language_conditioned import RuntimeState +from lerobot.policies.pi052.inference.pi052_adapter import PI052PolicyAdapter, split_plan_and_say + + +def test_pi052_adapter_builds_recipe_prompts_from_runtime_state(): + adapter = PI052PolicyAdapter(policy=object()) + state = RuntimeState( + task="clean the kitchen", + language_context={"memory": "cup moved", "plan": "pick then place"}, + extra={"prior_subtask": "pick the cup"}, + ) + + assert adapter.messages_for("subtask", state) == [{"role": "user", "content": "clean the kitchen"}] + assert adapter.messages_for("memory", state) == [ + {"role": "user", "content": "clean the kitchen"}, + {"role": "assistant", "content": "Previous memory: cup moved"}, + {"role": "user", "content": "Completed subtask: pick the cup"}, + ] + assert adapter.messages_for("interjection", state, user_text="wait") == [ + {"role": "user", "content": "clean the kitchen"}, + {"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_parses_say_tool_calls_and_plan_text(): + adapter = PI052PolicyAdapter(policy=object()) + text = "Move to the sink. heading to the sink" + + assert split_plan_and_say(text) == ("Move to the sink.", "heading to the sink") + assert adapter.parse_tool_calls(text)[0].name == "say" + assert adapter.parse_tool_calls(text)[0].arguments == {"text": "heading to the sink"} + assert adapter.plan_from_text(text) == "Move to the sink." + + +def test_pi052_runtime_cli_smoke_does_not_load_model(monkeypatch): + from lerobot.policies.pi052.inference import runtime_cli + + fake_policy = SimpleNamespace(config=SimpleNamespace(device="cpu")) + + monkeypatch.setattr( + runtime_cli, + "_load_policy_and_preprocessor", + lambda policy_path, dataset_repo_id: (fake_policy, None, None, None), + ) + monkeypatch.setattr(runtime_cli, "_build_tools", lambda no_tts, tts_voice: {}) + monkeypatch.setattr(runtime_cli, "_run_repl", lambda runtime, initial_task, max_ticks: 0) + + assert runtime_cli.main(["--policy.path=fake", "--no_robot", "--task=clean", "--max_ticks=0"]) == 0