feat(runtime): add interactive language rollouts

This commit is contained in:
Pepijn
2026-07-28 10:05:27 +02:00
parent 3f093d8927
commit f2b90e3ad6
21 changed files with 3074 additions and 49 deletions
+156
View File
@@ -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=<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 ## 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. 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.
+47 -1
View File
@@ -1,6 +1,6 @@
# Policy Deployment (lerobot-rollout) # 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 ## 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 ## Inference Backends
Select a backend with `--inference.type=<name>`. All strategies work with both backends. Select a backend with `--inference.type=<name>`. All strategies work with both backends.
+9 -1
View File
@@ -560,7 +560,13 @@ class RoboCasaEnv(EnvConfig):
kwargs["split"] = self.split kwargs["split"] = self.split
return kwargs 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 from .robocasa import create_robocasa_envs
if self.task is None: if self.task is None:
@@ -574,6 +580,8 @@ class RoboCasaEnv(EnvConfig):
env_cls=env_cls, env_cls=env_cls,
episode_length=self.episode_length, episode_length=self.episode_length,
obj_registries=tuple(self.obj_registries), obj_registries=tuple(self.obj_registries),
terminate_on_success=terminate_on_success,
horizon=horizon,
) )
+33 -11
View File
@@ -33,8 +33,8 @@ logger = logging.getLogger(__name__)
# Dimensions for the flat action/state vectors used by the LeRobot wrapper. # Dimensions for the flat action/state vectors used by the LeRobot wrapper.
# These correspond to the PandaOmron robot in RoboCasa365. # 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) OBS_STATE_DIM = 16 # ee_pos_rel(3) + ee_quat_rel(4) + base_pos(3) + base_quat(4) + gripper_qpos(2)
ACTION_DIM = 12 # base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1) ACTION_DIM = 12 # ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1)
ACTION_LOW = -1.0 ACTION_LOW = -1.0
ACTION_HIGH = 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]: def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
"""Split a flat (12,) action vector into a RoboCasa action dict. """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 { return {
"action.base_motion": flat_action[0:4], "action.end_effector_position": flat_action[0:3],
"action.control_mode": flat_action[4:5], "action.end_effector_rotation": flat_action[3:6],
"action.end_effector_position": flat_action[5:8], "action.gripper_close": flat_action[6:7],
"action.end_effector_rotation": flat_action[8:11], "action.base_motion": flat_action[7:11],
"action.gripper_close": flat_action[11:12], "action.control_mode": flat_action[11:12],
} }
@@ -136,9 +137,16 @@ class RoboCasaEnv(gym.Env):
episode_length: int | None = None, episode_length: int | None = None,
obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES, obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES,
episode_index: int = 0, episode_index: int = 0,
terminate_on_success: bool = True,
horizon: int | None = None,
): ):
super().__init__() super().__init__()
self.task = task 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.obs_type = obs_type
self.render_mode = render_mode self.render_mode = render_mode
self.observation_width = observation_width self.observation_width = observation_width
@@ -210,12 +218,16 @@ class RoboCasaEnv(gym.Env):
# (only None/"all"/"pretrain"/"target" are valid). Always pass a # (only None/"all"/"pretrain"/"target" are valid). Always pass a
# valid value so we don't hit that default. Extra kwargs are # valid value so we don't hit that default. Extra kwargs are
# forwarded to the underlying kitchen env via create_env/robosuite.make. # 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( self._env = RoboCasaGymEnv(
env_name=self.task, env_name=self.task,
camera_widths=self.observation_width, camera_widths=self.observation_width,
camera_heights=self.observation_height, camera_heights=self.observation_height,
split=self.split if self.split is not None else "all", split=self.split if self.split is not None else "all",
obj_registries=self.obj_registries, obj_registries=self.obj_registries,
**extra_kwargs,
) )
ep_meta = self._env.env.get_ep_meta() ep_meta = self._env.env.get_ep_meta()
@@ -230,12 +242,14 @@ class RoboCasaEnv(gym.Env):
return {"pixels": images} return {"pixels": images}
# `state.*` keys come from PandaOmronKeyConverter inside the wrapper. # `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( 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_position_relative", np.zeros(3)),
raw_obs.get("state.end_effector_rotation_relative", np.zeros(4)), 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)), raw_obs.get("state.gripper_qpos", np.zeros(2)),
], ],
axis=-1, axis=-1,
@@ -280,7 +294,7 @@ class RoboCasaEnv(gym.Env):
raw_obs, reward, done, truncated, info = self._env.step(action_dict) raw_obs, reward, done, truncated, info = self._env.step(action_dict)
is_success = bool(info.get("success", False)) 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}) info.update({"task": self.task, "done": done, "is_success": is_success})
observation = self._format_raw_obs(raw_obs) observation = self._format_raw_obs(raw_obs)
@@ -313,6 +327,8 @@ def _make_env_fns(
split: str | None, split: str | None,
episode_length: int | None, episode_length: int | None,
obj_registries: Sequence[str], obj_registries: Sequence[str],
terminate_on_success: bool = True,
horizon: int | None = None,
) -> list[Callable[[], RoboCasaEnv]]: ) -> list[Callable[[], RoboCasaEnv]]:
"""Build n_envs factory callables for a single task. """Build n_envs factory callables for a single task.
@@ -335,6 +351,8 @@ def _make_env_fns(
episode_length=episode_length, episode_length=episode_length,
obj_registries=obj_registries, obj_registries=obj_registries,
episode_index=episode_index, episode_index=episode_index,
terminate_on_success=terminate_on_success,
horizon=horizon,
) )
return [partial(_make_env, i) for i in range(n_envs)] 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, env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
episode_length: int | None = None, episode_length: int | None = None,
obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES, obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES,
terminate_on_success: bool = True,
horizon: int | None = None,
) -> dict[str, dict[int, Any]]: ) -> dict[str, dict[int, Any]]:
"""Create vectorized RoboCasa365 environments with a consistent return shape. """Create vectorized RoboCasa365 environments with a consistent return shape.
@@ -409,6 +429,8 @@ def create_robocasa_envs(
split=split, split=split,
episode_length=episode_length, episode_length=episode_length,
obj_registries=obj_registries, obj_registries=obj_registries,
terminate_on_success=terminate_on_success,
horizon=horizon,
) )
if is_async: if is_async:
+3 -1
View File
@@ -21,6 +21,8 @@ from lerobot.utils.import_utils import make_device_from_device_class
from .config import RobotConfig from .config import RobotConfig
from .robot import Robot from .robot import Robot
logger = logging.getLogger(__name__)
def make_robot_from_config(config: RobotConfig) -> Robot: def make_robot_from_config(config: RobotConfig) -> Robot:
# TODO(Steven): Consider just using the make_device_from_device_class for all types # 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: if warnings_dict:
logging.warning( logger.warning(
"Relative goal position magnitude had to be clamped to be safe.\n" "Relative goal position magnitude had to be clamped to be safe.\n"
f"{pformat(warnings_dict, indent=4)}" f"{pformat(warnings_dict, indent=4)}"
) )
+38
View File
@@ -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",
]
+165
View File
@@ -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 ``<say>`` 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 <say>speech</say>`` into ``(plan, speech)``."""
if not text:
return "", ""
match = _SAY_RE.search(text)
if not match:
return text.strip(), ""
speech = match.group(1).strip().strip('"').strip("'")
plan = (text[: match.start()] + text[match.end() :]).strip()
return plan, speech
File diff suppressed because it is too large Load Diff
+349
View File
@@ -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)
+39
View File
@@ -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)
+406
View File
@@ -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"<html><body style='margin:0;background:#111;text-align:center'>"
b"<img src='/stream' style='max-width:100vw;max-height:100vh;"
b"image-rendering:pixelated'></body></html>"
)
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} <host>) — 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)
+40 -2
View File
@@ -94,6 +94,19 @@ from lerobot.utils.utils import (
init_logging, init_logging,
inside_slurm, 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: def _env_features_to_dataset_features(env_features: dict) -> dict:
@@ -477,11 +490,36 @@ def eval_policy(
return return
n_to_render_now = min(max_episodes_rendered - n_episodes_rendered, env.num_envs) n_to_render_now = min(max_episodes_rendered - n_episodes_rendered, env.num_envs)
if isinstance(env, gym.vector.SyncVectorEnv): 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"): elif hasattr(env, "call"):
# Here we must render all frames and discard any we don't need. # Here we must render all frames and discard any we don't need.
# Covers AsyncVectorEnv and _LazyAsyncVectorEnv (which wraps one). # 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: if max_episodes_rendered > 0:
video_paths: list[str] = [] video_paths: list[str] = []
+63 -3
View File
@@ -151,6 +151,7 @@ Usage examples
""" """
import logging import logging
import sys
from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401
from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401 from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401
@@ -241,10 +242,69 @@ def rollout(cfg: RolloutConfig):
logger.info("Rollout finished") logger.info("Rollout finished")
def main(): _LANGUAGE_RUNTIME_FLAGS = {
"""CLI entry point for ``lerobot-rollout``.""" "--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() 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__": if __name__ == "__main__":
+45 -29
View File
@@ -23,6 +23,46 @@ logger = logging.getLogger(__name__)
JsonLike = str | int | float | bool | None | list["JsonLike"] | dict[str, "JsonLike"] | tuple["JsonLike", ...] 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: def load_json(fpath: Path) -> Any:
"""Load data from a JSON file. """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). stacked_frames: List of HWC uint8 numpy arrays (RGB).
fps: Frames per second for the output video. fps: Frames per second for the output video.
""" """
from .import_utils import require_package writer = StreamingVideoWriter(video_path, fps)
try:
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"
for frame_array in stacked_frames: for frame_array in stacked_frames:
if height != orig_height or width != orig_width: writer.add_frame(frame_array)
frame_array = frame_array[:height, :width] finally:
frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24") writer.close()
for packet in stream.encode(frame):
container.mux(packet)
for packet in stream.encode():
container.mux(packet)
def deserialize_json_into_object[T: JsonLike](fpath: Path, obj: T) -> T: def deserialize_json_into_object[T: JsonLike](fpath: Path, obj: T) -> T:
+9 -1
View File
@@ -38,7 +38,10 @@ def _is_scalar(x):
def init_rerun( 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: ) -> None:
""" """
Initializes the Rerun SDK for visualizing the control loop. Initializes the Rerun SDK for visualizing the control loop.
@@ -47,6 +50,7 @@ def init_rerun(
session_name: Name of the Rerun session. session_name: Name of the Rerun session.
ip: Optional IP for connecting to a Rerun server. ip: Optional IP for connecting to a Rerun server.
port: Optional port 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") 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%") memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%")
if ip and port: if ip and port:
rr.connect_grpc(url=f"rerun+http://{ip}:{port}/proxy") 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: else:
rr.spawn(memory_limit=memory_limit) rr.spawn(memory_limit=memory_limit)
+71
View File
@@ -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
+48
View File
@@ -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))
+105
View File
@@ -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 <say>heading left</say>"]})
state = RuntimeState(task="clean")
adapter.handle_interjection("turn", None, state)
assert state.language_context["plan"] == "turn to the left now"
def test_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
+75
View File
@@ -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})
+100
View File
@@ -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) == []
+82
View File
@@ -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]