diff --git a/pyproject.toml b/pyproject.toml index 1d2a05ff8..8d6b6c2ce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -346,8 +346,8 @@ lerobot-edit-dataset="lerobot.scripts.lerobot_edit_dataset:main" lerobot-setup-can="lerobot.scripts.lerobot_setup_can:main" lerobot-annotate="lerobot.scripts.lerobot_annotate:main" lerobot-rollout="lerobot.scripts.lerobot_rollout:main" -# Interactive hierarchical-VLA runtime for PI052 (PaliGemma backbone). -lerobot-pi052-runtime="lerobot.scripts.lerobot_pi052_runtime:main" +# Interactive high/low-level runtime for language-conditioned policies (pi052, ...). +lerobot-language-runtime="lerobot.scripts.lerobot_language_runtime:main" # ---------------- Tool Configurations ---------------- diff --git a/src/lerobot/policies/pi052/configuration_pi052.py b/src/lerobot/policies/pi052/configuration_pi052.py index 030f039b6..1e8911604 100644 --- a/src/lerobot/policies/pi052/configuration_pi052.py +++ b/src/lerobot/policies/pi052/configuration_pi052.py @@ -33,7 +33,7 @@ This is the dual-head co-training pattern from the paper: with α = 10.0 per § IV.D of arxiv:2504.16054. The π0.5 model splits inference into a text-prediction step followed by an action-prediction step, which the multi-rate runtime (``lerobot.runtime``, via the -``lerobot-pi052-runtime`` CLI) drives at separate rates. +``lerobot-language-runtime`` CLI) drives at separate rates. """ from dataclasses import dataclass diff --git a/src/lerobot/policies/pi052/inference/__init__.py b/src/lerobot/policies/pi052/inference/__init__.py index 21391157f..33a68afb0 100644 --- a/src/lerobot/policies/pi052/inference/__init__.py +++ b/src/lerobot/policies/pi052/inference/__init__.py @@ -16,7 +16,7 @@ The runtime, REPL, and CLI are policy-agnostic and live in :mod:`lerobot.runtime`. PI052 supplies only :class:`PI052PolicyAdapter`; -the ``lerobot-pi052-runtime`` entry point wires it into +the ``lerobot-language-runtime`` entry point wires it into :func:`lerobot.runtime.cli.run`. """ diff --git a/src/lerobot/policies/pi052/inference/pi052_adapter.py b/src/lerobot/policies/pi052/inference/pi052_adapter.py index 6b263c327..cd22c1ae6 100644 --- a/src/lerobot/policies/pi052/inference/pi052_adapter.py +++ b/src/lerobot/policies/pi052/inference/pi052_adapter.py @@ -12,28 +12,34 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""PI052 adapter for the generic language-conditioned runtime.""" +"""PI052 adapter for the generic language-conditioned runtime. + +Supplies only the PI052-specific primitives — acting, text generation, +and prompt templates. The high-level control loop (throttling, output +rejection, the subtask -> memory cascade) is inherited from +:class:`lerobot.runtime.adapter.BaseLanguageAdapter`. +""" from __future__ import annotations import logging -import re -from dataclasses import dataclass from typing import Any from lerobot.runtime import RuntimeState +from lerobot.runtime.adapter import BaseLanguageAdapter 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: +class PI052PolicyAdapter(BaseLanguageAdapter): """Runtime bridge for PI052 policies.""" - policy: Any + # PaliGemma's ```` prior dominates the first token on a small + # text-CE budget; suppress it for prose kinds (VQA would keep it, but + # the runtime no longer does interactive VQA). + LOC_SUPPRESS_KINDS = frozenset({"subtask", "memory", "interjection"}) def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: subtask = state.language_context.get("subtask") or state.task or "" @@ -49,99 +55,40 @@ class PI052PolicyAdapter: batch[OBS_LANGUAGE_ATTENTION_MASK] = text_batch["lang_masks"] return self.policy.predict_action_chunk(batch) - def select_text( + def generate_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) + messages = self.build_messages(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"}, + min_new_tokens=self.gen.min_new_tokens, + temperature=self.gen.temperature, + top_p=self.gen.top_p, + suppress_loc_tokens=kind in self.LOC_SUPPRESS_KINDS, ) - def plan_from_text(self, text: str) -> str: - plan, _speech = split_plan_and_say(text) - return "" if looks_like_gibberish(plan) else plan - - 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( + def build_messages( self, kind: str, state: RuntimeState, *, user_text: str | None = None, ) -> list[dict[str, Any]]: - if kind == "subtask": + if kind in ("subtask", "plan"): 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']}", - } + {"role": "assistant", "content": f"Previous memory: {state.language_context['memory']}"} ) if state.extra.get("prior_subtask"): messages.append( @@ -157,8 +104,6 @@ class PI052PolicyAdapter: if user_text: messages.append({"role": "user", "content": user_text}) return messages - if kind == "plan": - return [{"role": "user", "content": state.task or ""}] raise ValueError(f"Unknown PI052 text kind: {kind}") @@ -254,36 +199,3 @@ def _generate_with_policy( 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 diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py index 6907ddbe5..bd517d598 100644 --- a/src/lerobot/policies/pi052/modeling_pi052.py +++ b/src/lerobot/policies/pi052/modeling_pi052.py @@ -18,7 +18,7 @@ ``text_labels`` next to the flow loss (L = H(x, f_θ_text) + α·flow, α via ``config.flow_loss_weight``) and :meth:`select_message` for AR text generation. The multi-rate runtime in ``lerobot.policies.pi052.inference`` -(``lerobot-pi052-runtime`` CLI) drives ``predict_action_chunk`` + +(``lerobot-language-runtime`` CLI) drives ``predict_action_chunk`` + ``select_message``. See :class:`PI052Config` for the knobs. """ @@ -2014,10 +2014,9 @@ class PI052Policy(PreTrainedPolicy): return out def _generate_low_level_subtask(self, obs_i: dict[str, Tensor], task: str, i: int) -> str: - from .inference.pi052_adapter import ( # noqa: PLC0415 - _generate_with_policy, - looks_like_gibberish as _looks_like_gibberish, - ) + from lerobot.runtime.adapter import looks_like_gibberish as _looks_like_gibberish # noqa: PLC0415 + + from .inference.pi052_adapter import _generate_with_policy # noqa: PLC0415 msg = "" if task: diff --git a/src/lerobot/runtime/__init__.py b/src/lerobot/runtime/__init__.py index ce0aaa130..db38d9a5a 100644 --- a/src/lerobot/runtime/__init__.py +++ b/src/lerobot/runtime/__init__.py @@ -14,11 +14,13 @@ """Policy-agnostic high/low-level runtime for language-conditioned policies. -The tick loop, REPL, and interactive CLI here are policy-independent; a -policy plugs in by implementing :class:`LanguageConditionedPolicyAdapter` -and calling :func:`lerobot.runtime.cli.run` with an adapter factory. +The tick loop, REPL, and interactive CLI here are policy-independent. A +policy plugs in by subclassing :class:`BaseLanguageAdapter` (or satisfying +:class:`LanguageConditionedPolicyAdapter` directly) and registering it in +:mod:`lerobot.runtime.registry`; ``lerobot-language-runtime`` then serves it. """ +from .adapter import BaseLanguageAdapter, GenerationConfig, LanguageDiagnostics from .language_runtime import ( LanguageConditionedPolicyAdapter, LanguageConditionedRuntime, @@ -28,8 +30,11 @@ from .language_runtime import ( ) __all__ = [ + "BaseLanguageAdapter", + "GenerationConfig", "LanguageConditionedPolicyAdapter", "LanguageConditionedRuntime", + "LanguageDiagnostics", "RuntimeState", "Tick", "TickClock", diff --git a/src/lerobot/runtime/adapter.py b/src/lerobot/runtime/adapter.py new file mode 100644 index 000000000..fd9f10070 --- /dev/null +++ b/src/lerobot/runtime/adapter.py @@ -0,0 +1,195 @@ +# 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. + +"""Policy adapter base class for the language-conditioned runtime. + +The runtime loop drives the *control algorithm* (throttling, output +rejection, the subtask -> memory cascade, diagnostics) and delegates the +*policy primitives* (act, generate text) to an adapter. :class:`BaseLanguageAdapter` +implements the algorithm once; a policy subclasses it and supplies: + +* :meth:`select_action` — observation + language context -> action chunk +* :meth:`generate_text` — a text stream (``kind``) -> decoded string +* :meth:`build_messages` — the prompt for each ``kind`` + +A policy that needs full control can instead satisfy the +:class:`LanguageConditionedPolicyAdapter` protocol directly. +""" + +from __future__ import annotations + +import re +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any + +from .language_runtime import RuntimeState + +_SAY_RE = re.compile(r"<\s*say\s*>(.*?)<\s*/\s*say\s*>", re.IGNORECASE | re.DOTALL) + + +@dataclass +class GenerationConfig: + """Text-generation knobs, fixed for the lifetime of an adapter. + + These are configuration (set once from the CLI), not per-tick runtime + state — they live on the adapter, never in :class:`RuntimeState`. + """ + + min_new_tokens: int = 0 + temperature: float = 0.0 + top_p: float = 1.0 + chunks_per_regen: int = 1 # regenerate the language context every N action chunks + + +@dataclass +class LanguageDiagnostics: + """Rejection / repeat counters surfaced in the runtime panel. + + Keyed by text ``kind`` (``subtask`` / ``memory`` / ...) so the same + accounting works for any cascade shape. + """ + + last_raw: dict[str, str] = field(default_factory=dict) + empty: dict[str, int] = field(default_factory=dict) + gibberish: dict[str, int] = field(default_factory=dict) + repeat: int = 0 + + def _bump(self, table: dict[str, int], kind: str) -> int: + table[kind] = table.get(kind, 0) + 1 + return table[kind] + + +class BaseLanguageAdapter(ABC): + """Batteries-included adapter: generic high-level control, policy primitives abstract.""" + + def __init__(self, policy: Any, gen: GenerationConfig | None = None) -> None: + self.policy = policy + self.gen = gen or GenerationConfig() + self.diag = LanguageDiagnostics() + self._chunks_until_regen = 0 + + # --- policy primitives (subclass supplies) --------------------------- + + @abstractmethod + def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: + """Produce an action chunk from the observation + current language context.""" + + @abstractmethod + def generate_text( + self, + kind: str, + observation: dict[str, Any] | None, + state: RuntimeState, + user_text: str | None = None, + ) -> str: + """Generate one text stream (``kind``) and return the decoded string.""" + + # --- generic control algorithm (runtime calls these) ---------------- + + def update_language_state(self, observation: dict[str, Any] | None, state: RuntimeState) -> None: + """Throttled regeneration of the language context (subtask / memory / ...).""" + if self._chunks_until_regen > 0: + self._chunks_until_regen -= 1 + return + self._chunks_until_regen = max(1, self.gen.chunks_per_regen) - 1 + self._regenerate_context(observation, state) + + def handle_interjection( + self, user_text: str, observation: dict[str, Any] | None, state: RuntimeState + ) -> None: + """React to a mid-run user message by regenerating the plan.""" + out = self.generate_text("interjection", observation, state, user_text=user_text) + plan = self.plan_from_text(out) + if plan: + state.set_context("plan", plan, label="plan") + + def plan_from_text(self, text: str) -> str: + """Strip ```` speech markers and reject gibberish plans.""" + plan, _speech = split_plan_and_say(text) + return "" if looks_like_gibberish(plan) else plan + + # --- overridable cascade + shared helpers --------------------------- + + def _regenerate_context(self, observation: dict[str, Any] | None, state: RuntimeState) -> None: + """Default hierarchy: regenerate the subtask, then memory when it changes. + + Override for a policy with a different language hierarchy. + """ + subtask = self._generate_filtered("subtask", observation, state) + if subtask is None: + return + previous = state.language_context.get("subtask") + if not state.set_context("subtask", subtask, label="subtask"): + self.diag.repeat += 1 + return + self.diag.repeat = 0 + if previous: + state.extra["prior_subtask"] = previous + memory = self._generate_filtered("memory", observation, state) + if memory is not None: + state.set_context("memory", memory, label="memory") + + def _generate_filtered( + self, kind: str, observation: dict[str, Any] | None, state: RuntimeState + ) -> str | None: + """Generate one ``kind``, record diagnostics, drop empty / gibberish output.""" + text = self.generate_text(kind, observation, state) + self.diag.last_raw[kind] = text or "" + if not text: + count = self.diag._bump(self.diag.empty, kind) + if count == 1 or count % 5 == 0: + state.log(f" [info] {kind} gen returned empty (x{count})") + return None + if looks_like_gibberish(text): + count = self.diag._bump(self.diag.gibberish, kind) + if count == 1 or count % 30 == 0: + state.log(f" [info] {kind} gen rejected (gibberish x{count}): {text[:60]!r}") + return None + return text + + +def looks_like_gibberish(text: str) -> bool: + """Heuristic filter for malformed / collapsed LM-head output.""" + 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]: + """Split ``plan speech`` into ``(plan, 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 diff --git a/src/lerobot/runtime/cli.py b/src/lerobot/runtime/cli.py index 83537457c..413364fbe 100644 --- a/src/lerobot/runtime/cli.py +++ b/src/lerobot/runtime/cli.py @@ -17,7 +17,7 @@ Policy-agnostic CLI over :class:`lerobot.runtime.LanguageConditionedRuntime`. A policy wires it up with :func:`run`, passing an adapter factory (``policy -> LanguageConditionedPolicyAdapter``); see -``lerobot.scripts.lerobot_pi052_runtime`` for the PI052 entry point. +``lerobot.scripts.lerobot_language_runtime`` for the entry point. Stdin is the user channel: type a task, then natural-language interjections. The runtime prints state changes (plan / subtask / @@ -29,7 +29,7 @@ Examples Dry run on a Hub checkpoint, no robot connected — useful for sanity- checking text generation:: - uv run lerobot-pi052-runtime \\ + uv run lerobot-language-runtime \\ --policy.path= \\ --no_robot \\ --task="please clean the kitchen" @@ -37,7 +37,7 @@ checking text generation:: Same, but feed real frames from an annotated dataset so plan / subtask / memory generation runs against actual video + state:: - uv run lerobot-pi052-runtime \\ + uv run lerobot-language-runtime \\ --policy.path= \\ --dataset.repo_id= \\ --dataset.episode=0 \\ @@ -46,7 +46,7 @@ Same, but feed real frames from an annotated dataset so plan / subtask With a real robot:: - uv run lerobot-pi052-runtime \\ + uv run lerobot-language-runtime \\ --policy.path=... \\ --robot.type=so101 --robot.port=/dev/tty.usbmodem... @@ -63,6 +63,7 @@ from collections.abc import Callable from contextlib import suppress from typing import Any +from .adapter import GenerationConfig from .language_runtime import LanguageConditionedPolicyAdapter, LanguageConditionedRuntime from .repl import _emit @@ -1201,36 +1202,26 @@ def _make_state_panel_renderer( dispatched = int(st.get("actions_dispatched") or 0) console.print(f" [dim]queued actions: {queue_len} dispatched: {dispatched}[/]") - # Overfit / memorisation diagnostics. The high-level steps - # surface the raw generation each time they fire (even when - # 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}[/]") + # Overfit / memorisation diagnostics from the adapter. High repeat + # + fully cycling queue ⇒ stuck on one subtask (memorised a phase); + # climbing gibberish ⇒ LM head collapsed to chat-template salads. + diag = getattr(runtime.policy_adapter, "diag", None) + if diag is not None: + raw_subtask = diag.last_raw.get("subtask") + sub_rep = int(diag.repeat) + sub_gib = int(diag.gibberish.get("subtask", 0)) + sub_empty = int(diag.empty.get("subtask", 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}[/]" + ) + mem_gib = int(diag.gibberish.get("memory", 0)) + if mem_gib: + console.print(f" [dim]gen rejects memory:{mem_gib}[/]") console.rule(style="cyan") # Runtime scrollback — log lines pushed from generation steps # (warnings, gibberish rejections, plan speech). Last N lines, @@ -1294,16 +1285,18 @@ def _silence_noisy_loggers() -> None: def run( argv: list[str] | None = None, *, - adapter_factory: Callable[[Any], LanguageConditionedPolicyAdapter], - panel_label: str = "Runtime", - prog: str | None = None, + adapter_factory: Callable[[Any, GenerationConfig], LanguageConditionedPolicyAdapter] | None = None, + panel_label: str | None = None, + prog: str = "lerobot-language-runtime", ) -> int: """Run the interactive language-conditioned runtime CLI. - A policy wires this up by passing ``adapter_factory`` — a callable - that turns a loaded policy into a :class:`LanguageConditionedPolicyAdapter` - (typically the adapter class itself). ``panel_label`` names the state - panel; ``prog`` sets the argparse program name for ``--help``. + ``adapter_factory`` turns ``(policy, GenerationConfig)`` into a + :class:`LanguageConditionedPolicyAdapter` (typically the adapter class). + When ``None`` it is resolved from :mod:`lerobot.runtime.registry` by the + loaded policy's type, so a single ``lerobot-language-runtime`` entry + point serves every registered policy. ``panel_label`` defaults to the + policy type. """ args = _parse_args(argv, prog=prog) logging.basicConfig( @@ -1327,6 +1320,14 @@ def run( args.policy_path, args.dataset_repo_id ) + policy_type = getattr(policy.config, "type", None) + if adapter_factory is None: + from .registry import get_language_adapter_factory # noqa: PLC0415 + + adapter_factory = get_language_adapter_factory(policy_type) + if panel_label is None: + panel_label = str(policy_type or "runtime").upper() + # 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 @@ -1406,8 +1407,18 @@ def run( augment=getattr(args, "dataset_augment_at_inference", False), ) + # Text-generation knobs are fixed config, passed to the adapter at + # construction — not smuggled through per-tick runtime state. Lets the + # operator try e.g. ``--text_temperature=0.6 --subtask_chunks_per_gen=5`` + # on an under-trained checkpoint without recompiling. + gen_config = GenerationConfig( + min_new_tokens=int(args.text_min_new_tokens or 0), + temperature=float(args.text_temperature or 0.0), + top_p=float(args.text_top_p or 1.0), + chunks_per_regen=max(1, int(args.subtask_chunks_per_gen or 1)), + ) runtime = LanguageConditionedRuntime( - policy_adapter=adapter_factory(policy), + policy_adapter=adapter_factory(policy, gen_config), observation_provider=observation_provider, action_executor=robot_executor, # No background event collector — the REPL drives ticks @@ -1419,20 +1430,6 @@ def run( 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 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: diff --git a/src/lerobot/runtime/language_runtime.py b/src/lerobot/runtime/language_runtime.py index 7410d16fd..6c4bda8c2 100644 --- a/src/lerobot/runtime/language_runtime.py +++ b/src/lerobot/runtime/language_runtime.py @@ -117,17 +117,20 @@ class RuntimeState: class LanguageConditionedPolicyAdapter(Protocol): - """Policy-specific bridge used by :class:`LanguageConditionedRuntime`.""" + """The contract the runtime loop depends on. + + :class:`lerobot.runtime.adapter.BaseLanguageAdapter` provides a + batteries-included implementation; a policy can satisfy this protocol + directly for full control. + """ 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 update_language_state(self, observation: dict[str, Any] | None, state: RuntimeState) -> None: ... + + def handle_interjection( + self, user_text: str, observation: dict[str, Any] | None, state: RuntimeState + ) -> None: ... @dataclass @@ -260,12 +263,9 @@ class LanguageConditionedRuntime: 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) + self.policy_adapter.update_language_state(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}") @@ -279,12 +279,7 @@ class LanguageConditionedRuntime: 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 - plan = getattr(self.policy_adapter, "plan_from_text", lambda value: value)(out) - if plan: - self.state.set_context("plan", plan, label="plan") + self.policy_adapter.handle_interjection(text, observation, self.state) self.state.extra["recent_interjection"] = None def maybe_enqueue_action_chunk(self, *, force: bool = False) -> None: diff --git a/src/lerobot/runtime/registry.py b/src/lerobot/runtime/registry.py new file mode 100644 index 000000000..a4969a0a9 --- /dev/null +++ b/src/lerobot/runtime/registry.py @@ -0,0 +1,42 @@ +# 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. + +"""Registry mapping a policy type to its language-runtime adapter. + +Kept as import strings (resolved lazily) so ``lerobot-language-runtime`` +never imports a policy package until it actually loads that policy — the +same pattern as :mod:`lerobot.policies.factory`. +""" + +from __future__ import annotations + +import importlib +from collections.abc import Callable +from typing import Any + +_ADAPTERS: dict[str, str] = { + "pi052": "lerobot.policies.pi052.inference.pi052_adapter:PI052PolicyAdapter", +} + + +def get_language_adapter_factory(policy_type: str) -> Callable[..., Any]: + """Return the adapter class registered for ``policy_type``.""" + spec = _ADAPTERS.get(policy_type) + if spec is None: + raise ValueError( + f"No language-runtime adapter registered for policy type {policy_type!r}. " + f"Registered: {sorted(_ADAPTERS)}. Add an entry to lerobot.runtime.registry." + ) + module_path, class_name = spec.split(":") + return getattr(importlib.import_module(module_path), class_name) diff --git a/src/lerobot/scripts/lerobot_pi052_runtime.py b/src/lerobot/scripts/lerobot_language_runtime.py similarity index 65% rename from src/lerobot/scripts/lerobot_pi052_runtime.py rename to src/lerobot/scripts/lerobot_language_runtime.py index 2b4efe619..b3a5b0089 100644 --- a/src/lerobot/scripts/lerobot_pi052_runtime.py +++ b/src/lerobot/scripts/lerobot_language_runtime.py @@ -13,11 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Entry point for ``lerobot-pi052-runtime``. +"""Entry point for ``lerobot-language-runtime``. -Wires PI052's adapter into the generic runtime CLI. A new -language-conditioned policy adds its own such entry point with its -adapter — no runtime/REPL code to copy. +Policy-agnostic: the runtime resolves the right adapter from the loaded +policy's type via :mod:`lerobot.runtime.registry`. A new +language-conditioned policy just registers its adapter there — no new +script needed. """ from __future__ import annotations @@ -26,15 +27,9 @@ import sys def main(argv: list[str] | None = None) -> int: - from lerobot.policies.pi052.inference import PI052PolicyAdapter from lerobot.runtime.cli import run - return run( - argv, - adapter_factory=PI052PolicyAdapter, - panel_label="PI052", - prog="lerobot-pi052-runtime", - ) + return run(argv) if __name__ == "__main__": diff --git a/tests/policies/pi052/test_pi052_runtime_adapter.py b/tests/policies/pi052/test_pi052_runtime_adapter.py index cc5a8702a..ca4a1a109 100644 --- a/tests/policies/pi052/test_pi052_runtime_adapter.py +++ b/tests/policies/pi052/test_pi052_runtime_adapter.py @@ -1,7 +1,8 @@ from types import SimpleNamespace -from lerobot.policies.pi052.inference.pi052_adapter import PI052PolicyAdapter, split_plan_and_say +from lerobot.policies.pi052.inference.pi052_adapter import PI052PolicyAdapter from lerobot.runtime import RuntimeState +from lerobot.runtime.adapter import split_plan_and_say def test_pi052_adapter_builds_recipe_prompts_from_runtime_state(): @@ -12,13 +13,13 @@ def test_pi052_adapter_builds_recipe_prompts_from_runtime_state(): extra={"prior_subtask": "pick the cup"}, ) - assert adapter.messages_for("subtask", state) == [{"role": "user", "content": "clean the kitchen"}] - assert adapter.messages_for("memory", state) == [ + assert adapter.build_messages("subtask", state) == [{"role": "user", "content": "clean the kitchen"}] + assert adapter.build_messages("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") == [ + assert adapter.build_messages("interjection", state, user_text="wait") == [ {"role": "user", "content": "clean the kitchen"}, {"role": "assistant", "content": "Previous plan:\npick then place"}, {"role": "user", "content": "wait"}, @@ -33,12 +34,12 @@ def test_pi052_adapter_strips_say_markers_from_plan_text(): assert adapter.plan_from_text(text) == "Move to the sink." -def test_pi052_runtime_cli_smoke_does_not_load_model(monkeypatch): - """The pi052 entry wires its adapter into the generic runtime CLI.""" +def test_language_runtime_cli_smoke_does_not_load_model(monkeypatch): + """The general entry resolves the pi052 adapter from the registry by policy type.""" from lerobot.runtime import cli - from lerobot.scripts import lerobot_pi052_runtime + from lerobot.scripts import lerobot_language_runtime - fake_policy = SimpleNamespace(config=SimpleNamespace(device="cpu")) + fake_policy = SimpleNamespace(config=SimpleNamespace(device="cpu", type="pi052")) monkeypatch.setattr( cli, @@ -48,5 +49,6 @@ def test_pi052_runtime_cli_smoke_does_not_load_model(monkeypatch): monkeypatch.setattr(cli, "_run_repl", lambda runtime, **kwargs: 0) assert ( - lerobot_pi052_runtime.main(["--policy.path=fake", "--no_robot", "--task=clean", "--max_ticks=0"]) == 0 + lerobot_language_runtime.main(["--policy.path=fake", "--no_robot", "--task=clean", "--max_ticks=0"]) + == 0 ) diff --git a/tests/runtime/test_adapter.py b/tests/runtime/test_adapter.py new file mode 100644 index 000000000..78f19de68 --- /dev/null +++ b/tests/runtime/test_adapter.py @@ -0,0 +1,74 @@ +from lerobot.runtime import RuntimeState +from lerobot.runtime.adapter import BaseLanguageAdapter, GenerationConfig, looks_like_gibberish + + +class ScriptedAdapter(BaseLanguageAdapter): + """Base adapter whose text generation returns queued strings per kind.""" + + def __init__(self, scripts, gen=None): + super().__init__(policy=object(), gen=gen) + self.scripts = {k: list(v) for k, v in scripts.items()} + self.calls = [] + + def select_action(self, observation, state): + return None + + def generate_text(self, kind, observation, state, user_text=None): + self.calls.append(kind) + queue = self.scripts.get(kind, []) + return queue.pop(0) if queue else "" + + +def test_cascade_sets_subtask_then_memory(): + adapter = ScriptedAdapter({"subtask": ["pick the red cup"], "memory": ["the cup is grasped"]}) + state = RuntimeState(task="clean") + + adapter.update_language_state(None, state) + + assert state.language_context["subtask"] == "pick the red cup" + assert state.language_context["memory"] == "the cup is grasped" + assert adapter.calls == ["subtask", "memory"] + + +def test_gibberish_subtask_is_rejected_and_counted(): + adapter = ScriptedAdapter({"subtask": [":::: ::"], "memory": ["should not run"]}) + state = RuntimeState(task="clean") + + adapter.update_language_state(None, state) + + assert "subtask" not in state.language_context + assert adapter.diag.gibberish.get("subtask") == 1 + assert adapter.calls == ["subtask"] # memory never generated when subtask is rejected + + +def test_throttle_regenerates_every_n_chunks(): + adapter = ScriptedAdapter( + { + "subtask": ["pick the first cup", "pick the second cup"], + "memory": ["memory one two three", "memory four five six"], + }, + gen=GenerationConfig(chunks_per_regen=2), + ) + state = RuntimeState(task="clean") + + adapter.update_language_state(None, state) # generates + assert state.language_context["subtask"] == "pick the first cup" + adapter.update_language_state(None, state) # throttled — no generation + assert state.language_context["subtask"] == "pick the first cup" + adapter.update_language_state(None, state) # generates again + assert state.language_context["subtask"] == "pick the second cup" + + +def test_handle_interjection_sets_plan_and_strips_say(): + adapter = ScriptedAdapter({"interjection": ["turn to the left now heading left"]}) + state = RuntimeState(task="clean") + + adapter.handle_interjection("turn", None, state) + + assert state.language_context["plan"] == "turn to the left now" + + +def test_looks_like_gibberish_basic(): + assert looks_like_gibberish("") + assert looks_like_gibberish(":::: ::") + assert not looks_like_gibberish("pick up the red cube") diff --git a/tests/runtime/test_language_runtime.py b/tests/runtime/test_language_runtime.py index 8fd4e00fd..7d7dd8fba 100644 --- a/tests/runtime/test_language_runtime.py +++ b/tests/runtime/test_language_runtime.py @@ -7,21 +7,21 @@ from lerobot.runtime import ( class FakeAdapter: def __init__(self): self.updated = False - self.text_calls = [] + self.interjections = [] 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" - def update_language_state(self, observation, state): self.updated = True state.set_context("subtask", "pick cup", label="subtask") + def handle_interjection(self, user_text, observation, state): + self.interjections.append(user_text) + state.set_context("plan", "new plan", label="plan") + def test_runtime_tick_updates_language_enqueues_and_dispatches_action(): adapter = FakeAdapter() @@ -54,7 +54,7 @@ def test_runtime_handles_user_interjection(): runtime.step_once() - assert ("interjection", "please say ok") in adapter.text_calls + assert "please say ok" in adapter.interjections assert runtime.state.language_context["plan"] == "new plan"