mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-25 02:36:11 +00:00
refactor(runtime): reuse rollout context and remove dead code
This commit is contained in:
Pepijn
@@ -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 ""
|
||||
@@ -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 ""
|
||||
@@ -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():
|
||||
|
||||
+121
-599
@@ -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_<name>.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(<type>)`` (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 <task> to run · /pause to stop · "
|
||||
"/question <text> 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 <seconds>``): 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 <your task> 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.
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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}")
|
||||
@@ -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],
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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})
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user