refactor(runtime): template-method adapter base + policy registry; rename CLI

Make the policy adapter architecturally clean and set up a single general
entry point for any language-conditioned policy.

Adapter architecture (Template Method):
- New lerobot/runtime/adapter.py: BaseLanguageAdapter owns the generic
  control loop (throttle → generate → gibberish/empty reject → subtask→memory
  cascade → diagnostics) and plan_from_text/handle_interjection. A policy
  supplies only select_action + generate_text + build_messages. The
  subtask→memory cascade is an overridable hook (_regenerate_context).
- GenerationConfig (typed, constructor-time) replaces config smuggled through
  RuntimeState.extra (temperature/top_p/min_new_tokens/chunks_per_regen).
- LanguageDiagnostics (typed, keyed by kind) replaces ~8 loose state.extra
  counter keys; the panel reads it via the adapter.
- looks_like_gibberish + split_plan_and_say move to runtime (generic).

Contract:
- LanguageConditionedPolicyAdapter protocol now states the true contract
  (select_action, update_language_state, handle_interjection); the runtime
  drops both getattr fallbacks.
- PI052PolicyAdapter shrinks to just its primitives (132 → ~half).

General entry point:
- lerobot/runtime/registry.py maps policy type → adapter (lazy import).
- run() resolves the adapter from the registry by policy type and defaults
  the panel label to it, so one CLI serves every policy.
- Rename lerobot-pi052-runtime → lerobot-language-runtime (general script);
  a new policy just registers its adapter, no new script.

Tests: new tests/runtime/test_adapter.py covers throttle/reject/cascade/
interjection; adapter + runtime + CLI-smoke tests updated for the new shape.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Pepijn
2026-07-02 15:34:41 +02:00
parent 171e06c6ba
commit edc3a5eb4f
14 changed files with 438 additions and 222 deletions
+2 -2
View File
@@ -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 ----------------
@@ -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
@@ -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`.
"""
@@ -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 ``<locDDDD>`` 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
+4 -5
View File
@@ -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:
+8 -3
View File
@@ -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",
+195
View File
@@ -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 ``<say>`` 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 <say>speech</say>`` 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
+53 -56
View File
@@ -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=<repo-or-dir> \\
--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=<repo-or-dir> \\
--dataset.repo_id=<annotated-dataset> \\
--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:
+13 -18
View File
@@ -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:
+42
View File
@@ -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)
@@ -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__":
@@ -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
)
+74
View File
@@ -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 <say>heading left</say>"]})
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")
+6 -6
View File
@@ -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"