diff --git a/src/lerobot/policies/molmoact2/molmoact2_adapter.py b/src/lerobot/policies/molmoact2/molmoact2_adapter.py deleted file mode 100644 index 4675ed078..000000000 --- a/src/lerobot/policies/molmoact2/molmoact2_adapter.py +++ /dev/null @@ -1,58 +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. - -"""MolmoAct2 adapter for the language-conditioned runtime. - -MolmoAct2 is a flat VLA: it conditions on a single natural-language ``task`` -string (``"The task is to {task}. ..."``) that its processor packs — together -with the images and discretized state — into model inputs (``input_ids`` / -``pixel_values`` / ...). It has no subtask/memory generation head, so the runtime -just predicts an action chunk from the already-packed observation. - -Run with ``--direct_subtask`` (robot) or ``--sim.direct_subtask`` (sim): what you -type becomes the ``task`` the processor packs, and the runtime does not attempt -subtask/memory generation. The observation provider re-packs on every frame with -the live task (see ``_build_robot_observation_provider`` / the dynamic task -getter in ``runtime.cli``), so typing a new command switches the instruction -immediately. -""" - -from __future__ import annotations - -from typing import Any - -from lerobot.runtime import RuntimeState -from lerobot.runtime.adapter import BaseLanguageAdapter - - -class MolmoAct2PolicyAdapter(BaseLanguageAdapter): - """Runtime bridge for flat MolmoAct2 policies (direct task-text conditioning).""" - - def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: - # The current task/subtask was packed into the model inputs (input_ids, - # pixel_values, ...) by the policy processor, fed the live task by the - # observation provider. ``predict_action_chunk`` resolves the action mode - # from the checkpoint config (``inference_action_mode`` must be set to - # "continuous" or "discrete"). - return self.policy.predict_action_chunk(observation) - - def generate_text( - self, - kind: str, - observation: dict[str, Any] | None, - state: RuntimeState, - user_text: str | None = None, - ) -> str: - # MolmoAct2 has no text-generation head; direct-subtask mode skips this. - return "" diff --git a/src/lerobot/policies/pi05/pi05_adapter.py b/src/lerobot/policies/pi05/pi05_adapter.py deleted file mode 100644 index 757a52f1b..000000000 --- a/src/lerobot/policies/pi05/pi05_adapter.py +++ /dev/null @@ -1,50 +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. - -"""PI05 adapter for the language-conditioned runtime. - -PI05 is a flat VLA: it conditions the action expert directly on the task text, -which its preprocessor tokenizes into ``observation.language.tokens``. It has no -subtask/memory generation head, so the runtime simply predicts an action chunk -from the already-tokenized observation. Text generation is unsupported — run -with ``--sim.direct_subtask`` so the runtime doesn't attempt subtask/memory -generation (what you type becomes the task the preprocessor tokenizes). -""" - -from __future__ import annotations - -from typing import Any - -from lerobot.runtime import RuntimeState -from lerobot.runtime.adapter import BaseLanguageAdapter - - -class PI05PolicyAdapter(BaseLanguageAdapter): - """Runtime bridge for flat PI05 policies (direct task-text conditioning).""" - - def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: - # The task text was tokenized into observation.language.* by the policy - # preprocessor (fed the current task by the observation provider), so we - # just predict the action chunk from it. - return self.policy.predict_action_chunk(observation) - - def generate_text( - self, - kind: str, - observation: dict[str, Any] | None, - state: RuntimeState, - user_text: str | None = None, - ) -> str: - # PI05 has no text-generation head; direct-subtask mode skips this path. - return "" diff --git a/src/lerobot/runtime/adapter.py b/src/lerobot/runtime/adapter.py index eb9c4e79e..56ef4c187 100644 --- a/src/lerobot/runtime/adapter.py +++ b/src/lerobot/runtime/adapter.py @@ -168,6 +168,27 @@ class BaseLanguageAdapter(ABC): return text +class DirectTaskPolicyAdapter(BaseLanguageAdapter): + """Adapter for flat policies conditioned directly on the operator's task text. + + Policies such as PI0.5 and MolmoAct2 do not expose a language-generation + head. Their preprocessors pack the current task into the model inputs, so + the runtime only needs to request an action chunk. + """ + + def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: + return self.policy.predict_action_chunk(observation) + + def generate_text( + self, + kind: str, + observation: dict[str, Any] | None, + state: RuntimeState, + user_text: str | None = None, + ) -> str: + return "" + + def looks_like_gibberish(text: str) -> bool: """Heuristic filter for malformed / collapsed LM-head output.""" if not text or not text.strip(): diff --git a/src/lerobot/runtime/cli.py b/src/lerobot/runtime/cli.py index 5997b6e49..6d58c0531 100644 --- a/src/lerobot/runtime/cli.py +++ b/src/lerobot/runtime/cli.py @@ -60,7 +60,6 @@ import argparse import logging import sys from collections.abc import Callable -from contextlib import suppress from typing import Any from .adapter import GenerationConfig @@ -71,6 +70,7 @@ logger = logging.getLogger("lerobot.runtime") def _parse_args(argv: list[str] | None = None, *, prog: str | None = None) -> argparse.Namespace: + raw_argv = list(sys.argv[1:] if argv is None else argv) p = argparse.ArgumentParser( prog=prog, description="Interactive REPL runtime for a language-conditioned robot policy.", @@ -196,48 +196,6 @@ def _parse_args(argv: list[str] | None = None, *, prog: str | None = None) -> ar "``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( "--rerun", action="store_true", @@ -264,16 +222,6 @@ def _parse_args(argv: list[str] | None = None, *, prog: str | None = None) -> ar help="Direct-subtask mode (sim OR robot): your typed text IS the subtask " "fed to the action expert; the LM subtask generator is disabled.", ) - 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)." - ), - ) # --- RoboCasa simulation mode args ------------------------------- # Setting ``--sim`` flips the runtime into simulation mode: instead of # a real robot it drives a single RoboCasa mujoco scene, feeding the @@ -462,7 +410,12 @@ def _parse_args(argv: list[str] | None = None, *, prog: str | None = None) -> ar 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) + args, unknown = p.parse_known_args(raw_argv) + unsupported = [arg for arg in unknown if not arg.startswith(("--robot.", "--policy."))] + if unsupported: + p.error(f"unrecognized arguments: {' '.join(unsupported)}") + args.raw_argv = raw_argv + return args # Columns the runtime supplies itself via its own message stream — strip @@ -784,374 +737,122 @@ def _select_task_interactively( ) -def _dataset_features_from_robot(robot) -> dict[str, Any]: - """Build a LeRobot feature schema from a connected robot. +def _build_language_rollout_context(args: argparse.Namespace) -> Any: + """Build the canonical rollout context for a language-controlled robot.""" + import threading # noqa: PLC0415 - Used when no ``--dataset.repo_id`` is given so the runtime can assemble - observations and name action joints without a dataset (normalization stats - then come from the checkpoint). Mirrors ``lerobot-rollout``'s - ``build_rollout_context``: only ``.pos`` joints and camera features are - routed to the policy. - """ - from lerobot.utils.feature_utils import ( # noqa: PLC0415 - combine_feature_dicts, - hw_to_dataset_features, + import draccus # noqa: PLC0415 + + from lerobot.configs import parser # noqa: PLC0415 + from lerobot.rollout import RolloutConfig, build_rollout_context # noqa: PLC0415 + + # Importing the rollout entry point registers every bundled camera and + # robot config choice used by Draccus. Third-party choices were registered + # by the top-level entry point before reaching this function. + from lerobot.scripts import lerobot_rollout as _rollout_registrations # noqa: F401, PLC0415 + + rollout_argv = [arg for arg in args.raw_argv if arg.startswith(("--policy.", "--robot."))] + if args.task: + rollout_argv.append(f"--task={args.task}") + rollout_argv.extend( + ( + "--strategy.type=base", + f"--fps={args.ctrl_hz}", + "--return_to_initial_position=false", + ) ) - obs_hw = { - key: ft - for key, ft in robot.observation_features.items() - if isinstance(ft, tuple) or (ft is float and key.endswith(".pos")) - } - action_hw = {key: ft for key, ft in robot.action_features.items() if key.endswith(".pos")} - obs_features = hw_to_dataset_features(obs_hw, "observation") - action_features = hw_to_dataset_features(action_hw, "action") - return combine_feature_dicts(obs_features, action_features) - - -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) - + previous_argv = sys.argv 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 + # RolloutConfig resolves --policy.path and policy overrides through the + # shared parser helpers, which intentionally read sys.argv. + sys.argv = [previous_argv[0], *rollout_argv] + parsed_argv = parser.filter_path_args(RolloutConfig.__get_path_fields__(), rollout_argv) + cfg = draccus.parse(config_class=RolloutConfig, args=parsed_argv) + finally: + sys.argv = previous_argv - 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) + if getattr(cfg.policy, "compile_model", False): + cfg.policy.compile_model = False + if getattr(cfg.policy, "gradient_checkpointing", False): + cfg.policy.gradient_checkpointing = False + if args.fp8: + if hasattr(cfg.policy, "use_flashrt_fp8_mlp"): + cfg.policy.use_flashrt_fp8_mlp = True + else: + logger.warning("--fp8 ignored: %s does not support it", cfg.policy.type) - 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 + return build_rollout_context(cfg, threading.Event()) -def _build_robot_observation_provider( +def _build_rollout_runtime_io( + ctx: Any, *, - robot, - preprocessor: Any, - device: str, - task: str | None, - ds_features: dict[str, Any] | None, - rerun_log: bool = False, - get_task: Callable[[], str | None] | None = None, -) -> Callable[[], dict | None]: - """Closure reading from the robot each call: ``robot.get_observation()`` → - ``build_inference_frame`` (state vector + image tensors, batched, on device) - → ``EnvTransition``-wrapped preprocessor (rename, normalise) → flat - observation batch for ``select_action`` / ``select_message``. - - ``get_task`` (optional) is read every frame so the instruction packed into - the observation tracks the live task/subtask (e.g. MolmoAct2, whose processor - tokenizes the task into ``input_ids`` each frame). Falls back to the static - ``task`` when it returns nothing. - """ + rerun_log: bool, + get_task: Callable[[], str | None], +) -> tuple[Callable[[], dict | None], Callable[[Any], None]]: + """Adapt a rollout context to the language runtime's observation/action API.""" import torch # noqa: PLC0415 from lerobot.policies.utils import ( # noqa: PLC0415 - build_inference_frame, + make_robot_action, prepare_observation_for_inference, ) + from lerobot.utils.feature_utils import build_dataset_frame # noqa: PLC0415 - 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) - - # Camera-key → training (H, W) map from ``ds_features``. Live cameras - # rarely match the recorded resolution, and a different aspect ratio - # changes resize_with_pad's padding geometry — the flow head tolerates - # that, but the tightly-supervised LM head goes OOD and collapses. - _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)) + robot = ctx.hardware.robot_wrapper + device = torch.device(ctx.runtime.cfg.device or "cpu") + latest_raw: dict[str, Any] = {} def _provider() -> dict | None: - # Live task: re-read every frame so a typed command re-packs the prompt - # (falls back to the static startup task). - cur_task = (get_task() if get_task is not None else None) or task try: raw = robot.get_observation() + latest_raw.clear() + latest_raw.update(raw) + if rerun_log: + from lerobot.runtime import rerun_viz # noqa: PLC0415 + + camera_keys = list(robot.cameras) + state = {k: v for k, v in raw.items() if isinstance(v, (int, float))} + rerun_viz.log_robot_frame(raw, camera_keys, state=state, task=get_task()) + _strip_runtime_owned_language_cols(raw) + processed = ctx.processors.robot_observation_processor(raw) + observation = build_dataset_frame(ctx.data.dataset_features, processed, prefix="observation") + observation = prepare_observation_for_inference( + observation, + device, + task=get_task(), + robot_type=robot.robot_type, + ) + observation = ctx.policy.preprocessor(observation) + return _select_observation_to_device(observation, device) except Exception as exc: # noqa: BLE001 - logger.warning("robot.get_observation failed: %s", exc) + logger.warning("robot observation pipeline failed: %s", exc) return None - # Live camera view: log the raw frames + joint state to rerun before any - # resize (natural camera resolution). Best-effort — never blocks control. - if rerun_log: - from lerobot.runtime import rerun_viz # noqa: PLC0415 - - cam_keys = list(target_image_shapes.keys()) or [ - k for k, v in raw.items() if hasattr(v, "ndim") and getattr(v, "ndim", 0) == 3 - ] - state = {k: v for k, v in raw.items() if isinstance(v, (int, float)) and k not in cam_keys} - rerun_viz.log_robot_frame(raw, cam_keys, state=state, task=cur_task) - - # The runtime supplies messages itself; strip any language - # columns the robot stream may carry through. - _strip_runtime_owned_language_cols(raw) - - # Resize live frames to the training (H, W) so the downstream - # resize_with_pad geometry matches what the model saw in training. - 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 - # One-shot state-vector print so the operator can eyeball it - # against dataset stats (state OOD is a real VLA failure mode). - 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=cur_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=cur_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_features: dict[str, Any], - rerun_log: bool = False, -) -> 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_features) - elif isinstance(action, dict): - action_dict = action + processed_action = ctx.policy.postprocessor(action) + if isinstance(processed_action, torch.Tensor): + if processed_action.ndim == 1: + processed_action = processed_action.unsqueeze(0) + action_dict = make_robot_action(processed_action, ctx.data.dataset_features) + elif isinstance(processed_action, dict): + action_dict = processed_action else: - logger.warning("unsupported action type %r — skipping", type(action)) + logger.warning("unsupported action type %r — skipping", type(processed_action)) return - robot.send_action(action_dict) - # Smooth live view: log the cameras every control tick (buffered - # async_read is cheap). Best-effort — never blocks control. + raw = latest_raw or robot.get_observation() + robot_action = ctx.processors.robot_action_processor((action_dict, raw)) + robot.send_action(robot_action) if rerun_log: from lerobot.runtime import rerun_viz # noqa: PLC0415 rerun_viz.log_cameras(robot) except Exception as exc: # noqa: BLE001 - logger.error("robot.send_action failed: %s", exc, exc_info=True) + logger.error("robot action pipeline failed: %s", exc, exc_info=True) - return _executor + return _provider, _executor def _print_runtime_help() -> None: @@ -1256,168 +957,6 @@ def _handle_slash_command(runtime: Any, line: str) -> bool: return False -def _run_autonomous( - runtime: Any, - *, - robot, - auto_start: bool, - initial_task: str | None, - max_ticks: int | None, - panel_label: str = "Runtime", -) -> int: - """Drive the runtime continuously at ``ctrl_hz`` while accepting - stdin events in the foreground. - - 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( - "[runtime] Robot connected — starting in ACTION mode. Press ENTER to begin, Ctrl+C to abort. " - ) - except (EOFError, KeyboardInterrupt): - print("\n[runtime] 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="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", panel_label=panel_label, 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[runtime] 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="runtime-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( - "[runtime] no task yet — use /action to start", - flush=True, - ) - except KeyboardInterrupt: - print("\n[runtime] 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("[runtime] robot disconnected", flush=True) - except Exception as exc: # noqa: BLE001 - print(f"[runtime] WARNING: robot.disconnect raised {exc}", flush=True) - - return 0 - - def _make_state_panel_renderer( runtime: Any, *, @@ -1427,10 +966,7 @@ def _make_state_panel_renderer( ) -> 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. + Used by ``_run_repl`` for dataset-driven dry runs. """ from rich.console import Console # noqa: PLC0415 @@ -1621,17 +1157,27 @@ def run( render_size=args.sim_render_size, ) - print(f"[runtime] loading policy from {args.policy_path}", flush=True) - # Sim mode always loads processors from the checkpoint; robot mode does too - # when no dataset is supplied (stats come from the checkpoint / norm_tag). - load_processors_from_checkpoint = sim_mode or (autonomous_mode and not args.dataset_repo_id) - policy, preprocessor, postprocessor, ds_meta = _load_policy_and_preprocessor( - args.policy_path, - args.dataset_repo_id, - load_processors_from_checkpoint=load_processors_from_checkpoint, - fp8=args.fp8, - device=args.policy_device, - ) + rollout_ctx = None + if autonomous_mode: + print("[runtime] building rollout context (policy, processors, robot)", flush=True) + rollout_ctx = _build_language_rollout_context(args) + policy = rollout_ctx.policy.policy + preprocessor = rollout_ctx.policy.preprocessor + postprocessor = rollout_ctx.policy.postprocessor + ds_meta = None + if args.dataset_repo_id is not None: + from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata # noqa: PLC0415 + + ds_meta = LeRobotDatasetMetadata(args.dataset_repo_id) + else: + print(f"[runtime] loading policy from {args.policy_path}", flush=True) + policy, preprocessor, postprocessor, ds_meta = _load_policy_and_preprocessor( + args.policy_path, + args.dataset_repo_id, + load_processors_from_checkpoint=sim_mode, + fp8=args.fp8, + device=args.policy_device, + ) policy_type = getattr(policy.config, "type", None) if adapter_factory is None: @@ -1658,8 +1204,7 @@ def run( # 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. + # ``_run_repl`` still handles the no-default case. if not args.task: chosen = _select_task_interactively( ds_meta=ds_meta, @@ -1706,7 +1251,7 @@ def run( ) observation_provider = sim_backend.observation_provider robot_executor = sim_backend.action_executor - robot = sim_backend # reuse _run_autonomous cleanup (calls .disconnect()) + robot = sim_backend # Point the already-running live viewer at the backend and hand it the # server so disconnect() shuts it down cleanly. sim_holder["backend"] = sim_backend @@ -1721,35 +1266,13 @@ def run( grpc_port=args.rerun_grpc_port, web_port=args.rerun_web_port, ) - print( - f"[runtime] 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, - ) - # Feature schema: from the dataset when given, otherwise derived from the - # connected robot (mirrors lerobot-rollout) so no dataset is required. - robot_features = ds_meta.features if ds_meta is not None else _dataset_features_from_robot(robot) - observation_provider = _build_robot_observation_provider( - robot=robot, - preprocessor=preprocessor, - device=str(getattr(policy.config, "device", "cpu")), - task=args.task, - ds_features=robot_features, + robot = rollout_ctx.hardware.robot_wrapper.inner + print(f"[runtime] connected to {robot.name}", flush=True) + observation_provider, robot_executor = _build_rollout_runtime_io( + rollout_ctx, rerun_log=bool(args.rerun), get_task=_live_task, ) - robot_executor = _build_robot_action_executor( - robot=robot, - postprocessor=postprocessor, - ds_features=robot_features, - rerun_log=bool(args.rerun), - ) elif args.dataset_repo_id is not None: print( f"[runtime] streaming observations from {args.dataset_repo_id} " @@ -1865,9 +1388,8 @@ def _run_sim_interactive( ) -> int: """Main-thread control loop for the RoboCasa sim backend. - Unlike ``_run_autonomous`` (which runs ``runtime.run()`` in a daemon - thread), the tick loop — and therefore MuJoCo's EGL rendering — runs in the - MAIN thread. Driving the sim render from a background thread intermittently + The tick loop — and therefore MuJoCo's EGL rendering — runs in the MAIN + thread. Driving the sim render from a background thread intermittently corrupts the offscreen GL context (dark/garbled frames); main-thread stepping matches ``lerobot-eval`` and renders cleanly. Stdin is polled non-blockingly so typed commands still work while the sim runs. diff --git a/src/lerobot/runtime/registry.py b/src/lerobot/runtime/registry.py index 088ff4f56..d297cf656 100644 --- a/src/lerobot/runtime/registry.py +++ b/src/lerobot/runtime/registry.py @@ -27,8 +27,8 @@ from typing import Any _ADAPTERS: dict[str, str] = { "pi052": "lerobot.policies.pi052.inference.pi052_adapter:PI052PolicyAdapter", - "pi05": "lerobot.policies.pi05.pi05_adapter:PI05PolicyAdapter", - "molmoact2": "lerobot.policies.molmoact2.molmoact2_adapter:MolmoAct2PolicyAdapter", + "pi05": "lerobot.runtime.adapter:DirectTaskPolicyAdapter", + "molmoact2": "lerobot.runtime.adapter:DirectTaskPolicyAdapter", } diff --git a/src/lerobot/runtime/repl.py b/src/lerobot/runtime/repl.py index c7c07fead..8831350db 100644 --- a/src/lerobot/runtime/repl.py +++ b/src/lerobot/runtime/repl.py @@ -11,94 +11,13 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Stdin REPL event collector for the language-conditioned runtime. - -Reads non-blocking stdin lines, classifies each one heuristically: - - "stop" / "quit" / "exit" → state["stop"] = True - "/action" / "/pause" → set state["mode"] - starts with "task:" or first line → set runtime task - anything else → user_interjection event - -Plugged into the runtime via ``event_collector=StdinReader().poll``. - -Note: the shipped CLI drives stdin directly in its REPL / autonomous -loops and does *not* wire this collector; it's kept as the documented -embedding hook and for tests. -""" +"""Small event helper shared by the language-runtime command loops.""" from __future__ import annotations -import select -import sys -from dataclasses import dataclass, field from typing import Any -@dataclass -class StdinReader: - """Non-blocking stdin line collector for the runtime loop.""" - - prompt: str = "> " - _seen_first_line: bool = field(default=False, init=False) - _prompted: bool = field(default=False, init=False) - - def poll(self, state: dict[str, Any]) -> None: - """Drain pending stdin lines into runtime events.""" - # Print the input prompt once on every fresh tick if we don't - # already have a pending line; matches the expected REPL feel. - if not self._prompted: - print(self.prompt, end="", flush=True) - self._prompted = True - - # ``select`` with timeout=0 makes this non-blocking. Only works - # for actual TTY / pipe stdins; CI / scripted runs hit EOF. - try: - ready, _, _ = select.select([sys.stdin], [], [], 0) - except (ValueError, OSError): - return - if not ready: - return - - line = sys.stdin.readline() - if not line: # EOF - state["stop"] = True - return - line = line.strip() - self._prompted = False # we'll re-prompt next tick - if not line: - return - - lower = line.lower() - if lower in {"stop", "quit", "exit"}: - state["stop"] = True - return - - # Slash commands flip the run mode. ``/pause`` stops the action - # loop (the action steps gate on ``state["mode"]``); ``/action`` - # resumes it. - if lower.split(" ", 1)[0] in {"/action", "/act", "/run"}: - state["mode"] = "action" - return - if lower in {"/pause", "/p"}: - state["mode"] = "paused" - queue = state.get("action_queue") - if hasattr(queue, "clear"): - queue.clear() - return - - # First non-control line sets the task if no task is active. - if not state.get("task"): - task = line[5:].strip() if lower.startswith("task:") else line - state["task"] = task - print(f"[runtime] Task: {task}", flush=True) - self._seen_first_line = True - return - - state["recent_interjection"] = line - _emit(state, "user_interjection") - - def _emit(state: Any, event_name: str) -> None: if hasattr(state, "emit"): state.emit(event_name) diff --git a/src/lerobot/runtime/rerun_viz.py b/src/lerobot/runtime/rerun_viz.py index ae3098d55..1da696fcf 100644 --- a/src/lerobot/runtime/rerun_viz.py +++ b/src/lerobot/runtime/rerun_viz.py @@ -23,6 +23,7 @@ never interrupts robot control. from __future__ import annotations import logging +from contextlib import suppress from typing import Any logger = logging.getLogger(__name__) @@ -68,11 +69,9 @@ def log_cameras(robot: Any) -> None: cams = getattr(robot, "cameras", None) or {} for name, cam in cams.items(): - try: + with suppress(Exception): frame = cam.async_read(timeout_ms=1) rr.log(f"cameras/{name}", rr.Image(frame)) - except Exception: # noqa: BLE001 - pass except Exception as exc: # noqa: BLE001 logger.debug("[runtime] rerun camera log failed: %s", exc) @@ -97,10 +96,8 @@ def log_robot_frame( rr.log(f"cameras/{cam}", rr.Image(img)) if state: for name, val in state.items(): - try: + with suppress(Exception): rr.log(f"state/{name}", rr.Scalars(float(val))) - except Exception: # noqa: BLE001 - pass if task: rr.log("prompt/task", rr.TextLog(task)) if subtask: diff --git a/src/lerobot/runtime/sim_robocasa.py b/src/lerobot/runtime/sim_robocasa.py index e8278288c..870bcd8e3 100644 --- a/src/lerobot/runtime/sim_robocasa.py +++ b/src/lerobot/runtime/sim_robocasa.py @@ -22,6 +22,8 @@ from typing import Any import numpy as np import torch +from lerobot.utils.video_annotation import annotate_frame + logger = logging.getLogger(__name__) @@ -47,54 +49,6 @@ def _label_panel(img: np.ndarray, label: str) -> np.ndarray: return img -def _overlay_text(frame: np.ndarray, task: str | None, subtask: str | None, memory: str | None) -> np.ndarray: - """Draw task / subtask / memory lines onto an (H, W, 3) uint8 frame. - - Best-effort: returns the frame unchanged if OpenCV is unavailable. - """ - try: - import cv2 # noqa: PLC0415 - except ImportError: - return frame - - lines = [ - f"{label}: {val}" for label, val in (("Task", task), ("Subtask", subtask), ("Memory", memory)) if val - ] - if not lines: - return frame - - img = np.ascontiguousarray(frame).copy() - font, scale, margin = cv2.FONT_HERSHEY_SIMPLEX, 0.5, 6 - max_width = img.shape[1] - 2 * margin - wrapped_lines: list[str] = [] - for text in lines: - # naive width-based wrap so long memory strings stay on-frame - words, cur = text.split(), "" - for w in words: - cand = f"{cur} {w}".strip() - if cv2.getTextSize(cand, font, scale, 1)[0][0] > max_width and cur: - wrapped_lines.append(cur) - cur = w - else: - cur = cand - wrapped_lines.append(cur) - - # Use a dark translucent header for contrast, then draw each label once. - # The previous black-outline + white-foreground technique rendered every - # glyph twice and looked like offset duplicate text in the MJPEG stream. - line_height = 20 - header_height = min(img.shape[0], len(wrapped_lines) * line_height + 6) - backdrop = img.copy() - cv2.rectangle(backdrop, (0, 0), (img.shape[1], header_height), (0, 0, 0), -1) - cv2.addWeighted(backdrop, 0.55, img, 0.45, 0, dst=img) - - y = 18 - for line in wrapped_lines: - cv2.putText(img, line, (margin, y), font, scale, (255, 255, 255), 1, cv2.LINE_AA) - y += line_height - return img - - # RoboCasa's MuJoCo EGL offscreen renderer produces garbled/static frames when # only ONE worker env is running (reproducible with lerobot-eval --batch_size=1). # With >=2 workers the renderer is stable. We therefore run the interactive sim @@ -225,7 +179,7 @@ class RoboCasaSimBackend: Exposes ``observation_provider`` / ``action_executor`` closures matching the runtime's injected-callable contract, plus ``disconnect`` so the shared - ``_run_autonomous`` cleanup path can close the env (and flush the video). + The runtime cleanup path closes the env and flushes the video. The env must be created via :func:`create_sim_env` *before* the policy touches CUDA (see that function's note on the EGL/CUDA fork hazard). @@ -366,25 +320,6 @@ class RoboCasaSimBackend: except Exception as exc: # noqa: BLE001 logger.error("[sim] env.step failed: %s", exc, exc_info=True) - def _frontal_obs_image(self) -> np.ndarray | None: - """Return the current front agent-view camera image (H, W, 3) uint8. - - Uses the observation the policy already consumes rather than a separate - ``env.render()`` call: the render path's camera is intermittently - corrupted by the offscreen EGL context, whereas the policy's obs images - come straight through the eval pipeline and stay clean. - """ - pixels = (self._last_obs or {}).get("pixels") - if not isinstance(pixels, dict) or not pixels: - return None - cam = "robot0_agentview_left" if "robot0_agentview_left" in pixels else next(iter(pixels)) - img = np.asarray(pixels[cam]) - if img.ndim == 4: # vec env batches to (1, H, W, C) - img = img[0] - if img.ndim != 3 or img.shape[-1] != 3: - return None - return img.astype(np.uint8) - def _multiview_frame(self) -> np.ndarray | None: """Composite the configured camera views (env 0) side by side, labeled. @@ -424,7 +359,10 @@ class RoboCasaSimBackend: return subtask = self._subtask_getter() if self._subtask_getter else None memory = self._memory_getter() if self._memory_getter else None - annotated = _overlay_text(frame, self._current_task(), subtask, memory) + annotated = annotate_frame( + frame, + (("Task", self._current_task()), ("Subtask", subtask), ("Memory", memory)), + ) self._frames.append(annotated) self._latest_frame = annotated # served by the live MJPEG stream self._write_live_frame(annotated) @@ -476,7 +414,7 @@ class RoboCasaSimBackend: self._stream_server = server def disconnect(self) -> None: - """Match the robot backend's cleanup contract (called by _run_autonomous).""" + """Match the robot backend's cleanup contract.""" if self._stream_server is not None: try: self._stream_server.shutdown() diff --git a/src/lerobot/runtime/ui.py b/src/lerobot/runtime/ui.py deleted file mode 100644 index 2059b774c..000000000 --- a/src/lerobot/runtime/ui.py +++ /dev/null @@ -1,112 +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. -"""Rich-based REPL layout for the language-conditioned runtime. - -Two-zone terminal layout: - - [chat scrollback — user messages / robot responses, scrolls naturally] - - ┌── State ──────────────────────────────────────────┐ - │ task please clean up the kitchen │ - │ subtask grasp the handle of the sponge │ - │ plan 1. grasp sponge 2. wipe 3. tidy │ - │ memory sponge picked up; counter still dirty │ - └───────────────────────────────────────────────────┘ - > _ - -Chat lines print above a ``rich.Live`` region (natural scrollback); the -state panel re-renders on change, auto-suspending while input is pending. -""" - -from __future__ import annotations - -from typing import Any - -try: # rich is optional; only required for the interactive REPL. - from rich.console import Console - from rich.panel import Panel - from rich.table import Table - from rich.text import Text - - _HAS_RICH = True -except ImportError: # pragma: no cover - _HAS_RICH = False - Console = Any # type: ignore[assignment] - Panel = Any # type: ignore[assignment] - Table = Any # type: ignore[assignment] - Text = Any # type: ignore[assignment] - - -_STATE_KEYS = ( - ("task", "task"), - ("current_subtask", "subtask"), - ("current_plan", "plan"), - ("current_memory", "memory"), -) - - -def make_state_panel(state: dict[str, Any], *, title: str = "Runtime state") -> Any: - """Render the persistent state panel for the live region. - - Returns a :class:`rich.panel.Panel`. Caller passes it to - ``Live.update(panel)`` whenever the state changes. - """ - if not _HAS_RICH: - raise RuntimeError( - "rich is required for the interactive REPL. " - "`pip install rich` (it's a transitive dep of lerobot)." - ) - table = Table.grid(padding=(0, 2), expand=True) - table.add_column(justify="right", style="dim", no_wrap=True, width=10) - table.add_column(justify="left") - for key, label in _STATE_KEYS: - value = state.get(key) - 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 - footer = Text.assemble(("queued actions: ", "dim"), (str(queue_len), "bold cyan")) - table.add_row("", footer) - run_mode = state.get("mode", "action") - mode_tag = "[green]action[/]" if run_mode == "action" else "[yellow]paused[/]" - return Panel( - table, - title=f"[bold]{title}[/] · mode: {mode_tag}", - border_style="cyan", - ) - - -def print_user_line(console: Any, line: str) -> None: - """Append a user-typed line to the chat scrollback.""" - if not _HAS_RICH: - print(f"you: {line}", flush=True) - return - console.print(f"[bold cyan]you:[/] {line}") - - -def print_robot_lines(console: Any, lines: list[str]) -> None: - """Append robot/runtime log lines to the chat scrollback.""" - if not _HAS_RICH: - for line in lines: - print(f"robot: {line.lstrip()}", flush=True) - return - for line in lines: - # The runtime uses leading whitespace + "label: text"; render - # the label in green and the value in default for readability. - stripped = line.lstrip() - if ":" in stripped: - label, _, value = stripped.partition(":") - console.print(f"[bold green]robot[/] [dim]({label.strip()})[/] {value.strip()}") - else: - console.print(f"[bold green]robot:[/] {stripped}") diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index 0d6e998f3..da680b586 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -94,24 +94,7 @@ from lerobot.utils.utils import ( init_logging, inside_slurm, ) - - -def _wrap_text_to_width(text: str, cv2, font, scale: int, thickness: int, max_width: int) -> list[str]: - """Greedy word-wrap using measured pixel width so text fits the frame.""" - words = text.split() - lines: list[str] = [] - current = "" - for word in words: - candidate = f"{current} {word}".strip() - (w, _), _ = cv2.getTextSize(candidate, font, scale, thickness) - if w > max_width and current: - lines.append(current) - current = word - else: - current = candidate - if current: - lines.append(current) - return lines or [""] +from lerobot.utils.video_annotation import annotate_frame def _annotate_eval_frames(frames: np.ndarray, task: str | None, subtask: str | None) -> np.ndarray: @@ -123,36 +106,7 @@ def _annotate_eval_frames(frames: np.ndarray, task: str | None, subtask: str | N """ if frames.ndim != 4 or frames.shape[-1] != 3: return frames - try: - import cv2 # noqa: PLC0415 - except ImportError: - return frames - - width = frames.shape[2] - font = cv2.FONT_HERSHEY_SIMPLEX - scale = 0.5 - margin = 6 - max_width = width - 2 * margin - - lines: list[str] = [] - if task: - lines += _wrap_text_to_width(f"Task: {task}", cv2, font, scale, 1, max_width) - if subtask: - lines += _wrap_text_to_width(f"Subtask: {subtask}", cv2, font, scale, 1, max_width) - if not lines: - return frames - - out = frames.copy() - for i in range(out.shape[0]): - img = np.ascontiguousarray(out[i]) - y = 18 - for line in lines: - # Black outline then white fill so text stays legible on any scene. - cv2.putText(img, line, (margin, y), font, scale, (0, 0, 0), 3, cv2.LINE_AA) - cv2.putText(img, line, (margin, y), font, scale, (255, 255, 255), 1, cv2.LINE_AA) - y += 20 - out[i] = img - return out + return np.stack([annotate_frame(frame, (("Task", task), ("Subtask", subtask))) for frame in frames]) def _env_features_to_dataset_features(env_features: dict) -> dict: @@ -557,10 +511,7 @@ def eval_policy( subtask_scalar = getattr(policy, "last_subtask", None) annotated = [] for i in range(frames.shape[0]): - if subtasks is not None and i < len(subtasks): - subtask_i = subtasks[i] - else: - subtask_i = subtask_scalar + subtask_i = subtasks[i] if subtasks is not None and i < len(subtasks) else subtask_scalar annotated.append( _annotate_eval_frames( frames[i : i + 1], diff --git a/src/lerobot/utils/video_annotation.py b/src/lerobot/utils/video_annotation.py new file mode 100644 index 000000000..8fbc27ba5 --- /dev/null +++ b/src/lerobot/utils/video_annotation.py @@ -0,0 +1,71 @@ +# 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. + +"""Best-effort text overlays shared by evaluation and interactive rollouts.""" + +from __future__ import annotations + +from collections.abc import Iterable + +import numpy as np + + +def annotate_frame(frame: np.ndarray, fields: Iterable[tuple[str, str | None]]) -> np.ndarray: + """Return an RGB frame annotated with the non-empty labeled ``fields``.""" + if frame.ndim != 3 or frame.shape[-1] != 3: + return frame + try: + import cv2 # noqa: PLC0415 + except ImportError: + return frame + + text_rows = [f"{label}: {value}" for label, value in fields if value] + if not text_rows: + return frame + + image = np.ascontiguousarray(frame).copy() + font, scale, thickness, margin = cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1, 6 + max_width = image.shape[1] - 2 * margin + lines: list[str] = [] + for text in text_rows: + current = "" + for word in text.split(): + candidate = f"{current} {word}".strip() + width = cv2.getTextSize(candidate, font, scale, thickness)[0][0] + if width > max_width and current: + lines.append(current) + current = word + else: + current = candidate + if current: + lines.append(current) + + line_height = 20 + header_height = min(image.shape[0], len(lines) * line_height + 6) + backdrop = image.copy() + cv2.rectangle(backdrop, (0, 0), (image.shape[1], header_height), (0, 0, 0), -1) + cv2.addWeighted(backdrop, 0.55, image, 0.45, 0, dst=image) + + for index, line in enumerate(lines): + cv2.putText( + image, + line, + (margin, 18 + index * line_height), + font, + scale, + (255, 255, 255), + thickness, + cv2.LINE_AA, + ) + return image diff --git a/tests/runtime/test_adapter.py b/tests/runtime/test_adapter.py index 78f19de68..270fbf396 100644 --- a/tests/runtime/test_adapter.py +++ b/tests/runtime/test_adapter.py @@ -1,5 +1,10 @@ from lerobot.runtime import RuntimeState -from lerobot.runtime.adapter import BaseLanguageAdapter, GenerationConfig, looks_like_gibberish +from lerobot.runtime.adapter import ( + BaseLanguageAdapter, + DirectTaskPolicyAdapter, + GenerationConfig, + looks_like_gibberish, +) class ScriptedAdapter(BaseLanguageAdapter): @@ -72,3 +77,22 @@ def test_looks_like_gibberish_basic(): assert looks_like_gibberish("") assert looks_like_gibberish(":::: ::") assert not looks_like_gibberish("pick up the red cube") + + +def test_direct_task_adapter_delegates_action_chunk(): + class Policy: + def predict_action_chunk(self, observation): + return ("chunk", observation) + + observation = {"task": "pick up the cube"} + adapter = DirectTaskPolicyAdapter(Policy()) + + assert adapter.select_action(observation, RuntimeState()) == ("chunk", observation) + assert adapter.generate_text("subtask", observation, RuntimeState()) == "" + + +def test_flat_policy_registry_reuses_direct_task_adapter(): + from lerobot.runtime.registry import get_language_adapter_factory + + assert get_language_adapter_factory("pi05") is DirectTaskPolicyAdapter + assert get_language_adapter_factory("molmoact2") is DirectTaskPolicyAdapter diff --git a/tests/runtime/test_cli.py b/tests/runtime/test_cli.py new file mode 100644 index 000000000..16f5a867c --- /dev/null +++ b/tests/runtime/test_cli.py @@ -0,0 +1,55 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +import torch + +from lerobot.runtime.cli import _build_rollout_runtime_io, _parse_args + + +def test_parse_args_preserves_rollout_robot_overrides(): + args = _parse_args( + [ + "--policy.path=checkpoint", + "--robot.type=so101_follower", + "--robot.calibration_dir=/tmp/calibration", + ] + ) + + assert args.robot_type == "so101_follower" + assert "--robot.calibration_dir=/tmp/calibration" in args.raw_argv + + +def test_rollout_runtime_io_uses_context_processors(): + robot = MagicMock() + robot.robot_type = "mock_robot" + robot.cameras = {} + robot.get_observation.return_value = {"joint.pos": 1.5} + ctx = SimpleNamespace( + hardware=SimpleNamespace(robot_wrapper=robot), + runtime=SimpleNamespace(cfg=SimpleNamespace(device="cpu")), + processors=SimpleNamespace( + robot_observation_processor=lambda observation: observation, + robot_action_processor=lambda pair: pair[0], + ), + policy=SimpleNamespace( + preprocessor=lambda observation: observation, + postprocessor=lambda action: action, + ), + data=SimpleNamespace( + dataset_features={ + "observation.state": { + "dtype": "float32", + "shape": (1,), + "names": ["joint.pos"], + }, + "action": {"dtype": "float32", "shape": (1,), "names": ["joint.pos"]}, + } + ), + ) + provider, executor = _build_rollout_runtime_io(ctx, rerun_log=False, get_task=lambda: "move") + + observation = provider() + executor(torch.tensor([[2.0]])) + + assert observation["observation.state"].shape == (1, 1) + robot.send_action.assert_called_once_with({"joint.pos": 2.0}) diff --git a/tests/runtime/test_sim_robocasa.py b/tests/runtime/test_sim_robocasa.py index 3e60cb2bd..07ad6375c 100644 --- a/tests/runtime/test_sim_robocasa.py +++ b/tests/runtime/test_sim_robocasa.py @@ -3,7 +3,7 @@ from types import SimpleNamespace import numpy as np -from lerobot.runtime.sim_robocasa import _overlay_text +from lerobot.utils.video_annotation import annotate_frame def test_overlay_draws_each_label_once(monkeypatch): @@ -33,7 +33,10 @@ def test_overlay_draws_each_label_once(monkeypatch): monkeypatch.setitem(sys.modules, "cv2", fake_cv2) frame = np.full((120, 480, 3), 200, dtype=np.uint8) - annotated = _overlay_text(frame, "close the fridge", "reach for the handle", None) + annotated = annotate_frame( + frame, + (("Task", "close the fridge"), ("Subtask", "reach for the handle"), ("Memory", None)), + ) assert [call[0] for call in put_text_calls] == [ "Task: close the fridge",