From f2b90e3ad63714929ed7db2bd0ea15b12249672a Mon Sep 17 00:00:00 2001 From: Pepijn Date: Tue, 28 Jul 2026 10:05:27 +0200 Subject: [PATCH] feat(runtime): add interactive language rollouts --- docs/source/bring_your_own_policies.mdx | 156 +++ docs/source/inference.mdx | 48 +- src/lerobot/envs/configs.py | 10 +- src/lerobot/envs/robocasa.py | 44 +- src/lerobot/robots/utils.py | 4 +- src/lerobot/runtime/__init__.py | 38 + src/lerobot/runtime/adapter.py | 165 +++ src/lerobot/runtime/cli.py | 1191 ++++++++++++++++++++++ src/lerobot/runtime/language_runtime.py | 349 +++++++ src/lerobot/runtime/registry.py | 39 + src/lerobot/runtime/sim_robocasa.py | 406 ++++++++ src/lerobot/scripts/lerobot_eval.py | 42 +- src/lerobot/scripts/lerobot_rollout.py | 66 +- src/lerobot/utils/io_utils.py | 74 +- src/lerobot/utils/rerun_visualization.py | 10 +- src/lerobot/utils/video_annotation.py | 71 ++ tests/envs/test_robocasa_ordering.py | 48 + tests/runtime/test_adapter.py | 105 ++ tests/runtime/test_cli.py | 75 ++ tests/runtime/test_language_runtime.py | 100 ++ tests/runtime/test_sim_robocasa.py | 82 ++ 21 files changed, 3074 insertions(+), 49 deletions(-) create mode 100644 src/lerobot/runtime/__init__.py create mode 100644 src/lerobot/runtime/adapter.py create mode 100644 src/lerobot/runtime/cli.py create mode 100644 src/lerobot/runtime/language_runtime.py create mode 100644 src/lerobot/runtime/registry.py create mode 100644 src/lerobot/runtime/sim_robocasa.py create mode 100644 src/lerobot/utils/video_annotation.py create mode 100644 tests/envs/test_robocasa_ordering.py create mode 100644 tests/runtime/test_adapter.py create mode 100644 tests/runtime/test_cli.py create mode 100644 tests/runtime/test_language_runtime.py create mode 100644 tests/runtime/test_sim_robocasa.py diff --git a/docs/source/bring_your_own_policies.mdx b/docs/source/bring_your_own_policies.mdx index 697a07691..4b09186ed 100644 --- a/docs/source/bring_your_own_policies.mdx +++ b/docs/source/bring_your_own_policies.mdx @@ -191,6 +191,162 @@ def make_my_policy_pre_post_processors( --- +## Adding high- and low-level language control + +The policy API above is sufficient for training and standard evaluation. To use a language-conditioned policy with interactive `lerobot-rollout`, also register a runtime adapter. The adapter keeps policy-specific prompting and tokenization out of the generic control loop. + +The runtime supports two policy shapes: + +| Policy shape | Behavior | Adapter | +| ---------------- | ----------------------------------------------------------------------- | ---------------------------------------------- | +| Low-level / flat | The operator's task or subtask directly conditions action prediction. | Reuse `DirectTaskPolicyAdapter`. | +| High + low level | The policy generates subtasks or memory, then conditions actions on it. | Subclass `BaseLanguageAdapter`, as PI052 does. | + +During a rollout, `RuntimeState` stores the high-level task and the active language context: + +```text +task ──> adapter.generate_text("subtask", ...) ──> state.language_context["subtask"] + │ +observation ──> processors ──> adapter.select_action() ─┴─> action chunk ──> robot +``` + +The generic runtime handles generation frequency, pause/resume, prompt replacement, action queues, and dispatch. The adapter only translates between that runtime contract and your policy. + +### Low-level policies + +If your policy already consumes the live task through its normal preprocessor and implements `predict_action_chunk`, register the shared direct adapter. PI0.5 and MolmoAct2 use this path: + +```python +# src/lerobot/runtime/registry.py +_ADAPTERS = { + # ... + "my_policy": "lerobot.runtime.adapter:DirectTaskPolicyAdapter", +} +``` + +Run it with direct-subtask mode so the operator supplies the instruction used by the action policy: + +```bash +lerobot-rollout \ + --language \ + --policy.path=user/my_policy_checkpoint \ + --robot.type=so101_follower \ + --robot.port=/dev/ttyACM0 \ + --direct_subtask +``` + +The rollout context builds the observation batch with the current instruction before `DirectTaskPolicyAdapter` calls `policy.predict_action_chunk(observation)`. No text-generation method is required. + +### Hierarchical policies + +For a policy that generates language and actions, subclass [`BaseLanguageAdapter`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/runtime/adapter.py) and implement two methods: + +- `generate_text(kind, observation, state, user_text=None) -> str` generates a `subtask`, `memory`, or interjection response. +- `select_action(observation, state)` builds the low-level prompt from the active context and returns an action chunk. + +This abbreviated adapter follows [`PI052PolicyAdapter`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/inference/pi052_adapter.py): + +```python +# inference/my_policy_adapter.py +from typing import Any + +from lerobot.runtime import RuntimeState +from lerobot.runtime.adapter import BaseLanguageAdapter +from lerobot.utils.constants import ( + OBS_LANGUAGE_ATTENTION_MASK, + OBS_LANGUAGE_TOKENS, +) + + +class MyPolicyAdapter(BaseLanguageAdapter): + def select_action(self, observation: dict[str, Any], state: RuntimeState): + instruction = state.language_context.get("subtask") or state.task or "" + tokens, attention_mask = tokenize_instruction(instruction) + + batch = dict(observation) + batch[OBS_LANGUAGE_TOKENS] = tokens + batch[OBS_LANGUAGE_ATTENTION_MASK] = attention_mask + return self.policy.predict_action_chunk(batch) + + def generate_text( + self, + kind: str, + observation: dict[str, Any] | None, + state: RuntimeState, + user_text: str | None = None, + ) -> str: + messages = self.build_messages(kind, state, user_text) + batch, tokenizer = tokenize_messages(messages, observation) + return self.policy.select_message( + batch, + tokenizer=tokenizer, + min_new_tokens=self.gen.min_new_tokens, + temperature=self.gen.temperature, + top_p=self.gen.top_p, + ) + + def build_messages( + self, kind: str, state: RuntimeState, user_text: str | None + ) -> list[dict[str, str]]: + if kind == "subtask": + return [{"role": "user", "content": state.task or ""}] + if kind == "memory": + return [ + {"role": "user", "content": state.task or ""}, + { + "role": "user", + "content": f"Completed subtask: {state.extra.get('prior_subtask', '')}", + }, + ] + if kind == "interjection": + return [ + {"role": "user", "content": state.task or ""}, + {"role": "user", "content": user_text or ""}, + ] + raise ValueError(f"Unsupported text kind: {kind}") +``` + +`tokenize_instruction` and `tokenize_messages` are policy-specific helpers. They must reproduce the prompt format used during training; PI052, for example, adds the discretized robot state to its low-level subtask prompt and uses the same PaliGemma formatting for `select_message`. + +`BaseLanguageAdapter` provides the default hierarchy: regenerate a subtask at action-chunk boundaries, update memory when the subtask changes, and handle user interjections. Override `_regenerate_context` only if your policy uses a different hierarchy. + +Register the adapter with a lazy import so importing LeRobot does not load the model or its optional dependencies: + +```python +# src/lerobot/runtime/registry.py +_ADAPTERS = { + # ... + "my_policy": "lerobot.policies.my_policy.inference.my_policy_adapter:MyPolicyAdapter", +} +``` + +The key must match the policy's registered type. Once registered, the same checkpoint works through the shared entry point: + +```bash +lerobot-rollout \ + --language \ + --policy.path=user/my_hierarchical_checkpoint \ + --robot.type=so101_follower \ + --robot.port=/dev/ttyACM0 \ + --task="put the cup in the sink" +``` + +For RoboCasa-compatible policies, replace the robot arguments with `--sim --sim.task=`. Without `--direct_subtask`, the adapter generates the low-level subtask; with it, the operator bypasses high-level generation and supplies each subtask. + +### Keep training and deployment aligned + +The adapter is intentionally small, but its prompts are part of the model contract: + +- Use the same tokenizer, role formatting, special tokens, image ordering, and state encoding as training. +- Condition `select_action` on `state.language_context["subtask"]`, falling back to `state.task` for direct or not-yet-generated prompts. +- Return a full action chunk from `select_action`; the runtime handles control-rate dispatch. +- Keep optional model dependencies inside lazy imports. +- Test adapter selection, generated-message routing, action-batch construction, and direct-subtask behavior with a lightweight fake policy. + +PI052 is the complete in-tree reference: its [processor](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/processor_pi052.py) renders the training recipe, its [policy](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/modeling_pi052.py) exposes text and action generation, and its [adapter](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/inference/pi052_adapter.py) reconstructs those same prompts at deployment. + +--- + ## Path A: Out-of-tree plugin The fastest way to ship a policy: package it as a standalone Python distribution and install it alongside LeRobot. No PR required, you own the release cycle, and you can publish to PyPI under your own namespace. diff --git a/docs/source/inference.mdx b/docs/source/inference.mdx index 31405b5de..bdb47e978 100644 --- a/docs/source/inference.mdx +++ b/docs/source/inference.mdx @@ -1,6 +1,6 @@ # Policy Deployment (lerobot-rollout) -`lerobot-rollout` is the single CLI for deploying trained policies on real robots. It supports multiple execution strategies and inference backends, from quick evaluation to continuous recording and human-in-the-loop data collection. +`lerobot-rollout` is the single CLI for deploying trained policies on real robots or in an interactive simulator. It supports multiple execution strategies and inference backends, from quick evaluation to continuous recording, language-driven control, and human-in-the-loop data collection. ## Quick Start @@ -197,6 +197,52 @@ Teleop is optional — if omitted the robot holds its position during the reset --- +## Interactive language control + +Language-conditioned policies can expose a high-level text head in addition to +their action head. Add `--language` to open-prompt one of these policies on a +real robot. Language-only flags such as `--direct_subtask` select this mode +automatically. + +MolmoAct2 has no high-level planner, so use direct-subtask mode and type each +next low-level instruction yourself: + +```bash +lerobot-rollout \ + --policy.path=lerobot/MolmoAct2-SO100_101-LeRobot \ + --policy.device=cuda \ + --robot.type=so101_follower \ + --robot.port=/dev/ttyACM1 \ + --robot.cameras='{"cam0":{"type":"opencv","index_or_path":"/dev/video0","width":640,"height":480,"fps":30,"fourcc":"MJPG","backend":200},"cam1":{"type":"opencv","index_or_path":"/dev/video2","width":640,"height":480,"fps":30,"fourcc":"MJPG","backend":200}}' \ + --direct_subtask \ + --robot.max_relative_target='{"shoulder_pan":5,"shoulder_lift":5,"elbow_flex":5,"wrist_flex":5,"wrist_roll":5,"gripper":5}' +``` + +The robot starts paused. Type a subtask, then use `/resume` and `/pause` to +control action dispatch. Check the workspace and motion limits before resuming. +Without `--direct_subtask`, a policy such as PI052 generates its active subtask +from the high-level `--task` itself. + +RoboCasa uses the same runtime and processor path. `--sim` selects it +automatically, so no robot configuration is needed: + +```bash +MUJOCO_GL=egl lerobot-rollout \ + --policy.path=lerobot/pi052_robocasa \ + --sim --sim.task=CloseFridge --sim.split=pretrain \ + --task="close the fridge" \ + --disable_memory \ + --sim.render_size=384 \ + --sim.views=robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right \ + --mode=action --ctrl_hz=20 +``` + +Open `http://localhost:8010` for the live simulator view. Add +`--sim.direct_subtask` to bypass the language planner and make each typed prompt +the action policy's current subtask. + +--- + ## Inference Backends Select a backend with `--inference.type=`. All strategies work with both backends. diff --git a/src/lerobot/envs/configs.py b/src/lerobot/envs/configs.py index 3f6fd75f9..63d5e427f 100644 --- a/src/lerobot/envs/configs.py +++ b/src/lerobot/envs/configs.py @@ -560,7 +560,13 @@ class RoboCasaEnv(EnvConfig): kwargs["split"] = self.split return kwargs - def create_envs(self, n_envs: int, use_async_envs: bool = False): + def create_envs( + self, + n_envs: int, + use_async_envs: bool = False, + terminate_on_success: bool = True, + horizon: int | None = None, + ): from .robocasa import create_robocasa_envs if self.task is None: @@ -574,6 +580,8 @@ class RoboCasaEnv(EnvConfig): env_cls=env_cls, episode_length=self.episode_length, obj_registries=tuple(self.obj_registries), + terminate_on_success=terminate_on_success, + horizon=horizon, ) diff --git a/src/lerobot/envs/robocasa.py b/src/lerobot/envs/robocasa.py index a84a7c766..559ad2e01 100644 --- a/src/lerobot/envs/robocasa.py +++ b/src/lerobot/envs/robocasa.py @@ -33,8 +33,8 @@ logger = logging.getLogger(__name__) # Dimensions for the flat action/state vectors used by the LeRobot wrapper. # These correspond to the PandaOmron robot in RoboCasa365. -OBS_STATE_DIM = 16 # base_pos(3) + base_quat(4) + ee_pos_rel(3) + ee_quat_rel(4) + gripper_qpos(2) -ACTION_DIM = 12 # base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1) +OBS_STATE_DIM = 16 # ee_pos_rel(3) + ee_quat_rel(4) + base_pos(3) + base_quat(4) + gripper_qpos(2) +ACTION_DIM = 12 # ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1) ACTION_LOW = -1.0 ACTION_HIGH = 1.0 @@ -101,14 +101,15 @@ def _resolve_tasks(task: str) -> tuple[list[str], str | None]: def convert_action(flat_action: np.ndarray) -> dict[str, Any]: """Split a flat (12,) action vector into a RoboCasa action dict. - Layout: base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1) + Layout (openpi / robocasa.utils.env_utils.convert_action order): + ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1) """ return { - "action.base_motion": flat_action[0:4], - "action.control_mode": flat_action[4:5], - "action.end_effector_position": flat_action[5:8], - "action.end_effector_rotation": flat_action[8:11], - "action.gripper_close": flat_action[11:12], + "action.end_effector_position": flat_action[0:3], + "action.end_effector_rotation": flat_action[3:6], + "action.gripper_close": flat_action[6:7], + "action.base_motion": flat_action[7:11], + "action.control_mode": flat_action[11:12], } @@ -136,9 +137,16 @@ class RoboCasaEnv(gym.Env): episode_length: int | None = None, obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES, episode_index: int = 0, + terminate_on_success: bool = True, + horizon: int | None = None, ): super().__init__() self.task = task + # When False, a task-success does NOT end/reset the episode — used by the + # interactive sim so one kitchen persists across sequential prompts. + self.terminate_on_success = terminate_on_success + # Underlying robosuite horizon (steps before truncation). None -> default. + self.horizon = horizon self.obs_type = obs_type self.render_mode = render_mode self.observation_width = observation_width @@ -210,12 +218,16 @@ class RoboCasaEnv(gym.Env): # (only None/"all"/"pretrain"/"target" are valid). Always pass a # valid value so we don't hit that default. Extra kwargs are # forwarded to the underlying kitchen env via create_env/robosuite.make. + extra_kwargs: dict[str, Any] = {} + if self.horizon is not None: + extra_kwargs["horizon"] = int(self.horizon) self._env = RoboCasaGymEnv( env_name=self.task, camera_widths=self.observation_width, camera_heights=self.observation_height, split=self.split if self.split is not None else "all", obj_registries=self.obj_registries, + **extra_kwargs, ) ep_meta = self._env.env.get_ep_meta() @@ -230,12 +242,14 @@ class RoboCasaEnv(gym.Env): return {"pixels": images} # `state.*` keys come from PandaOmronKeyConverter inside the wrapper. + # openpi state order: ee first, then base, then gripper (matches the + # openpi robocasa pipeline / examples/robocasa/main.py state layout). agent_pos = np.concatenate( [ - raw_obs.get("state.base_position", np.zeros(3)), - raw_obs.get("state.base_rotation", np.zeros(4)), raw_obs.get("state.end_effector_position_relative", np.zeros(3)), raw_obs.get("state.end_effector_rotation_relative", np.zeros(4)), + raw_obs.get("state.base_position", np.zeros(3)), + raw_obs.get("state.base_rotation", np.zeros(4)), raw_obs.get("state.gripper_qpos", np.zeros(2)), ], axis=-1, @@ -280,7 +294,7 @@ class RoboCasaEnv(gym.Env): raw_obs, reward, done, truncated, info = self._env.step(action_dict) is_success = bool(info.get("success", False)) - terminated = done or is_success + terminated = done or (is_success and self.terminate_on_success) info.update({"task": self.task, "done": done, "is_success": is_success}) observation = self._format_raw_obs(raw_obs) @@ -313,6 +327,8 @@ def _make_env_fns( split: str | None, episode_length: int | None, obj_registries: Sequence[str], + terminate_on_success: bool = True, + horizon: int | None = None, ) -> list[Callable[[], RoboCasaEnv]]: """Build n_envs factory callables for a single task. @@ -335,6 +351,8 @@ def _make_env_fns( episode_length=episode_length, obj_registries=obj_registries, episode_index=episode_index, + terminate_on_success=terminate_on_success, + horizon=horizon, ) return [partial(_make_env, i) for i in range(n_envs)] @@ -348,6 +366,8 @@ def create_robocasa_envs( env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None, episode_length: int | None = None, obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES, + terminate_on_success: bool = True, + horizon: int | None = None, ) -> dict[str, dict[int, Any]]: """Create vectorized RoboCasa365 environments with a consistent return shape. @@ -409,6 +429,8 @@ def create_robocasa_envs( split=split, episode_length=episode_length, obj_registries=obj_registries, + terminate_on_success=terminate_on_success, + horizon=horizon, ) if is_async: diff --git a/src/lerobot/robots/utils.py b/src/lerobot/robots/utils.py index f897a560e..b3874e5cf 100644 --- a/src/lerobot/robots/utils.py +++ b/src/lerobot/robots/utils.py @@ -21,6 +21,8 @@ from lerobot.utils.import_utils import make_device_from_device_class from .config import RobotConfig from .robot import Robot +logger = logging.getLogger(__name__) + def make_robot_from_config(config: RobotConfig) -> Robot: # TODO(Steven): Consider just using the make_device_from_device_class for all types @@ -118,7 +120,7 @@ def ensure_safe_goal_position( } if warnings_dict: - logging.warning( + logger.warning( "Relative goal position magnitude had to be clamped to be safe.\n" f"{pformat(warnings_dict, indent=4)}" ) diff --git a/src/lerobot/runtime/__init__.py b/src/lerobot/runtime/__init__.py new file mode 100644 index 000000000..da537b93b --- /dev/null +++ b/src/lerobot/runtime/__init__.py @@ -0,0 +1,38 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Policy-agnostic runtime for language-conditioned policies. + +Adapters registered in :mod:`lerobot.runtime.registry` are served by ``lerobot-rollout --language``. +""" + +from .adapter import BaseLanguageAdapter, GenerationConfig, LanguageDiagnostics +from .language_runtime import ( + LanguageConditionedPolicyAdapter, + LanguageConditionedRuntime, + RuntimeState, + Tick, + TickClock, +) + +__all__ = [ + "BaseLanguageAdapter", + "GenerationConfig", + "LanguageConditionedPolicyAdapter", + "LanguageConditionedRuntime", + "LanguageDiagnostics", + "RuntimeState", + "Tick", + "TickClock", +] diff --git a/src/lerobot/runtime/adapter.py b/src/lerobot/runtime/adapter.py new file mode 100644 index 000000000..e94a19a4a --- /dev/null +++ b/src/lerobot/runtime/adapter.py @@ -0,0 +1,165 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Policy adapters for the language runtime. + +The base adapter owns generation control and diagnostics while subclasses provide policy-specific actions and text. +""" + +from __future__ import annotations + +import re +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any + +from .language_runtime import RuntimeState + +_SAY_RE = re.compile(r"<\s*say\s*>(.*?)<\s*/\s*say\s*>", re.IGNORECASE | re.DOTALL) + + +@dataclass +class GenerationConfig: + """Text-generation settings fixed for the adapter's lifetime.""" + + min_new_tokens: int = 0 + temperature: float = 0.0 + top_p: float = 1.0 + chunks_per_regen: int = 1 # regenerate the language context every N action chunks + enable_memory: bool = True # generate a running memory note on subtask change + enable_subtask: bool = True # generate the low-level subtask (off => use the given text directly) + + +@dataclass +class LanguageDiagnostics: + """Runtime-panel generation counters keyed by text kind.""" + + last_raw: dict[str, str] = field(default_factory=dict) + empty: dict[str, int] = field(default_factory=dict) + repeat: int = 0 + + def _bump(self, table: dict[str, int], kind: str) -> int: + table[kind] = table.get(kind, 0) + 1 + return table[kind] + + +class BaseLanguageAdapter(ABC): + """Batteries-included adapter: generic high-level control, policy primitives abstract.""" + + def __init__(self, policy: Any, gen: GenerationConfig | None = None) -> None: + self.policy = policy + self.gen = gen or GenerationConfig() + self.diag = LanguageDiagnostics() + self._chunks_until_regen = 0 + + @abstractmethod + def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: + """Produce an action chunk from the observation + current language context.""" + + @abstractmethod + def generate_text( + self, + kind: str, + observation: dict[str, Any] | None, + state: RuntimeState, + user_text: str | None = None, + ) -> str: + """Generate one text stream (``kind``) and return the decoded string.""" + + def update_language_state(self, observation: dict[str, Any] | None, state: RuntimeState) -> None: + """Throttled regeneration of the language context (subtask / memory / ...).""" + if self._chunks_until_regen > 0: + self._chunks_until_regen -= 1 + return + self._chunks_until_regen = max(1, self.gen.chunks_per_regen) - 1 + self._regenerate_context(observation, state) + + def handle_interjection( + self, user_text: str, observation: dict[str, Any] | None, state: RuntimeState + ) -> None: + """React to a mid-run user message by regenerating the plan.""" + out = self.generate_text("interjection", observation, state, user_text=user_text) + plan = self.plan_from_text(out) + if plan: + state.set_context("plan", plan, label="plan") + + def plan_from_text(self, text: str) -> str: + """Strip ```` speech markers from a generated plan.""" + plan, _speech = split_plan_and_say(text) + return plan + + def _regenerate_context(self, observation: dict[str, Any] | None, state: RuntimeState) -> None: + """Default hierarchy: regenerate the subtask, then memory when it changes. + + Override for a policy with a different language hierarchy. + """ + if not self.gen.enable_subtask: + # Preserve operator-provided subtasks in direct mode. + return + subtask = self._generate_filtered("subtask", observation, state) + if subtask is None: + return + previous = state.language_context.get("subtask") + if not state.set_context("subtask", subtask, label="subtask"): + self.diag.repeat += 1 + return + self.diag.repeat = 0 + if previous: + state.extra["prior_subtask"] = previous + if not self.gen.enable_memory: + return + memory = self._generate_filtered("memory", observation, state) + if memory is not None: + state.set_context("memory", memory, label="memory") + + def _generate_filtered( + self, kind: str, observation: dict[str, Any] | None, state: RuntimeState + ) -> str | None: + """Generate one ``kind``, record diagnostics, and drop empty output.""" + text = self.generate_text(kind, observation, state) + self.diag.last_raw[kind] = text or "" + if not text: + count = self.diag._bump(self.diag.empty, kind) + if count == 1 or count % 5 == 0: + state.log(f" [info] {kind} gen returned empty (x{count})") + return None + return text + + +class DirectTaskPolicyAdapter(BaseLanguageAdapter): + """Adapter for flat policies whose preprocessors condition actions on the operator's task.""" + + 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 split_plan_and_say(text: str) -> tuple[str, str]: + """Split ``plan speech`` into ``(plan, speech)``.""" + if not text: + return "", "" + match = _SAY_RE.search(text) + if not match: + return text.strip(), "" + speech = match.group(1).strip().strip('"').strip("'") + plan = (text[: match.start()] + text[match.end() :]).strip() + return plan, speech diff --git a/src/lerobot/runtime/cli.py b/src/lerobot/runtime/cli.py new file mode 100644 index 000000000..0be5b0485 --- /dev/null +++ b/src/lerobot/runtime/cli.py @@ -0,0 +1,1191 @@ +#!/usr/bin/env python +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Interactive CLI for language-conditioned policy rollouts. + +It supports a text-only REPL, real robots, and RoboCasa with local or Hub checkpoints. +""" + +from __future__ import annotations + +import argparse +import logging +import sys +from collections.abc import Callable +from contextlib import nullcontext +from typing import Any + +from .adapter import GenerationConfig +from .language_runtime import LanguageConditionedPolicyAdapter, LanguageConditionedRuntime + +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.", + ) + p.add_argument( + "--policy.path", + dest="policy_path", + type=str, + required=True, + help="Local directory or Hugging Face Hub repo id pointing at a trained ``pretrained_model``.", + ) + p.add_argument( + "--policy.device", + dest="policy_device", + type=str, + default=None, + help=( + "Override the checkpoint's ``config.device`` (e.g. ``cuda``, ``cpu``). " + "Some checkpoints ship ``device=cpu``; pass ``cuda`` to run on GPU." + ), + ) + p.add_argument( + "--task", + dest="task", + type=str, + default=None, + help=("Initial task. If omitted, enter a task at the interactive prompt."), + ) + p.add_argument( + "--mode", + dest="mode", + type=str, + choices=["action", "paused"], + default=None, + help=( + "Start-up run mode. 'action' runs the robot immediately on " + "--task; 'paused' (the default) comes up at the command line " + "with the robot idle. Flip any time with /action and /pause." + ), + ) + p.add_argument( + "--no_robot", + action="store_true", + help="Skip robot connection and open a language-only REPL.", + ) + # ``--robot.type`` enables real-time control while stdin remains interactive. + p.add_argument( + "--robot.type", + dest="robot_type", + type=str, + default=None, + help=( + "Robot config choice (e.g. ``so101``, ``so101_follower``). " + "When set, the runtime drives the actual robot at " + "``--ctrl_hz`` instead of the no-robot REPL. Implies " + "``--autonomous`` unless ``--no_robot`` is also " + "passed (in which case the flag is ignored). See " + "``lerobot.robots`` for available choices." + ), + ) + p.add_argument( + "--rerun", + action="store_true", + help="Live rerun viewer for the robot cameras (real-robot mode). Serves a " + "headless web viewer; forward --rerun.web_port and --rerun.grpc_port over SSH.", + ) + p.add_argument( + "--rerun.web_port", + dest="rerun_web_port", + type=int, + default=9090, + help="rerun web-viewer port (default 9090).", + ) + p.add_argument( + "--rerun.grpc_port", + dest="rerun_grpc_port", + type=int, + default=9876, + help="rerun gRPC data port (default 9876).", + ) + p.add_argument( + "--direct_subtask", + action="store_true", + help="Direct-subtask mode (sim OR robot): your typed text IS the subtask " + "fed to the action expert; the LM subtask generator is disabled.", + ) + # ``--sim`` uses the eval pipeline and is mutually exclusive with a robot. + p.add_argument( + "--sim", + action="store_true", + help=( + "Run the policy in the RoboCasa simulator instead of on a real " + "robot. Select the scene with --sim.task; type prompts with " + "/action to have the policy execute them in that scene." + ), + ) + p.add_argument( + "--sim.task", + dest="sim_task", + type=str, + default="CloseFridge", + help="RoboCasa task/scene to instantiate (e.g. OpenDrawer, LoadDishwasher).", + ) + p.add_argument( + "--sim.split", + dest="sim_split", + type=str, + default="pretrain", + help="RoboCasa scene split (all/pretrain/target). Default: pretrain.", + ) + p.add_argument( + "--sim.obj_registries", + dest="sim_obj_registries", + type=str, + default="objaverse,lightwheel", + help="Comma-separated object-mesh registries. Default: objaverse,lightwheel.", + ) + p.add_argument( + "--sim.seed", + dest="sim_seed", + type=int, + default=1000, + help="Seed for RoboCasa scene reset (default: 1000, matches eval).", + ) + p.add_argument( + "--sim.record", + dest="sim_record", + type=str, + choices=["mp4", "off"], + default="mp4", + help="Record an annotated mp4 (task/subtask/memory overlay) of the sim session. Default: mp4.", + ) + p.add_argument( + "--sim.output_dir", + dest="sim_output_dir", + type=str, + default="outputs/runtime_sim", + help="Directory for the recorded sim video (default: outputs/runtime_sim).", + ) + p.add_argument( + "--sim.render_size", + dest="sim_render_size", + type=int, + default=384, + help=( + "Resolution (px) of the observation cameras used for the display " + "(default 384; try 512 for sharper, 256 for faster). The policy is " + "unaffected — it resizes to 224 internally." + ), + ) + p.add_argument( + "--sim.views", + dest="sim_views", + type=str, + default="robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right", + help=( + "Comma-separated camera views to show side by side. Default shows " + "left, wrist (eye-in-hand), right. Use e.g. 'robot0_eye_in_hand' " + "for wrist-only." + ), + ) + p.add_argument( + "--sim.stream_port", + dest="sim_stream_port", + type=int, + default=8010, + help=( + "Port for the live MJPEG viewer (default: 8010; 0 disables). " + "Open http://localhost: in a browser; over SSH forward it with " + "ssh -L :localhost: ." + ), + ) + p.add_argument( + "--chunk_hz", + type=float, + default=1.0, + help=( + "Action-chunk generation rate (Hz). Default ``1.0`` — one " + "new chunk per second. Lower = less inference cost / " + "smoother behaviour but longer reaction time to changes. " + "Higher = fresher actions / more inference cost; cap at " + "~1/(forward-pass latency)." + ), + ) + p.add_argument("--ctrl_hz", type=float, default=50.0, help="Action dispatch rate.") + p.add_argument( + "--high_level_hz", + type=float, + default=1.0, + help="High-level subtask generation rate.", + ) + p.add_argument( + "--sim.direct_subtask", + dest="sim_direct_subtask", + action="store_true", + help=( + "Direct-subtask mode: what you type IS the subtask fed to the action " + "expert (no LM subtask generation). Good when the model's subtask " + "head is weak — you steer the policy with exact imperatives." + ), + ) + p.add_argument( + "--disable_memory", + action="store_true", + help=( + "Skip the memory-note generation on subtask change. Use for " + "subtask-only checkpoints (no memory head) — avoids a wasted LM " + "decode and a meaningless memory line." + ), + ) + p.add_argument( + "--fp8", + action="store_true", + help=( + "PI052 only: enable FlashRT FP8 MLP kernels. Calibrates on the " + "first inference call and swaps every Gemma + SigLIP MLP to fused " + "FP8 (needs the `kernels` package and CUDA SM>=8.9; degrades to " + "BF16 otherwise). Speeds up the forward pass at a small accuracy " + "cost — see policies/pi052/flashrt_fp8.py." + ), + ) + p.add_argument( + "--subtask_chunks_per_gen", + type=int, + default=1, + help=( + "Throttle subtask gen to once every N action-chunk boundaries. " + "Default 1 = regenerate the subtask on every chunk refresh. " + "Set to 5 to run ~5 flow-matching action chunks per LM-head " + "subtask gen — saves compute and avoids re-planning trajectories " + "mid-grasp when a subtask is still valid across multiple chunks." + ), + ) + p.add_argument( + "--max_ticks", + type=int, + default=None, + help="Stop after N ticks (debug / smoke-test).", + ) + p.add_argument( + "--text_min_new_tokens", + type=int, + default=0, + help=( + "Debug knob for under-trained checkpoints: force the LM head " + "to emit at least N non-EOS tokens before EOS is allowed. " + "Use when the head's prior at position 0 still favours EOS " + "(short training run on a chat-pretrained backbone). 3-5 " + "is usually enough to reveal whether the model has real " + "subtask-token mass under the EOS argmax." + ), + ) + p.add_argument( + "--text_temperature", + type=float, + default=0.0, + help=( + "Sampling temperature for high-level text gen. 0 = greedy " + "argmax (default, matches training). Set 0.3-0.7 with an " + "under-trained checkpoint to escape stuck-at-EOS argmax." + ), + ) + p.add_argument( + "--text_top_p", + type=float, + default=1.0, + help="Nucleus filtering for high-level text gen.", + ) + p.add_argument("-v", "--verbose", action="store_true", help="Enable DEBUG logging.") + 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 +# them so the recipe render + text-tokenizer processor steps are no-ops. +_RUNTIME_OWNED_LANGUAGE_COLS = ("language_persistent", "language_events") + + +def _strip_runtime_owned_language_cols(sample: dict) -> None: + """In-place drop of language columns the runtime owns at inference.""" + for k in _RUNTIME_OWNED_LANGUAGE_COLS: + sample.pop(k, None) + + +# Non-observation model inputs emitted by processors such as MolmoAct2's. +_MODEL_INPUT_PASSTHROUGH_KEYS = ( + "input_ids", + "attention_mask", + "token_type_ids", + "pixel_values", + "image_token_pooling", + "image_grids", + "image_num_crops", + "pixel_values_videos", + "video_token_pooling", + "video_grids", +) + + +def _select_observation_to_device(sample: dict, device: Any) -> dict: + """Keep ``observation.*`` (+ model-input passthrough) keys, move tensors to ``device``.""" + import torch # noqa: PLC0415 + + return { + k: v.to(device) if isinstance(v, torch.Tensor) else v + for k, v in sample.items() + if isinstance(k, str) and (k.startswith("observation.") or k in _MODEL_INPUT_PASSTHROUGH_KEYS) + } + + +def _load_policy_and_preprocessor( + policy_path: str, + *, + load_processors_from_checkpoint: bool = False, + fp8: bool = False, + device: str | None = None, +) -> tuple[Any, Any, Any]: + """Load a local or Hub policy, optionally with its eval processors.""" + from lerobot.configs import PreTrainedConfig # noqa: PLC0415 + from lerobot.policies.factory import get_policy_class, make_pre_post_processors # noqa: PLC0415 + + cfg = PreTrainedConfig.from_pretrained(policy_path) + cfg.pretrained_path = policy_path + + # Optional device override — some checkpoints ship device=cpu. + if device: + cfg.device = device + + # Variable prompts trigger recompilation, and checkpointing only adds inference overhead. + if getattr(cfg, "compile_model", False): + cfg.compile_model = False + if getattr(cfg, "gradient_checkpointing", False): + cfg.gradient_checkpointing = False + + # Opt-in FP8: only PI052 has the FlashRT MLP swap. The policy calibrates and + # swaps on its first predict_action_chunk when this flag is set. + if fp8: + if hasattr(cfg, "use_flashrt_fp8_mlp"): + cfg.use_flashrt_fp8_mlp = True + logging.info("[runtime] FP8 MLP kernels enabled (FlashRT) for %s", cfg.type) + else: + logging.warning( + "[runtime] --fp8 ignored: %s has no use_flashrt_fp8_mlp (PI052 only).", + cfg.type, + ) + + preprocessor = None + postprocessor = None + policy_cls = get_policy_class(cfg.type) + policy = policy_cls.from_pretrained(policy_path, config=cfg) + policy.to(cfg.device) + if load_processors_from_checkpoint: + preprocessor, postprocessor = make_pre_post_processors( + cfg, + pretrained_path=cfg.pretrained_path, + preprocessor_overrides={"device_processor": {"device": str(cfg.device)}}, + ) + + policy.eval() + return policy, preprocessor, postprocessor + + +def _build_language_rollout_context(args: argparse.Namespace) -> Any: + """Build the canonical rollout context for a language-controlled robot.""" + import threading # noqa: PLC0415 + + import draccus # noqa: PLC0415 + + from lerobot.configs import parser # noqa: PLC0415 + from lerobot.rollout import RolloutConfig, build_rollout_context # noqa: PLC0415 + + # Import for bundled Draccus camera and robot registrations. + 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", + ) + ) + + previous_argv = sys.argv + try: + # 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 + + 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) + + return build_rollout_context(cfg, threading.Event()) + + +def _build_rollout_runtime_io( + ctx: Any, + *, + 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 + make_robot_action, + prepare_observation_for_inference, + ) + from lerobot.utils.feature_utils import build_dataset_frame # noqa: PLC0415 + from lerobot.utils.rerun_visualization import log_rerun_data # noqa: PLC0415 + + robot = ctx.hardware.robot_wrapper + device = torch.device(ctx.runtime.cfg.device or "cpu") + latest_raw: dict[str, Any] = {} + + def _provider() -> dict | None: + try: + raw = robot.get_observation() + latest_raw.clear() + latest_raw.update(raw) + if rerun_log: + try: + log_rerun_data(observation=raw) + except Exception as exc: # noqa: BLE001 + logger.debug("rerun observation log failed: %s", exc) + _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 observation pipeline failed: %s", exc) + return None + + def _executor(action: Any) -> None: + try: + 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(processed_action)) + return + raw = latest_raw or robot.get_observation() + robot_action = ctx.processors.robot_action_processor((action_dict, raw)) + robot.send_action(robot_action) + except Exception as exc: # noqa: BLE001 + logger.error("robot action pipeline failed: %s", exc, exc_info=True) + + return _provider, _executor + + +def _print_runtime_help() -> None: + """Print the slash-command reference.""" + print( + "[runtime] commands (arguments need no quotes):\n" + " /action run the robot; an argument switches to that task\n" + " /action resume the robot on the current task\n" + " /action run the robot for N seconds, then auto-pause\n" + " /pause pause the action loop — robot holds position\n" + " /help show this help\n" + " stop | quit | exit end the session", + flush=True, + ) + + +def _is_number(text: str) -> bool: + """True if ``text`` parses as a float (a ``/action`` duration arg).""" + try: + float(text) + return True + except ValueError: + return False + + +def _strip_quotes(text: str) -> str: + """Strip one pair of surrounding quotes from a command argument.""" + text = text.strip() + if len(text) >= 2 and text[0] == text[-1] and text[0] in {'"', "'"}: + return text[1:-1].strip() + return text + + +def _clear_action_queue(runtime: Any) -> None: + """Drop any queued action chunk so nothing fires while paused.""" + lock = getattr(runtime.state, "lock", nullcontext()) + with lock: + queue = runtime.state.get("action_queue") + if hasattr(queue, "clear"): + queue.clear() + + +def _handle_slash_command(runtime: Any, line: str) -> bool: + """Dispatch the runtime slash commands. + + ``/action ["task"]`` run the robot; a quoted/bare argument sets a + new task, a bare number is a timed burst + (seconds), no argument resumes the current + task. + ``/pause`` pause the action loop — the robot holds. + ``/help`` print the command reference. + + Returns ``True`` when ``line`` was a recognised command (consumed). + """ + stripped = line.strip() + if not stripped.startswith("/"): + return False + head, _, rest = stripped.partition(" ") + cmd = head.lower() + rest = _strip_quotes(rest) + + if cmd in {"/action", "/act", "/run"}: + runtime.state["mode"] = "action" + if rest and _is_number(rest): + import time as _time # noqa: PLC0415 + + secs = float(rest) + runtime.state["action_deadline"] = _time.monotonic() + secs + print( + f"[runtime] action — running {secs:g}s, then auto-pause", + flush=True, + ) + else: + runtime.state["action_deadline"] = None + if rest: + runtime.set_task(rest) + # New task → drop the stale subtask so the high-level + # loop regenerates one for the new goal. + runtime.state.set_context("subtask", None) + print(f"[runtime] action — task: {rest!r}", flush=True) + elif runtime.state.get("task"): + print( + f"[runtime] action — resuming: {runtime.state['task']!r}", + flush=True, + ) + else: + runtime.state["mode"] = "paused" + print( + "[runtime] no task set — use /action ", + flush=True, + ) + return True + + if cmd in {"/pause", "/p"}: + runtime.state["mode"] = "paused" + runtime.state["action_deadline"] = None + _clear_action_queue(runtime) + print("[runtime] paused — robot holding position", flush=True) + return True + + if cmd in {"/help", "/?"}: + _print_runtime_help() + return True + return False + + +def _make_state_panel_renderer( + runtime: Any, + *, + mode_label: str, + panel_label: str = "Runtime", + scrollback: list[str] | None = None, +) -> Callable[[list[str] | None], None]: + """Return a closure that prints the task/subtask/plan/memory panel. + + Used by ``_run_repl`` for the no-robot language REPL. + """ + from rich.console import Console # noqa: PLC0415 + + console = Console(highlight=False) + + def _redraw(robot_lines: list[str] | None = None) -> None: + console.clear() + st = runtime.state + run_mode = st.get("mode", "action") + mode_tag = "[green]mode: action[/]" if run_mode == "action" else "[yellow]mode: paused[/]" + console.rule(f"[bold]{panel_label}[/] · {mode_label} · {mode_tag}", style="cyan") + # Always-visible command hint so the operator never has to + # remember the slash commands. + if run_mode == "action": + console.print(" [dim]commands:[/] [bold]/pause[/] stop · [bold]/help[/] · [bold]stop[/]") + else: + console.print( + " [dim]commands:[/] [bold]/action[/] run · [bold]/help[/] · [bold]stop[/]" + ) + display_values = { + "task": st.get("task"), + **(st.get("language_context") or {}), + } + for key in ("task", "subtask", "plan", "memory"): + value = display_values.get(key) + if value: + console.print(f" [bold cyan]{key:<8}[/] {value}") + else: + console.print(f" [dim]{key:<8} (not set)[/]") + queue_len = ( + len(st["action_queue"]) + if isinstance(st.get("action_queue"), list | tuple) or hasattr(st.get("action_queue"), "__len__") + else 0 + ) + dispatched = int(st.get("actions_dispatched") or 0) + console.print(f" [dim]queued actions: {queue_len} dispatched: {dispatched}[/]") + + # Surface repeated or empty generations as overfitting diagnostics. + diag = getattr(runtime.policy_adapter, "diag", None) + if diag is not None: + raw_subtask = diag.last_raw.get("subtask") + sub_rep = int(diag.repeat) + sub_empty = int(diag.empty.get("subtask", 0)) + if raw_subtask is not None or sub_rep or sub_empty: + raw_display = (raw_subtask or "(empty)")[:80] + color = "yellow" if (sub_rep >= 3 or sub_empty >= 3) else "dim" + console.print( + f" [{color}]subtask diag repeat:{sub_rep} empty:{sub_empty} " + f"last_raw: {raw_display!r}[/]" + ) + console.rule(style="cyan") + # Show recent generation warnings and speech oldest-first. + if scrollback: + for line in scrollback: + console.print(f" [magenta]{line.rstrip()}[/]") + console.rule(style="cyan") + if robot_lines: + for line in robot_lines: + console.print(f" [magenta]{line.strip()}[/]") + console.print() + if not st.get("task"): + console.print( + " [dim]Type [bold]/action [/bold] to begin, /help for commands, stop to exit.[/]" + ) + + return _redraw + + +def _silence_noisy_loggers() -> None: + """Keep request-level third-party logs out of the interactive prompt.""" + for name in ( + "httpcore", + "httpcore.connection", + "httpcore.http11", + "httpcore.proxy", + "httpx", + "urllib3", + "urllib3.connectionpool", + "huggingface_hub", + "huggingface_hub.repocard", + "huggingface_hub.file_download", + "transformers", + "transformers.modeling_utils", + "transformers.tokenization_utils_base", + "datasets", + "filelock", + ): + logging.getLogger(name).setLevel(logging.WARNING) + + # Clamp warnings can fire every control tick and flood the panel. + logging.getLogger("lerobot.robots.utils").setLevel(logging.ERROR) + + +def run( + argv: list[str] | None = None, + *, + adapter_factory: Callable[[Any, GenerationConfig], LanguageConditionedPolicyAdapter] | None = None, + panel_label: str | None = None, + prog: str = "lerobot-rollout", +) -> int: + """Run the interactive language-conditioned runtime CLI. + + ``adapter_factory`` turns ``(policy, GenerationConfig)`` into a + :class:`LanguageConditionedPolicyAdapter` (typically the adapter class). + When ``None`` it is resolved from :mod:`lerobot.runtime.registry` by the + loaded policy's type, so the ``lerobot-rollout`` entry + point serves every registered policy. ``panel_label`` defaults to the + policy type. + """ + args = _parse_args(argv, prog=prog) + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(asctime)s %(levelname)s %(message)s", + ) + _silence_noisy_loggers() + + sim_mode = bool(getattr(args, "sim", False)) and not args.no_robot + autonomous_mode = bool(args.robot_type) and not args.no_robot + if sim_mode and autonomous_mode: + print( + "[runtime] ERROR: --sim and --robot.type are mutually exclusive " + "(pick a simulator scene OR a real robot).", + file=sys.stderr, + ) + return 2 + # Fork the simulator before CUDA initialization to avoid inherited EGL corruption. + sim_env = None + sim_obs = None + sim_stream_server = None + sim_holder: dict[str, Any] = {"backend": None} + if sim_mode: + from lerobot.runtime.sim_robocasa import create_sim_env, start_mjpeg_server # noqa: PLC0415 + + # Start the live viewer first so the port listens during the ~60s model + # load (browsers get a loading page instead of connection-refused). + if args.sim_stream_port: + sim_stream_server = start_mjpeg_server( + args.sim_stream_port, + lambda: sim_holder["backend"]._latest_frame if sim_holder["backend"] else None, + ) + print( + f"[runtime] starting RoboCasa sim scene={args.sim_task!r} split={args.sim_split!r}", + flush=True, + ) + sim_env, sim_obs = create_sim_env( + task=args.sim_task, + split=args.sim_split, + obj_registries=[r.strip() for r in args.sim_obj_registries.split(",") if r.strip()], + seed=args.sim_seed, + render_size=args.sim_render_size, + ) + + 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 + else: + print(f"[runtime] loading policy from {args.policy_path}", flush=True) + policy, preprocessor, postprocessor = _load_policy_and_preprocessor( + args.policy_path, + 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: + from .registry import get_language_adapter_factory # noqa: PLC0415 + + adapter_factory = get_language_adapter_factory(policy_type) + if panel_label is None: + panel_label = str(policy_type or "runtime").upper() + + # Default to idle until the operator supplies a command. + startup_mode = args.mode or "paused" + + observation_provider: Callable[[], dict | None] | None = None + robot_executor: Callable[[Any], None] | None = None + robot = None + sim_backend = None + # Late-bound handle to the runtime so the robot observation provider can read + # the live task/subtask each frame (the runtime is created further below). + runtime_box: dict[str, Any] = {} + + def _live_task() -> str | None: + rt = runtime_box.get("rt") + if rt is None: + return args.task + return rt.state.language_context.get("subtask") or rt.state.task or args.task + + if sim_mode: + from lerobot.runtime.sim_robocasa import RoboCasaSimBackend # noqa: PLC0415 + + sim_backend = RoboCasaSimBackend( + env=sim_env, + last_obs=sim_obs, + task=args.sim_task, + seed=args.sim_seed, + device=str(getattr(policy.config, "device", "cpu")), + preprocessor=preprocessor, + postprocessor=postprocessor, + record=(args.sim_record == "mp4"), + output_dir=args.sim_output_dir, + view_cams=[v.strip() for v in args.sim_views.split(",") if v.strip()], + ) + observation_provider = sim_backend.observation_provider + robot_executor = sim_backend.action_executor + 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 + if sim_stream_server is not None: + sim_backend.attach_stream_server(sim_stream_server) + elif autonomous_mode: + rerun_log = False + if args.rerun: + from lerobot.utils.rerun_visualization import init_rerun # noqa: PLC0415 + + try: + init_rerun( + session_name=f"lerobot_{policy_type or 'runtime'}", + port=args.rerun_grpc_port, + web_port=args.rerun_web_port, + ) + rerun_log = True + print( + f"[runtime] rerun live view: http://localhost:{args.rerun_web_port}", + flush=True, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("could not start rerun: %s", exc) + 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=rerun_log, + get_task=_live_task, + ) + # Generation settings belong to the adapter rather than mutable runtime state. + gen_config = GenerationConfig( + min_new_tokens=int(args.text_min_new_tokens or 0), + temperature=float(args.text_temperature or 0.0), + top_p=float(args.text_top_p or 1.0), + chunks_per_regen=max(1, int(args.subtask_chunks_per_gen or 1)), + enable_memory=not bool(getattr(args, "disable_memory", False)), + enable_subtask=not _direct_subtask_enabled(args), + ) + runtime = LanguageConditionedRuntime( + policy_adapter=adapter_factory(policy, gen_config), + observation_provider=observation_provider, + action_executor=robot_executor, + event_collector=None, + chunk_hz=args.chunk_hz, + ctrl_hz=args.ctrl_hz, + high_level_hz=args.high_level_hz, + ) + # Let the robot observation provider read the live task/subtask each frame. + runtime_box["rt"] = runtime + # Apply the configured startup mode. + runtime.state["mode"] = startup_mode + if args.task: + runtime.set_task(args.task) + + # Let the sim backend read live task/subtask/memory for the video overlay. + if sim_backend is not None: + sim_backend.bind_runtime(runtime) + # Keep EGL rendering on the main thread. + return _run_sim_interactive( + runtime, + sim_backend, + initial_task=args.task, + max_ticks=args.max_ticks, + panel_label=panel_label, + direct_subtask=_direct_subtask_enabled(args), + ) + + if autonomous_mode: + return _run_robot_interactive( + runtime, + robot, + initial_task=args.task, + max_ticks=args.max_ticks, + direct_subtask=_direct_subtask_enabled(args), + panel_label=panel_label, + ) + return _run_repl(runtime, initial_task=args.task, max_ticks=args.max_ticks, panel_label=panel_label) + + +def _direct_subtask_enabled(args: Any) -> bool: + """Direct-subtask mode via either the general or sim-scoped flag.""" + return bool(getattr(args, "direct_subtask", False) or getattr(args, "sim_direct_subtask", False)) + + +def _run_sim_interactive( + runtime: Any, + sim_backend: Any, + *, + initial_task: str | None, + max_ticks: int | None, + panel_label: str = "Runtime", + direct_subtask: bool = False, +) -> int: + """Keep RoboCasa rendering on the main thread while polling stdin.""" + import select # noqa: PLC0415 + import time # noqa: PLC0415 + + import torch # noqa: PLC0415 + + if initial_task: + runtime.set_task(initial_task) + # In direct-subtask mode the typed text IS the subtask; otherwise clear + # it so the model generates one. + runtime.state.set_context("subtask", initial_task if direct_subtask else None) + runtime.state["mode"] = "action" + + # Keep the terminal quiet while the browser renders the rollout. + _mode_line = ( + " Mode: DIRECT subtask (your text drives the action expert as-is)\n" + if direct_subtask + else " Mode: task (the model generates a subtask from your text)\n" + ) + print( + f"\n{'=' * 64}\n" + f" {panel_label} — RoboCasa interactive sim (one persistent kitchen)\n" + f"{_mode_line}" + f" Type a command + Enter to run it, e.g. open the fridge\n" + f" Commands: /pause · /resume · /reset (new kitchen) · stop\n" + f"{'=' * 64}", + flush=True, + ) + + def _prompt() -> None: + print("\n> ", end="", flush=True) + + _prompt() + ticks_done = 0 + stdin_open = True + try: + while True: + # Non-blocking stdin: a full line (canonical-mode terminal) is read + # only when Enter is pressed, so line editing works normally. + if stdin_open and select.select([sys.stdin], [], [], 0)[0]: + line = sys.stdin.readline() + if line == "": # EOF — keep running the sim, stop reading stdin + stdin_open = False + else: + cmd = line.strip() + if cmd: + low = cmd.lower() + if low in {"stop", "quit", "exit"}: + break + elif low in {"/pause", "pause", "/p"}: + runtime.state["mode"] = "paused" + _clear_action_queue(runtime) + print("[paused] robot holding", flush=True) + elif low in {"/resume", "resume", "/run"}: + runtime.state["mode"] = "action" + print("[running]", flush=True) + elif low in {"/reset", "reset"}: + sim_backend.reset_scene() + _clear_action_queue(runtime) + runtime.state.set_context("subtask", None) + if hasattr(runtime.policy, "reset"): + runtime.policy.reset() + print("[reset] new kitchen scene", flush=True) + else: + # Clear queued actions and rearm generation for a new command. + runtime.set_task(cmd) + # Direct mode: the typed text is the subtask itself; + # otherwise clear it so the model regenerates one. + runtime.state.set_context("subtask", cmd if direct_subtask else None) + _clear_action_queue(runtime) + adapter = getattr(runtime, "policy_adapter", None) + if adapter is not None and hasattr(adapter, "_chunks_until_regen"): + adapter._chunks_until_regen = 0 + gate = getattr(runtime, "_language_gate", None) + if gate is not None and hasattr(gate, "rearm"): + gate.rearm() + runtime.state["mode"] = "action" + print(f"[running] {cmd}", flush=True) + _prompt() + + # Match lerobot-eval's inference context on the main thread. + if runtime.state.get("mode", "paused") == "action": + with torch.inference_mode(): + runtime.step_once() + ticks_done += 1 + else: + time.sleep(0.05) # idle only while paused (robot not moving) + if runtime.state.stop: + break + if max_ticks is not None and ticks_done >= max_ticks: + break + except KeyboardInterrupt: + print("\n[stopping]", flush=True) + finally: + runtime.stop() + try: + sim_backend.disconnect() + except Exception as exc: # noqa: BLE001 + print(f"[runtime] WARNING: sim disconnect raised {exc}", flush=True) + return 0 + + +def _run_robot_interactive( + runtime: Any, + robot: Any, + *, + initial_task: str | None, + max_ticks: int | None, + direct_subtask: bool = False, + panel_label: str = "Runtime", +) -> int: + """Run steady robot control in the background and commands in the foreground.""" + import threading # noqa: PLC0415 + import time # noqa: PLC0415 + + if initial_task: + runtime.set_task(initial_task) + runtime.state.set_context("subtask", initial_task if direct_subtask else None) + # An explicit initial task starts immediately; otherwise the robot stays paused. + runtime.state["mode"] = "action" + + mode_line = ( + "DIRECT subtask (your text drives the action expert)" + if direct_subtask + else "task (the model generates a subtask from your text)" + ) + starting_action = runtime.state.get("mode", "paused") == "action" + start_line = ( + f" Starting in ACTION — the ARM WILL MOVE NOW on: {initial_task!r}\n" + if starting_action + else " Starts PAUSED. Type a command + Enter to run it — the ARM WILL MOVE.\n" + ) + print( + f"\n{'=' * 64}\n" + f" {panel_label} — OMX robot runtime · Mode: {mode_line}\n" + f"{start_line}" + f" Commands: /pause · /resume · stop\n" + f"{'=' * 64}", + flush=True, + ) + + thread = threading.Thread( + target=runtime.run, kwargs={"max_ticks": max_ticks}, name="runtime-loop", daemon=True + ) + thread.start() + try: + while thread.is_alive(): + try: + line = input("\n> ").strip() + except EOFError: + break + if not line: + continue + low = line.lower() + if low in {"stop", "quit", "exit"}: + break + elif low in {"/pause", "pause", "/p"}: + runtime.state["mode"] = "paused" + _clear_action_queue(runtime) + print("[paused] robot holding", flush=True) + elif low in {"/resume", "resume", "/run"}: + runtime.state["mode"] = "action" + print("[running]", flush=True) + else: + # New command: switch task/subtask immediately and regenerate. + runtime.set_task(line) + runtime.state.set_context("subtask", line if direct_subtask else None) + _clear_action_queue(runtime) + adapter = getattr(runtime, "policy_adapter", None) + if adapter is not None and hasattr(adapter, "_chunks_until_regen"): + adapter._chunks_until_regen = 0 + gate = getattr(runtime, "_language_gate", None) + if gate is not None and hasattr(gate, "rearm"): + gate.rearm() + runtime.state["mode"] = "action" + print(f"[running] {line}", flush=True) + except KeyboardInterrupt: + print("\n[stopping]", flush=True) + finally: + runtime.stop() + 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 _run_repl( + runtime: Any, *, initial_task: str | None, max_ticks: int | None, panel_label: str = "Runtime" +) -> int: + """Redraw the status block and logs once per REPL turn.""" + try: + from rich.console import Console # noqa: PLC0415 + except ImportError: + print( + "[runtime] rich is required for the interactive REPL. `pip install rich` and re-run.", + file=sys.stderr, + ) + return 2 + + _redraw = _make_state_panel_renderer(runtime, mode_label="no robot", panel_label=panel_label) + console = Console(highlight=False) + + last_logs: list[str] = [] + _redraw() + if initial_task is None: + # Already shown the help line in _redraw when task is None. + pass + ticks_done = 0 + try: + while True: + try: + line = console.input("[bold cyan]> [/]").strip() + except EOFError: + break + if not line: + _redraw(last_logs) + continue + lower = line.lower() + if lower in {"stop", "quit", "exit"}: + break + + # Slash commands, including VQA questions, run inline. + if _handle_slash_command(runtime, line): + last_logs = list(runtime.state.get("log_lines") or []) + _redraw(last_logs) + ticks_done += 1 + if max_ticks is not None and ticks_done >= max_ticks: + break + continue + + # A bare (non-slash) line is a user interjection — needs a + # task to be meaningful. + if not runtime.state.get("task"): + print( + "[runtime] no task yet — use /action ", + flush=True, + ) + _redraw(last_logs) + continue + runtime.state["recent_interjection"] = line + runtime.state.emit("user_interjection") + + last_logs = runtime.step_once() or [] + _redraw(last_logs) + + ticks_done += 1 + if max_ticks is not None and ticks_done >= max_ticks: + break + except KeyboardInterrupt: + console.print("\n[dim]interrupted[/]") + console.print("[dim]runtime stopped[/]") + return 0 diff --git a/src/lerobot/runtime/language_runtime.py b/src/lerobot/runtime/language_runtime.py new file mode 100644 index 000000000..419028c90 --- /dev/null +++ b/src/lerobot/runtime/language_runtime.py @@ -0,0 +1,349 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Small reusable runtime for language-conditioned robot policies.""" + +from __future__ import annotations + +import logging +import threading +import time +from collections import deque +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any, Protocol + +logger = logging.getLogger(__name__) + + +@dataclass +class RuntimeState: + """Explicit state shared by the runtime and policy adapter.""" + + task: str = "" + language_context: dict[str, str] = field(default_factory=dict) + action_queue: deque[Any] = field(default_factory=deque) + events: set[str] = field(default_factory=set) + log_lines: list[str] = field(default_factory=list) + mode: str = "action" + stop: bool = False + tick: Tick | None = None + actions_dispatched: int = 0 + action_deadline: float | None = None + extra: dict[str, Any] = field(default_factory=dict) + revision: int = 0 + lock: Any = field(default_factory=threading.RLock, repr=False) + + def emit(self, event_name: str) -> None: + self.events.add(event_name) + + def take_event(self, event_name: str) -> bool: + if event_name not in self.events: + return False + self.events.remove(event_name) + return True + + def log(self, line: str) -> None: + self.log_lines.append(line) + + def set_context(self, key: str, value: str | None, *, label: str | None = None) -> bool: + with self.lock: + previous = self.language_context.get(key) + if previous == value: + return False + if value is None: + self.language_context.pop(key, None) + else: + self.language_context[key] = value + self.revision += 1 + if label is not None and value: + self.log(f" {label}: {value}") + return True + + def get(self, key: str, default: Any = None) -> Any: + try: + return self[key] + except KeyError: + return default + + def setdefault(self, key: str, default: Any = None) -> Any: + current = self.get(key, None) + if current is not None: + return current + self[key] = default + return default + + def __getitem__(self, key: str) -> Any: + if hasattr(self, key): + return getattr(self, key) + if key in self.extra: + return self.extra[key] + raise KeyError(key) + + def __setitem__(self, key: str, value: Any) -> None: + with self.lock: + if hasattr(self, key): + if key == "mode" and self.mode != value: + self.revision += 1 + setattr(self, key, value) + else: + self.extra[key] = value + + +class LanguageConditionedPolicyAdapter(Protocol): + """Runtime policy contract, implemented directly or through ``BaseLanguageAdapter``.""" + + def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: ... + + def update_language_state(self, observation: dict[str, Any] | None, state: RuntimeState) -> None: ... + + def handle_interjection( + self, user_text: str, observation: dict[str, Any] | None, state: RuntimeState + ) -> None: ... + + +@dataclass +class Tick: + index: int + monotonic_seconds: float + + +@dataclass +class TickClock: + max_rate_hz: float = 50.0 + _index: int = field(default=0, init=False) + _last_seconds: float | None = field(default=None, init=False) + + def advance(self) -> Tick: + period = 1.0 / max(self.max_rate_hz, 0.1) + now = time.monotonic() + if self._last_seconds is not None: + sleep_for = (self._last_seconds + period) - now + if sleep_for > 0: + time.sleep(sleep_for) + now = time.monotonic() + self._last_seconds = now + self._index += 1 + return Tick(index=self._index, monotonic_seconds=now) + + +@dataclass +class _RateGate: + hz: float + _last_seconds: float | None = None + + def due(self, tick: Tick, *, force: bool = False) -> bool: + if force: + self._last_seconds = tick.monotonic_seconds + return True + period = 1.0 / max(self.hz, 1e-6) + if self._last_seconds is None or tick.monotonic_seconds - self._last_seconds >= period: + self._last_seconds = tick.monotonic_seconds + return True + return False + + def rearm(self) -> None: + self._last_seconds = None + + +@dataclass +class LanguageConditionedRuntime: + """Generic tick loop for language-conditioned robot policies.""" + + policy_adapter: LanguageConditionedPolicyAdapter + observation_provider: Callable[[], dict[str, Any] | None] | None = None + action_executor: Callable[[Any], None] | None = None + event_collector: Callable[[RuntimeState], None] | None = None + chunk_hz: float = 4.0 + ctrl_hz: float = 50.0 + high_level_hz: float = 1.0 + max_rate_hz: float = 50.0 + + state: RuntimeState = field(default_factory=RuntimeState) + _chunk_gate: _RateGate = field(init=False) + _ctrl_gate: _RateGate = field(init=False) + _language_gate: _RateGate = field(init=False) + _stop: bool = field(default=False, init=False) + _last_dispatch_seconds: float | None = field(default=None, init=False) + + def __post_init__(self) -> None: + self._chunk_gate = _RateGate(self.chunk_hz) + self._ctrl_gate = _RateGate(self.ctrl_hz) + self._language_gate = _RateGate(self.high_level_hz) + + @property + def policy(self) -> Any: + return getattr(self.policy_adapter, "policy", self.policy_adapter) + + def set_task(self, task: str) -> None: + with self.state.lock: + if self.state.task != task: + self.state.revision += 1 + self.state.task = task + self.state.log(f"Task: {task}") + + def stop(self) -> None: + self._stop = True + self.state.stop = True + + def run(self, *, max_ticks: int | None = None) -> None: + clock = TickClock(max_rate_hz=self.max_rate_hz) + while not self._stop: + tick = clock.advance() + self._run_tick(tick) + self._flush_logs() + if self.state.stop: + self._stop = True + if max_ticks is not None and tick.index >= max_ticks: + break + self._on_shutdown() + + def step_once(self) -> list[str]: + previous = self.state.tick.index if self.state.tick is not None else 0 + tick = Tick(index=previous + 1, monotonic_seconds=time.monotonic()) + self._run_tick(tick, force_rates=True) + return list(self.state.log_lines) + + def _run_tick(self, tick: Tick, *, force_rates: bool = False) -> None: + self.state.tick = tick + self.state.log_lines = [] + if self.event_collector is not None: + self.event_collector(self.state) + self._handle_action_deadline() + if self.state.stop: + return + self.maybe_update_language_state(force=force_rates) + self.maybe_handle_user_events() + self.maybe_enqueue_action_chunk(force=force_rates) + self.dispatch_action(force=force_rates) + self.state.events.clear() + + def _current_observation(self) -> dict[str, Any] | None: + if self.observation_provider is None: + return None + try: + return self.observation_provider() + except Exception as exc: # noqa: BLE001 + logger.debug("observation_provider failed: %s", exc) + return None + + def maybe_update_language_state(self, *, force: bool = False) -> None: + if self.state.mode != "action" or not self.state.task: + return + if self.state.action_queue: + self._language_gate.rearm() + return + if self.state.tick is None or not self._language_gate.due(self.state.tick, force=force): + return + observation = self._current_observation() + try: + self.policy_adapter.update_language_state(observation, self.state) + except Exception as exc: # noqa: BLE001 + logger.warning("language update failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG)) + self.state.log(f" [warn] language update failed: {type(exc).__name__}: {exc}") + + def maybe_handle_user_events(self) -> None: + if self.state.take_event("user_interjection"): + self._handle_user_interjection() + + def _handle_user_interjection(self) -> None: + text = str(self.state.extra.get("recent_interjection") or "") + if not text: + return + observation = self._current_observation() + self.policy_adapter.handle_interjection(text, observation, self.state) + self.state.extra["recent_interjection"] = None + + def maybe_enqueue_action_chunk(self, *, force: bool = False) -> None: + with self.state.lock: + if self.state.mode != "action" or not self.state.task: + return + if self.state.action_queue: + return + if self.state.tick is None or not self._chunk_gate.due(self.state.tick, force=force): + return + revision = self.state.revision + observation = self._current_observation() + if observation is None: + return + try: + chunk = self.policy_adapter.select_action(observation, self.state) + except Exception as exc: # noqa: BLE001 + logger.warning("select_action failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG)) + self.state.log(f" [warn] select_action failed: {type(exc).__name__}: {exc}") + return + with self.state.lock: + if ( + self.state.revision != revision + or self.state.mode != "action" + or self.state.stop + or self._stop + ): + logger.info("Discarded an action chunk invalidated during inference.") + return + self._enqueue_chunk(chunk) + + def _enqueue_chunk(self, chunk: Any) -> None: + if chunk is None: + return + chunk_iter = chunk[0] if getattr(chunk, "ndim", None) == 3 else chunk + if getattr(chunk_iter, "ndim", None) == 1: + chunk_iter = chunk_iter.unsqueeze(0) + for step in chunk_iter: + self.state.action_queue.append(step.unsqueeze(0) if hasattr(step, "unsqueeze") else step) + try: + self.state.extra["last_chunk_size"] = int(chunk_iter.shape[0]) + except Exception: # noqa: BLE001 + self.state.extra["last_chunk_size"] = len(self.state.action_queue) + + def dispatch_action(self, *, force: bool = False) -> None: + if self.state.mode != "action": + self._last_dispatch_seconds = None + return + if self.state.tick is None or not self._ctrl_gate.due(self.state.tick, force=force): + return + queue = self.state.action_queue + if not queue: + self._last_dispatch_seconds = None + return + now = time.monotonic() + if self._last_dispatch_seconds is None or self.ctrl_hz <= 0: + n_to_pop = 1 + else: + n_to_pop = max(1, min(len(queue), int(round((now - self._last_dispatch_seconds) * self.ctrl_hz)))) + self._last_dispatch_seconds = now + latest = None + for _ in range(n_to_pop): + if not queue: + break + latest = queue.popleft() + self.state.actions_dispatched += 1 + if latest is not None and self.action_executor is not None: + self.action_executor(latest) + + def _handle_action_deadline(self) -> None: + deadline = self.state.action_deadline + if self.state.mode == "action" and deadline is not None and time.monotonic() >= deadline: + self.state.mode = "paused" + self.state.action_deadline = None + self.state.action_queue.clear() + self.state.log("timed action elapsed — paused") + + def _flush_logs(self) -> None: + for line in self.state.log_lines: + print(f"[runtime] {line}", flush=True) + + def _on_shutdown(self) -> None: + self.state.action_queue.clear() + print("[runtime] stopped", flush=True) diff --git a/src/lerobot/runtime/registry.py b/src/lerobot/runtime/registry.py new file mode 100644 index 000000000..98cb7e492 --- /dev/null +++ b/src/lerobot/runtime/registry.py @@ -0,0 +1,39 @@ +# 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. + +"""Lazy mapping from policy types to language-runtime adapters.""" + +from __future__ import annotations + +import importlib +from collections.abc import Callable +from typing import Any + +_ADAPTERS: dict[str, str] = { + "pi052": "lerobot.policies.pi052.inference.pi052_adapter:PI052PolicyAdapter", + "pi05": "lerobot.runtime.adapter:DirectTaskPolicyAdapter", + "molmoact2": "lerobot.runtime.adapter:DirectTaskPolicyAdapter", +} + + +def get_language_adapter_factory(policy_type: str) -> Callable[..., Any]: + """Return the adapter class registered for ``policy_type``.""" + spec = _ADAPTERS.get(policy_type) + if spec is None: + raise ValueError( + f"No language-runtime adapter registered for policy type {policy_type!r}. " + f"Registered: {sorted(_ADAPTERS)}. Add an entry to lerobot.runtime.registry." + ) + module_path, class_name = spec.split(":") + return getattr(importlib.import_module(module_path), class_name) diff --git a/src/lerobot/runtime/sim_robocasa.py b/src/lerobot/runtime/sim_robocasa.py new file mode 100644 index 000000000..d1a5380b1 --- /dev/null +++ b/src/lerobot/runtime/sim_robocasa.py @@ -0,0 +1,406 @@ +# 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. + +"""RoboCasa backend for interactive language-conditioned rollouts. + +It reuses the eval observation/action pipeline while prompts control a persistent selected scene. +""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from datetime import datetime +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from lerobot.utils.io_utils import StreamingVideoWriter +from lerobot.utils.video_annotation import annotate_frame + +logger = logging.getLogger(__name__) + + +def _short_cam_name(cam: str) -> str: + """Human-friendly view label for a RoboCasa camera name.""" + c = cam.replace("robot0_", "") + return { + "agentview_left": "left", + "agentview_right": "right", + "eye_in_hand": "wrist", + }.get(c, c) + + +def _label_panel(img: np.ndarray, label: str) -> np.ndarray: + """Draw a small camera-view label in the bottom-left corner of a panel.""" + try: + import cv2 # noqa: PLC0415 + except ImportError: + return img + y = img.shape[0] - 6 + cv2.putText(img, label, (5, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 0), 3, cv2.LINE_AA) + cv2.putText(img, label, (5, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 0), 1, cv2.LINE_AA) + return img + + +# Two workers avoid broken single-worker EGL rendering; only env 0 is displayed. +_SIM_N_ENVS = 2 + + +def create_sim_env( + *, + task: str, + split: str | None, + obj_registries: list[str], + seed: int | None, + render_size: int = 384, +) -> tuple[Any, dict]: + """Create and reset the vectorized RoboCasa environment before CUDA initializes. + + Two workers keep EGL stable, while only env 0 is driven and displayed. + """ + from lerobot.envs.configs import RoboCasaEnv as RoboCasaEnvConfig # noqa: PLC0415 + + # The policy resizes inputs, so render_size only affects display quality and cost. + env_cfg = RoboCasaEnvConfig( + task=task, + split=split, + obj_registries=list(obj_registries), + observation_height=render_size, + observation_width=render_size, + ) + # Keep one kitchen alive across sequential prompts. + envs = env_cfg.create_envs( + n_envs=_SIM_N_ENVS, + use_async_envs=True, + terminate_on_success=False, + horizon=100_000, + ) + env = envs[next(iter(envs))][0] + logger.info("[sim] resetting RoboCasa scene task=%r split=%r (n_envs=%d)", task, split, _SIM_N_ENVS) + seeds = None if seed is None else [seed + i for i in range(_SIM_N_ENVS)] + obs, _ = env.reset(seed=seeds) + return env, obs + + +def start_mjpeg_server(port: int, get_frame: Callable[[], np.ndarray | None]) -> Any: + """Start an MJPEG server that shows a placeholder until ``get_frame`` returns frames.""" + import io # noqa: PLC0415 + import threading # noqa: PLC0415 + import time # noqa: PLC0415 + from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer # noqa: PLC0415 + + from PIL import Image # noqa: PLC0415 + + _placeholder = Image.new("RGB", (256, 256), (17, 17, 17)) + + class _Handler(BaseHTTPRequestHandler): + def log_message(self, *args): # silence per-request logging + pass + + def do_GET(self): # noqa: N802 + if self.path in ("/", "/index.html"): + self.send_response(200) + self.send_header("Content-Type", "text/html") + self.end_headers() + self.wfile.write( + b"" + b"" + ) + return + if self.path != "/stream": + self.send_response(404) + self.end_headers() + return + self.send_response(200) + self.send_header("Content-Type", "multipart/x-mixed-replace; boundary=frame") + self.end_headers() + try: + while True: + frame = get_frame() + buf = io.BytesIO() + img = Image.fromarray(frame) if frame is not None else _placeholder + img.save(buf, format="JPEG", quality=80) + data = buf.getvalue() + self.wfile.write( + b"--frame\r\nContent-Type: image/jpeg\r\nContent-Length: " + + str(len(data)).encode() + + b"\r\n\r\n" + + data + + b"\r\n" + ) + time.sleep(0.05) + except (BrokenPipeError, ConnectionResetError): + pass + + try: + # Bind all interfaces intentionally so the viewer remains reachable + # through the documented SSH port-forwarding workflow. + server = ThreadingHTTPServer(("0.0.0.0", port), _Handler) # nosec B104 + except OSError as exc: + logger.warning("[sim] could not start live stream on port %d: %s", port, exc) + print(f"[runtime] WARNING: live stream port {port} unavailable ({exc})", flush=True) + return None + threading.Thread(target=server.serve_forever, daemon=True, name="sim-mjpeg").start() + print( + f"[runtime] live view: http://localhost:{port} " + f"(over SSH: ssh -L {port}:localhost:{port} ) — loading until scene is ready", + flush=True, + ) + return server + + +class RoboCasaSimBackend: + """Expose a RoboCasa environment through the runtime observation/action contract. + + The environment must be created before the policy initializes CUDA. + """ + + def __init__( + self, + *, + env: Any, + last_obs: dict, + task: str, + seed: int | None, + device: str, + preprocessor: Any, + postprocessor: Any, + record: bool = True, + output_dir: str | None = None, + view_cams: list[str] | None = None, + ) -> None: + self.env = env + self._last_obs = last_obs + self._scene_task = task + self._view_cams = view_cams or [ + "robot0_agentview_left", + "robot0_eye_in_hand", + "robot0_agentview_right", + ] + self.device = torch.device(device) if isinstance(device, str) else device + self.preprocessor = preprocessor + self.postprocessor = postprocessor + self.seed = seed + self.record = record + self.output_dir = Path(output_dir) if output_dir else Path("outputs/runtime_sim") + + self._video_writer: StreamingVideoWriter | None = None + self._video_path: Path | None = None + self._live_counter = 0 + self._latest_frame: np.ndarray | None = None + self._stream_server: Any = None + self._reset_count = 0 + # Bind these after runtime construction for live annotations. + self._task_getter: Callable[[], str | None] | None = None + self._subtask_getter: Callable[[], str | None] | None = None + self._memory_getter: Callable[[], str | None] | None = None + logger.info("[sim] scene ready — task_description=%r", self._scene_description()) + + def bind_runtime(self, runtime: Any) -> None: + """Wire live task/subtask/memory getters from the runtime state.""" + self._task_getter = lambda: runtime.state.get("task") + self._subtask_getter = lambda: runtime.state.language_context.get("subtask") + self._memory_getter = lambda: (runtime.state.get("language_context") or {}).get("memory") + + def _scene_description(self) -> str: + try: + return str(self.env.get_attr("task_description")[0]) or self._scene_task + except Exception: # noqa: BLE001 + return self._scene_task + + def _current_task(self) -> str: + task = self._task_getter() if self._task_getter else None + return task or self._scene_description() or self._scene_task + + def reset_scene(self) -> None: + """Re-roll the kitchen: reset the env to a fresh scene (new layout/style). + + Uses a new seed each call so ``/reset`` explores different kitchens. + """ + self._reset_count += 1 + n = self.env.num_envs + if self.seed is None: + seeds = None + else: + base = self.seed + self._reset_count * 1000 + seeds = [base + i for i in range(n)] + obs, _ = self.env.reset(seed=seeds) + self._last_obs = obs + logger.info("[sim] scene reset (#%d)", self._reset_count) + + def _env0_obs(self) -> dict: + """Slice env 0 out of the batched vec-env observation (batch of 1).""" + raw = self._last_obs or {} + pixels = raw.get("pixels") + out: dict[str, Any] = {} + if isinstance(pixels, dict): + out["pixels"] = {k: np.asarray(v)[0:1] for k, v in pixels.items()} + agent_pos = raw.get("agent_pos") + if agent_pos is not None: + out["agent_pos"] = np.asarray(agent_pos)[0:1] + return out + + def observation_provider(self) -> dict | None: + from lerobot.envs.utils import preprocess_observation # noqa: PLC0415 + + try: + obs = preprocess_observation(self._env0_obs()) + except Exception as exc: # noqa: BLE001 + logger.warning("[sim] preprocess_observation failed: %s", exc) + return None + # The adapter later replaces this recipe input with its generated subtask. + obs["task"] = [self._current_task()] + if self.preprocessor is not None: + try: + obs = self.preprocessor(obs) + except Exception as exc: # noqa: BLE001 + logger.warning("[sim] preprocessor failed: %s", exc) + return None + return { + k: (v.to(self.device) if isinstance(v, torch.Tensor) else v) + for k, v in obs.items() + if isinstance(k, str) and k.startswith("observation.") + } + + def action_executor(self, action: Any) -> None: + try: + if self.postprocessor is not None: + action = self.postprocessor(action) + if isinstance(action, torch.Tensor): + if action.ndim > 1 and action.shape[0] == 1: + action = action.squeeze(0) + action = action.detach().to("cpu").numpy() + # Tile env 0's action because the extra workers exist only for EGL stability. + action_row = np.asarray(action, dtype=np.float32).reshape(-1) + action_np = np.tile(action_row, (self.env.num_envs, 1)) + obs, _reward, terminated, truncated, _info = self.env.step(action_np) + self._last_obs = obs + self._capture_frame() + # AsyncVectorEnv resets terminated sub-environments automatically. + if bool(np.any(terminated)) or bool(np.any(truncated)): + logger.info("[sim] episode ended — scene auto-reset") + except Exception as exc: # noqa: BLE001 + logger.error("[sim] env.step failed: %s", exc, exc_info=True) + + def _multiview_frame(self) -> np.ndarray | None: + """Label and compose env 0's existing observation views without extra rendering.""" + pixels = (self._last_obs or {}).get("pixels") + if not isinstance(pixels, dict) or not pixels: + return None + panels: list[np.ndarray] = [] + for cam in self._view_cams: + v = pixels.get(cam) + if v is None: + continue + img = np.asarray(v) + if img.ndim == 4: # (n_envs, H, W, C) -> env 0 + img = img[0] + if img.ndim != 3 or img.shape[-1] != 3: + continue + panels.append(_label_panel(np.ascontiguousarray(img.astype(np.uint8)), _short_cam_name(cam))) + if not panels: + return None + h = min(p.shape[0] for p in panels) + panels = [p[:h] for p in panels] + return np.concatenate(panels, axis=1) + + def _capture_frame(self) -> None: + frame = self._multiview_frame() + if frame is None: # fallback to single env.render() + try: + rendered = self.env.call("render")[0] + if isinstance(rendered, np.ndarray) and rendered.ndim == 3: + frame = rendered + except Exception as exc: # noqa: BLE001 + logger.debug("[sim] render failed: %s", exc) + if frame is None: + return + subtask = self._subtask_getter() if self._subtask_getter else None + memory = self._memory_getter() if self._memory_getter else None + annotated = annotate_frame( + frame, + (("Task", self._current_task()), ("Subtask", subtask), ("Memory", memory)), + ) + self._latest_frame = annotated # served by the live MJPEG stream + self._write_live_frame(annotated) + if self.record: + self._write_recording_frame(annotated) + + def _write_recording_frame(self, frame: np.ndarray) -> None: + try: + if self._video_writer is None: + self.output_dir.mkdir(parents=True, exist_ok=True) + stamp = datetime.now().strftime("%Y%m%d_%H%M%S") + self._video_path = self.output_dir / f"sim_{stamp}.mp4" + fps = int((getattr(self.env, "metadata", None) or {}).get("render_fps", 20)) + self._video_writer = StreamingVideoWriter(self._video_path, fps) + self._video_writer.add_frame(frame) + except Exception as exc: # noqa: BLE001 + logger.warning("[sim] video encoding failed: %s", exc) + self.record = False + + def _write_live_frame(self, frame: np.ndarray) -> None: + """Write a rolling latest.png every few frames for live viewing over SSH. + + Open ``{output_dir}/latest.png`` in an editor/viewer and refresh to watch + the rollout in near-real-time without a GUI window. Written atomically + (temp + replace) so a reader never sees a half-written file. + """ + self._live_counter += 1 + if self._live_counter % 3 != 0: + return + try: + import os # noqa: PLC0415 + + from PIL import Image # noqa: PLC0415 + + self.output_dir.mkdir(parents=True, exist_ok=True) + tmp = self.output_dir / ".latest.tmp.png" + Image.fromarray(frame).save(tmp) + os.replace(tmp, self.output_dir / "latest.png") + except Exception as exc: # noqa: BLE001 + logger.debug("[sim] live frame write failed: %s", exc) + + def _flush_video(self) -> None: + if self._video_writer is None: + return + writer = self._video_writer + self._video_writer = None + try: + writer.close() + logger.info("[sim] wrote video (%d frames) to %s", writer.frames_written, self._video_path) + print(f"[runtime] sim video saved to {self._video_path}", flush=True) + except Exception as exc: # noqa: BLE001 + logger.warning("[sim] video close failed: %s", exc) + + def attach_stream_server(self, server: Any) -> None: + """Attach an already-running MJPEG server so disconnect() can stop it.""" + self._stream_server = server + + def disconnect(self) -> None: + """Match the robot backend's cleanup contract.""" + if self._stream_server is not None: + try: + self._stream_server.shutdown() + except Exception as exc: # noqa: BLE001 + logger.debug("[sim] stream server shutdown raised %s", exc) + self._flush_video() + try: + self.env.close() + except Exception as exc: # noqa: BLE001 + logger.debug("[sim] env.close raised %s", exc) diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index c4ed35145..b219f9bb1 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -94,6 +94,19 @@ from lerobot.utils.utils import ( init_logging, inside_slurm, ) +from lerobot.utils.video_annotation import annotate_frame + + +def _annotate_eval_frames(frames: np.ndarray, task: str | None, subtask: str | None) -> np.ndarray: + """Overlay the high-level task and predicted subtask onto rendered frames. + + ``frames`` is ``(n_envs, H, W, C)`` uint8. Best-effort: if OpenCV isn't + available the frames are returned unchanged so eval never fails over a + visualization concern. + """ + if frames.ndim != 4 or frames.shape[-1] != 3: + return frames + return np.stack([annotate_frame(frame, (("Task", task), ("Subtask", subtask))) for frame in frames]) def _env_features_to_dataset_features(env_features: dict) -> dict: @@ -477,11 +490,36 @@ def eval_policy( return n_to_render_now = min(max_episodes_rendered - n_episodes_rendered, env.num_envs) if isinstance(env, gym.vector.SyncVectorEnv): - ep_frames.append(np.stack([env.envs[i].render() for i in range(n_to_render_now)])) # noqa: B023 + frames = np.stack([env.envs[i].render() for i in range(n_to_render_now)]) # noqa: B023 elif hasattr(env, "call"): # Here we must render all frames and discard any we don't need. # Covers AsyncVectorEnv and _LazyAsyncVectorEnv (which wraps one). - ep_frames.append(np.stack(env.call("render")[:n_to_render_now])) + frames = np.stack(env.call("render")[:n_to_render_now]) + else: + return + + # Overlay the high-level task and (for hierarchical policies like + # pi052) the predicted low-level subtask onto each frame. Both are + # best-effort: missing values just skip that line. + try: + tasks = list(env.call("task_description")) + except (AttributeError, NotImplementedError): + try: + tasks = list(env.call("task")) + except (AttributeError, NotImplementedError): + tasks = None + subtasks = getattr(policy, "last_subtasks", None) + annotated = [] + for i in range(frames.shape[0]): + subtask_i = subtasks[i] if subtasks is not None and i < len(subtasks) else None + annotated.append( + _annotate_eval_frames( + frames[i : i + 1], + tasks[i] if tasks is not None and i < len(tasks) else None, + subtask_i, + )[0] + ) + ep_frames.append(np.stack(annotated)) if max_episodes_rendered > 0: video_paths: list[str] = [] diff --git a/src/lerobot/scripts/lerobot_rollout.py b/src/lerobot/scripts/lerobot_rollout.py index 879070721..06023e6e8 100644 --- a/src/lerobot/scripts/lerobot_rollout.py +++ b/src/lerobot/scripts/lerobot_rollout.py @@ -151,6 +151,7 @@ Usage examples """ import logging +import sys from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401 @@ -241,10 +242,69 @@ def rollout(cfg: RolloutConfig): logger.info("Rollout finished") -def main(): - """CLI entry point for ``lerobot-rollout``.""" +_LANGUAGE_RUNTIME_FLAGS = { + "--language", + "--no_robot", + "--sim", + "--direct_subtask", + "--sim.direct_subtask", + "--disable_memory", + "--fp8", +} +_LANGUAGE_RUNTIME_PREFIXES = ( + "--sim.", + "--chunk_hz", + "--ctrl_hz", + "--high_level_hz", + "--subtask_chunks_per_gen", + "--text_min_new_tokens", + "--text_temperature", + "--text_top_p", +) + + +def _uses_language_runtime(argv: list[str]) -> bool: + """Return whether *argv* selects the interactive language runtime. + + ``--language`` is the explicit selector for real-robot runs whose other + options overlap with the standard rollout CLI. Language-only options also + select it automatically, which keeps the former language-runtime examples + working after replacing their command name with ``lerobot-rollout``. + """ + return any( + arg.split("=", 1)[0] in _LANGUAGE_RUNTIME_FLAGS or arg.startswith(_LANGUAGE_RUNTIME_PREFIXES) + for arg in argv + ) + + +def main(argv: list[str] | None = None): + """CLI entry point for ``lerobot-rollout``. + + Standard policy deployment continues through :class:`RolloutConfig`. + Interactive language-conditioned and RoboCasa runs share this entry point + and are selected with ``--language`` or any language-runtime-only option. + """ register_third_party_plugins() - rollout() + cli_args = list(sys.argv[1:] if argv is None else argv) + if _uses_language_runtime(cli_args): + from lerobot.runtime.cli import run as run_language_runtime + + # ``--language`` is a dispatcher flag, not part of the runtime's own + # argparse surface. All other arguments pass through unchanged. + runtime_args = [arg for arg in cli_args if arg != "--language"] + return run_language_runtime(runtime_args, prog="lerobot-rollout") + + if argv is None: + return rollout() + + # draccus reads sys.argv. Supporting an explicit argv keeps this entry + # point easy to smoke-test and mirrors the language-runtime branch above. + previous_argv = sys.argv + try: + sys.argv = [previous_argv[0], *cli_args] + return rollout() + finally: + sys.argv = previous_argv if __name__ == "__main__": diff --git a/src/lerobot/utils/io_utils.py b/src/lerobot/utils/io_utils.py index e037b412c..7deeb2d70 100644 --- a/src/lerobot/utils/io_utils.py +++ b/src/lerobot/utils/io_utils.py @@ -23,6 +23,46 @@ logger = logging.getLogger(__name__) JsonLike = str | int | float | bool | None | list["JsonLike"] | dict[str, "JsonLike"] | tuple["JsonLike", ...] +class StreamingVideoWriter: + """Incrementally encode RGB frames to an MP4 without retaining them in memory.""" + + def __init__(self, video_path: str | Path, fps: int) -> None: + from .import_utils import require_package + + require_package("av", extra="av-dep") + import av + + self._av = av + self._container = av.open(str(video_path), mode="w") + self._stream = self._container.add_stream("libx264", rate=fps) + self._shape: tuple[int, int] | None = None + self.frames_written = 0 + + def add_frame(self, frame_array) -> None: + orig_height, orig_width = frame_array.shape[:2] + height = orig_height - orig_height % 2 + width = orig_width - orig_width % 2 + if self._shape is None: + self._shape = (height, width) + self._stream.width = width + self._stream.height = height + self._stream.pix_fmt = "yuv420p" + elif self._shape != (height, width): + raise ValueError(f"Video frame shape changed from {self._shape} to {(height, width)}") + frame = self._av.VideoFrame.from_ndarray(frame_array[:height, :width], format="rgb24") + for packet in self._stream.encode(frame): + self._container.mux(packet) + self.frames_written += 1 + + def close(self) -> None: + if self._container is None: + return + for packet in self._stream.encode(): + self._container.mux(packet) + self._container.close() + self._container = None + + def load_json(fpath: Path) -> Any: """Load data from a JSON file. @@ -58,36 +98,12 @@ def write_video(video_path: str | Path, stacked_frames: list, fps: int) -> None: stacked_frames: List of HWC uint8 numpy arrays (RGB). fps: Frames per second for the output video. """ - from .import_utils import require_package - - require_package("av", extra="av-dep") - import av - - with av.open(str(video_path), mode="w") as container: - orig_height, orig_width = stacked_frames[0].shape[:2] - # yuv420p requires even dimensions; crop by one pixel if needed - height = orig_height if orig_height % 2 == 0 else orig_height - 1 - width = orig_width if orig_width % 2 == 0 else orig_width - 1 - if height != orig_height or width != orig_width: - logger.warning( - "Frame dimensions %dx%d are not even; cropping to %dx%d for yuv420p compatibility.", - orig_width, - orig_height, - width, - height, - ) - stream = container.add_stream("libx264", rate=fps) - stream.width = width - stream.height = height - stream.pix_fmt = "yuv420p" + writer = StreamingVideoWriter(video_path, fps) + try: for frame_array in stacked_frames: - if height != orig_height or width != orig_width: - frame_array = frame_array[:height, :width] - frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24") - for packet in stream.encode(frame): - container.mux(packet) - for packet in stream.encode(): - container.mux(packet) + writer.add_frame(frame_array) + finally: + writer.close() def deserialize_json_into_object[T: JsonLike](fpath: Path, obj: T) -> T: diff --git a/src/lerobot/utils/rerun_visualization.py b/src/lerobot/utils/rerun_visualization.py index 46f2c0b4b..5dd3eb174 100644 --- a/src/lerobot/utils/rerun_visualization.py +++ b/src/lerobot/utils/rerun_visualization.py @@ -38,7 +38,10 @@ def _is_scalar(x): def init_rerun( - session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None + session_name: str = "lerobot_control_loop", + ip: str | None = None, + port: int | None = None, + web_port: int | None = None, ) -> None: """ Initializes the Rerun SDK for visualizing the control loop. @@ -47,6 +50,7 @@ def init_rerun( session_name: Name of the Rerun session. ip: Optional IP for connecting to a Rerun server. port: Optional port for connecting to a Rerun server. + web_port: Serve a headless web viewer on this port, using ``port`` for gRPC. """ require_package("rerun-sdk", extra="viz", import_name="rerun") @@ -60,6 +64,10 @@ def init_rerun( memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%") if ip and port: rr.connect_grpc(url=f"rerun+http://{ip}:{port}/proxy") + elif web_port is not None: + grpc_port = port or 9876 + url = rr.serve_grpc(grpc_port=grpc_port) + rr.serve_web_viewer(web_port=web_port, open_browser=False, connect_to=url) else: rr.spawn(memory_limit=memory_limit) 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/envs/test_robocasa_ordering.py b/tests/envs/test_robocasa_ordering.py new file mode 100644 index 000000000..c4fbe9041 --- /dev/null +++ b/tests/envs/test_robocasa_ordering.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np + +from lerobot.envs.robocasa import RoboCasaEnv, convert_action + + +def test_robocasa_action_uses_openpi_checkpoint_order(): + action = np.arange(12, dtype=np.float32) + + converted = convert_action(action) + + np.testing.assert_array_equal(converted["action.end_effector_position"], [0, 1, 2]) + np.testing.assert_array_equal(converted["action.end_effector_rotation"], [3, 4, 5]) + np.testing.assert_array_equal(converted["action.gripper_close"], [6]) + np.testing.assert_array_equal(converted["action.base_motion"], [7, 8, 9, 10]) + np.testing.assert_array_equal(converted["action.control_mode"], [11]) + + +def test_robocasa_state_uses_openpi_checkpoint_order(): + env = object.__new__(RoboCasaEnv) + env.obs_type = "pixels_agent_pos" + env.camera_name = [] + raw_observation = { + "state.end_effector_position_relative": np.arange(0, 3), + "state.end_effector_rotation_relative": np.arange(3, 7), + "state.base_position": np.arange(7, 10), + "state.base_rotation": np.arange(10, 14), + "state.gripper_qpos": np.arange(14, 16), + } + + observation = env._format_raw_obs(raw_observation) + + np.testing.assert_array_equal(observation["agent_pos"], np.arange(16, dtype=np.float32)) diff --git a/tests/runtime/test_adapter.py b/tests/runtime/test_adapter.py new file mode 100644 index 000000000..470ef13ab --- /dev/null +++ b/tests/runtime/test_adapter.py @@ -0,0 +1,105 @@ +# 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. + +from lerobot.runtime import RuntimeState +from lerobot.runtime.adapter import ( + BaseLanguageAdapter, + DirectTaskPolicyAdapter, + GenerationConfig, +) + + +class ScriptedAdapter(BaseLanguageAdapter): + """Base adapter whose text generation returns queued strings per kind.""" + + def __init__(self, scripts, gen=None): + super().__init__(policy=object(), gen=gen) + self.scripts = {k: list(v) for k, v in scripts.items()} + self.calls = [] + + def select_action(self, observation, state): + return None + + def generate_text(self, kind, observation, state, user_text=None): + self.calls.append(kind) + queue = self.scripts.get(kind, []) + return queue.pop(0) if queue else "" + + +def test_cascade_sets_subtask_then_memory(): + adapter = ScriptedAdapter({"subtask": ["pick the red cup"], "memory": ["the cup is grasped"]}) + state = RuntimeState(task="clean") + + adapter.update_language_state(None, state) + + assert state.language_context["subtask"] == "pick the red cup" + assert state.language_context["memory"] == "the cup is grasped" + assert adapter.calls == ["subtask", "memory"] + + +def test_nonempty_generation_is_used_verbatim(): + adapter = ScriptedAdapter({"subtask": [":::: ::"], "memory": ["memory"]}) + state = RuntimeState(task="clean") + + adapter.update_language_state(None, state) + + assert state.language_context["subtask"] == ":::: ::" + assert state.language_context["memory"] == "memory" + assert adapter.calls == ["subtask", "memory"] + + +def test_throttle_regenerates_every_n_chunks(): + adapter = ScriptedAdapter( + { + "subtask": ["pick the first cup", "pick the second cup"], + "memory": ["memory one two three", "memory four five six"], + }, + gen=GenerationConfig(chunks_per_regen=2), + ) + state = RuntimeState(task="clean") + + adapter.update_language_state(None, state) # generates + assert state.language_context["subtask"] == "pick the first cup" + adapter.update_language_state(None, state) # throttled — no generation + assert state.language_context["subtask"] == "pick the first cup" + adapter.update_language_state(None, state) # generates again + assert state.language_context["subtask"] == "pick the second cup" + + +def test_handle_interjection_sets_plan_and_strips_say(): + adapter = ScriptedAdapter({"interjection": ["turn to the left now heading left"]}) + state = RuntimeState(task="clean") + + adapter.handle_interjection("turn", None, state) + + assert state.language_context["plan"] == "turn to the left now" + + +def test_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..33d0414a1 --- /dev/null +++ b/tests/runtime/test_cli.py @@ -0,0 +1,75 @@ +# 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. + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +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_parse_args_rejects_removed_dataset_replay_flags(): + with pytest.raises(SystemExit): + _parse_args(["--policy.path=checkpoint", "--dataset.repo_id=dataset"]) + + +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_language_runtime.py b/tests/runtime/test_language_runtime.py new file mode 100644 index 000000000..38ffc77c5 --- /dev/null +++ b/tests/runtime/test_language_runtime.py @@ -0,0 +1,100 @@ +# 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. + +import threading +import time + +from lerobot.runtime import LanguageConditionedRuntime, Tick + + +class FakeAdapter: + def __init__(self): + self.updated = False + self.interjections = [] + + def select_action(self, observation, state): + assert observation == {"observation.state": 1} + assert state.task == "clean" + return ["a0", "a1"] + + def update_language_state(self, observation, state): + self.updated = True + state.set_context("subtask", "pick cup", label="subtask") + + def handle_interjection(self, user_text, observation, state): + self.interjections.append(user_text) + state.set_context("plan", "new plan", label="plan") + + +def test_runtime_tick_updates_language_enqueues_and_dispatches_action(): + adapter = FakeAdapter() + executed = [] + runtime = LanguageConditionedRuntime( + policy_adapter=adapter, + observation_provider=lambda: {"observation.state": 1}, + action_executor=executed.append, + ) + runtime.set_task("clean") + + logs = runtime.step_once() + + assert adapter.updated + assert runtime.state.language_context["subtask"] == "pick cup" + assert executed == ["a0"] + assert list(runtime.state.action_queue) == ["a1"] + assert " subtask: pick cup" in logs + + +def test_runtime_handles_user_interjection(): + adapter = FakeAdapter() + runtime = LanguageConditionedRuntime( + policy_adapter=adapter, + observation_provider=lambda: {"observation.state": 1}, + ) + runtime.set_task("clean") + runtime.state.extra["recent_interjection"] = "please say ok" + runtime.state.emit("user_interjection") + + runtime.step_once() + + assert "please say ok" in adapter.interjections + assert runtime.state.language_context["plan"] == "new plan" + + +def test_prompt_change_discards_in_flight_action_chunk(): + started = threading.Event() + release = threading.Event() + + class BlockingAdapter(FakeAdapter): + def select_action(self, observation, state): + started.set() + assert release.wait(timeout=2) + return ["stale"] + + runtime = LanguageConditionedRuntime( + policy_adapter=BlockingAdapter(), + observation_provider=lambda: {"observation.state": 1}, + ) + runtime.set_task("old task") + runtime.state.tick = Tick(index=1, monotonic_seconds=time.monotonic()) + inference = threading.Thread(target=runtime.maybe_enqueue_action_chunk, kwargs={"force": True}) + inference.start() + assert started.wait(timeout=2) + + runtime.set_task("new task") + release.set() + inference.join(timeout=2) + + assert not inference.is_alive() + assert list(runtime.state.action_queue) == [] diff --git a/tests/runtime/test_sim_robocasa.py b/tests/runtime/test_sim_robocasa.py new file mode 100644 index 000000000..5f60bfd6e --- /dev/null +++ b/tests/runtime/test_sim_robocasa.py @@ -0,0 +1,82 @@ +# 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. + +import sys +from types import SimpleNamespace + +import numpy as np + +from lerobot.runtime.sim_robocasa import RoboCasaSimBackend +from lerobot.utils.video_annotation import annotate_frame + + +def test_overlay_draws_each_label_once(monkeypatch): + put_text_calls = [] + rectangle_calls = [] + + def put_text(image, text, origin, font, scale, color, thickness, line_type): + put_text_calls.append((text, color, thickness)) + return image + + def rectangle(image, start, end, color, thickness): + rectangle_calls.append((start, end, color, thickness)) + return image + + def add_weighted(src1, alpha, src2, beta, gamma, *, dst): + dst[:] = src1 * alpha + src2 * beta + gamma + return dst + + fake_cv2 = SimpleNamespace( + FONT_HERSHEY_SIMPLEX=0, + LINE_AA=16, + getTextSize=lambda text, font, scale, thickness: ((len(text) * 7, 10), 0), + putText=put_text, + rectangle=rectangle, + addWeighted=add_weighted, + ) + monkeypatch.setitem(sys.modules, "cv2", fake_cv2) + + frame = np.full((120, 480, 3), 200, dtype=np.uint8) + 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", + "Subtask: reach for the handle", + ] + assert all(color == (255, 255, 255) and thickness == 1 for _, color, thickness in put_text_calls) + assert len(rectangle_calls) == 1 + assert not np.shares_memory(annotated, frame) + + +def test_capture_updates_live_frame_when_recording_is_disabled(monkeypatch): + backend = object.__new__(RoboCasaSimBackend) + frame = np.full((8, 8, 3), 42, dtype=np.uint8) + written = [] + backend.record = False + backend.runtime_state = None + backend._multiview_frame = lambda: frame + backend._current_task = lambda: "task" + backend._subtask_getter = None + backend._memory_getter = None + backend._latest_frame = None + backend._write_live_frame = written.append + monkeypatch.setattr("lerobot.runtime.sim_robocasa.annotate_frame", lambda image, labels: image) + + backend._capture_frame() + + assert backend._latest_frame is frame + assert written == [frame]