diff --git a/src/lerobot/envs/configs.py b/src/lerobot/envs/configs.py index 3624357e2..77a01ba32 100644 --- a/src/lerobot/envs/configs.py +++ b/src/lerobot/envs/configs.py @@ -556,7 +556,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: @@ -570,6 +576,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 5d8932f03..559ad2e01 100644 --- a/src/lerobot/envs/robocasa.py +++ b/src/lerobot/envs/robocasa.py @@ -137,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 @@ -211,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() @@ -283,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) @@ -316,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. @@ -338,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)] @@ -351,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. @@ -412,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/policies/pi05/pi05_adapter.py b/src/lerobot/policies/pi05/pi05_adapter.py new file mode 100644 index 000000000..757a52f1b --- /dev/null +++ b/src/lerobot/policies/pi05/pi05_adapter.py @@ -0,0 +1,50 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""PI05 adapter for the language-conditioned runtime. + +PI05 is a flat VLA: it conditions the action expert directly on the task text, +which its preprocessor tokenizes into ``observation.language.tokens``. It has no +subtask/memory generation head, so the runtime simply predicts an action chunk +from the already-tokenized observation. Text generation is unsupported — run +with ``--sim.direct_subtask`` so the runtime doesn't attempt subtask/memory +generation (what you type becomes the task the preprocessor tokenizes). +""" + +from __future__ import annotations + +from typing import Any + +from lerobot.runtime import RuntimeState +from lerobot.runtime.adapter import BaseLanguageAdapter + + +class PI05PolicyAdapter(BaseLanguageAdapter): + """Runtime bridge for flat PI05 policies (direct task-text conditioning).""" + + def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: + # The task text was tokenized into observation.language.* by the policy + # preprocessor (fed the current task by the observation provider), so we + # just predict the action chunk from it. + return self.policy.predict_action_chunk(observation) + + def generate_text( + self, + kind: str, + observation: dict[str, Any] | None, + state: RuntimeState, + user_text: str | None = None, + ) -> str: + # PI05 has no text-generation head; direct-subtask mode skips this path. + return "" diff --git a/src/lerobot/policies/pi052/inference/pi052_adapter.py b/src/lerobot/policies/pi052/inference/pi052_adapter.py index 4d70679ab..1cbc0bf97 100644 --- a/src/lerobot/policies/pi052/inference/pi052_adapter.py +++ b/src/lerobot/policies/pi052/inference/pi052_adapter.py @@ -37,14 +37,31 @@ class PI052PolicyAdapter(BaseLanguageAdapter): """Runtime bridge for PI052 policies.""" def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: + import torch # noqa: PLC0415 + + from lerobot.utils.constants import ( # noqa: PLC0415 + OBS_LANGUAGE_ATTENTION_MASK, + OBS_LANGUAGE_TOKENS, + OBS_STATE, + ) + subtask = state.language_context.get("subtask") or state.task or "" + # Condition the action expert on subtask + discretized state, matching + # training and lerobot-eval's low-level prompt ("{subtask}, State: {..};"). + # Without the state the action expert is off-distribution. + content = subtask + obs_state = observation.get(OBS_STATE) + if isinstance(obs_state, torch.Tensor) and obs_state.numel() > 0: + from lerobot.policies.pi052.text_processor_pi052 import discretize_state_str # noqa: PLC0415 + + state_row = obs_state[0] if obs_state.ndim > 1 else obs_state + content = f"{subtask}, State: {discretize_state_str(state_row)};" + text_batch = _build_text_batch( self.policy, - [{"role": "user", "content": subtask}], + [{"role": "user", "content": content}], add_generation_prompt=False, ) - from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS # noqa: PLC0415 - batch = dict(observation) batch[OBS_LANGUAGE_TOKENS] = text_batch["lang_tokens"] batch[OBS_LANGUAGE_ATTENTION_MASK] = text_batch["lang_masks"] diff --git a/src/lerobot/runtime/adapter.py b/src/lerobot/runtime/adapter.py index fd9f10070..eb9c4e79e 100644 --- a/src/lerobot/runtime/adapter.py +++ b/src/lerobot/runtime/adapter.py @@ -51,6 +51,8 @@ class GenerationConfig: 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 @@ -127,6 +129,10 @@ class BaseLanguageAdapter(ABC): Override for a policy with a different language hierarchy. """ + if not self.gen.enable_subtask: + # Direct-subtask mode: the operator supplies the subtask; don't + # generate (and thus don't overwrite) it. + return subtask = self._generate_filtered("subtask", observation, state) if subtask is None: return @@ -137,6 +143,8 @@ class BaseLanguageAdapter(ABC): 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") diff --git a/src/lerobot/runtime/cli.py b/src/lerobot/runtime/cli.py index 413364fbe..74b149e58 100644 --- a/src/lerobot/runtime/cli.py +++ b/src/lerobot/runtime/cli.py @@ -238,6 +238,97 @@ def _parse_args(argv: list[str] | None = None, *, prog: str | None = None) -> ar "wrong robot, robot not at home pose)." ), ) + # --- RoboCasa simulation mode args ------------------------------- + # Setting ``--sim`` flips the runtime into simulation mode: instead of + # a real robot it drives a single RoboCasa mujoco scene, feeding the + # eval observation/action pipeline. The operator still types prompts + # (/action ) that the policy executes inside the chosen scene. + # Mutually exclusive with ``--robot.type``. + 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, @@ -257,6 +348,25 @@ def _parse_args(argv: list[str] | None = None, *, prog: str | None = None) -> ar 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( "--subtask_chunks_per_gen", type=int, @@ -333,6 +443,8 @@ def _select_observation_to_device(sample: dict, device: Any) -> dict: def _load_policy_and_preprocessor( policy_path: str, dataset_repo_id: str | None, + *, + load_processors_from_checkpoint: bool = False, ) -> tuple[Any, Any, Any, Any]: """Load a policy checkpoint (local path or Hub repo id). @@ -340,6 +452,12 @@ def _load_policy_and_preprocessor( ``preprocessor`` / ``postprocessor`` / ``ds_meta`` are ``None`` when no dataset is provided (rare — needed for autonomous robot mode to have action-denormalisation stats). + + When ``load_processors_from_checkpoint`` is set and no dataset is + given, the pre/post processors are loaded from the checkpoint exactly + like ``lerobot-eval`` (normalizer stats from the saved safetensors, + recipe from ``cfg.recipe_path``). This is what the RoboCasa sim + backend uses so it needs no dataset to match eval-time processing. """ from lerobot.configs import PreTrainedConfig # noqa: PLC0415 from lerobot.policies.factory import make_policy, make_pre_post_processors # noqa: PLC0415 @@ -370,6 +488,12 @@ def _load_policy_and_preprocessor( 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: + # Eval-matching processors: stats from the checkpoint safetensors, + # recipe from cfg.recipe_path. No dataset needed. + preprocessor, postprocessor = make_pre_post_processors( + cfg, pretrained_path=cfg.pretrained_path + ) policy.eval() return policy, preprocessor, postprocessor, ds_meta @@ -1305,7 +1429,15 @@ def run( ) _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 if autonomous_mode and not args.dataset_repo_id: print( "[runtime] ERROR: autonomous robot mode requires --dataset.repo_id " @@ -1315,9 +1447,40 @@ def run( ) return 2 + # Create the sim env subprocess BEFORE the policy initialises CUDA — the + # env worker inherits a corrupt EGL/GL context if forked from a CUDA parent + # (dark/garbled renders). This mirrors eval's make_env-before-make_policy. + 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, + ) + print(f"[runtime] loading policy from {args.policy_path}", flush=True) policy, preprocessor, postprocessor, ds_meta = _load_policy_and_preprocessor( - args.policy_path, args.dataset_repo_id + args.policy_path, + args.dataset_repo_id, + load_processors_from_checkpoint=sim_mode, ) policy_type = getattr(policy.config, "type", None) @@ -1365,8 +1528,32 @@ def run( observation_provider: Callable[[], dict | None] | None = None robot_executor: Callable[[Any], None] | None = None robot = None + sim_backend = None - if autonomous_mode: + 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 # reuse _run_autonomous cleanup (calls .disconnect()) + # 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: print( f"[runtime] connecting to robot.type={args.robot_type} port={args.robot_port}", flush=True, @@ -1416,6 +1603,8 @@ def run( 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 bool(getattr(args, "sim_direct_subtask", False)), ) runtime = LanguageConditionedRuntime( policy_adapter=adapter_factory(policy, gen_config), @@ -1444,6 +1633,20 @@ def run( if bootstrap_state.get("subtask"): runtime.state["current_subtask"] = bootstrap_state["subtask"] + # Let the sim backend read live task/subtask/memory for the video overlay. + if sim_backend is not None: + sim_backend.bind_runtime(runtime) + # Sim runs its control/render loop in the MAIN thread (see + # _run_sim_interactive) — background-thread rendering corrupts EGL. + return _run_sim_interactive( + runtime, + sim_backend, + initial_task=args.task, + max_ticks=args.max_ticks, + panel_label=panel_label, + direct_subtask=bool(args.sim_direct_subtask), + ) + if autonomous_mode: return _run_autonomous( runtime, @@ -1471,6 +1674,132 @@ def run( return _run_repl(runtime, initial_task=args.task, max_ticks=args.max_ticks, panel_label=panel_label) +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: + """Main-thread control loop for the RoboCasa sim backend. + + Unlike ``_run_autonomous`` (which runs ``runtime.run()`` in a daemon + thread), the tick loop — and therefore MuJoCo's EGL rendering — runs in the + MAIN thread. Driving the sim render from a background thread intermittently + corrupts the offscreen GL context (dark/garbled frames); main-thread + stepping matches ``lerobot-eval`` and renders cleanly. Stdin is polled + non-blockingly so typed commands still work while the sim runs. + """ + 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["current_subtask"] = initial_task if direct_subtask else None + runtime.state["mode"] = "action" + + # Clean chat-style prompt. The control loop steps in the MAIN thread (clean + # EGL rendering); the browser live-view shows the rollout, so the terminal + # stays a quiet command line. Nothing is printed mid-step, so typing is never + # clobbered — you can queue the next command any time. + _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["current_subtask"] = None + if hasattr(runtime.policy, "reset"): + runtime.policy.reset() + print("[reset] new kitchen scene", flush=True) + else: + # A bare line is a new command: switch the robot to it + # immediately (clear the in-flight chunk + subtask) and + # force the subtask to regenerate on the very next tick + # (reset the adapter throttle + high-level rate gate). + runtime.set_task(cmd) + # Direct mode: the typed text is the subtask itself; + # otherwise clear it so the model regenerates one. + runtime.state["current_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() + + # One tick in the MAIN thread: subtask/action gen + env.step + render. + # inference_mode matches lerobot-eval's forward context. + 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_repl( runtime: Any, *, initial_task: str | None, max_ticks: int | None, panel_label: str = "Runtime" ) -> int: diff --git a/src/lerobot/runtime/registry.py b/src/lerobot/runtime/registry.py index a4969a0a9..eaf0d7246 100644 --- a/src/lerobot/runtime/registry.py +++ b/src/lerobot/runtime/registry.py @@ -27,6 +27,7 @@ from typing import Any _ADAPTERS: dict[str, str] = { "pi052": "lerobot.policies.pi052.inference.pi052_adapter:PI052PolicyAdapter", + "pi05": "lerobot.policies.pi05.pi05_adapter:PI05PolicyAdapter", } diff --git a/src/lerobot/runtime/sim_robocasa.py b/src/lerobot/runtime/sim_robocasa.py new file mode 100644 index 000000000..43011e895 --- /dev/null +++ b/src/lerobot/runtime/sim_robocasa.py @@ -0,0 +1,479 @@ +"""RoboCasa simulation backend for the interactive language runtime. + +Lets an operator type open-ended prompts (``/action ``) and have a +language-conditioned policy (e.g. PI052) execute them inside a RoboCasa mujoco +kitchen scene. The observation/action pipeline mirrors ``lerobot-eval`` exactly +so behaviour matches offline evaluation; only the *source* of observations and +the *sink* of actions differ from the real-robot backend, which is left +untouched. + +A RoboCasa episode always instantiates a concrete scene (objects + layout) from +its task name, so ``--sim.task`` selects the scene while the prompt typed at the +prompt drives what the policy is asked to do inside it. +""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +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 + + +def _overlay_text( + frame: np.ndarray, task: str | None, subtask: str | None, memory: str | None +) -> np.ndarray: + """Draw task / subtask / memory lines onto an (H, W, 3) uint8 frame. + + Best-effort: returns the frame unchanged if OpenCV is unavailable. + """ + try: + import cv2 # noqa: PLC0415 + except ImportError: + return frame + + lines = [f"{label}: {val}" for label, val in + (("Task", task), ("Subtask", subtask), ("Memory", memory)) if val] + if not lines: + return frame + + img = np.ascontiguousarray(frame) + font, scale, margin = cv2.FONT_HERSHEY_SIMPLEX, 0.5, 6 + max_width = img.shape[1] - 2 * margin + y = 18 + for text in lines: + # naive width-based wrap so long memory strings stay on-frame + words, cur = text.split(), "" + wrapped: list[str] = [] + for w in words: + cand = f"{cur} {w}".strip() + if cv2.getTextSize(cand, font, scale, 1)[0][0] > max_width and cur: + wrapped.append(cur) + cur = w + else: + cur = cand + wrapped.append(cur) + for line in wrapped: + cv2.putText(img, line, (margin, y), font, scale, (0, 0, 0), 3, cv2.LINE_AA) + cv2.putText(img, line, (margin, y), font, scale, (255, 255, 255), 1, cv2.LINE_AA) + y += 20 + return img + + +# RoboCasa's MuJoCo EGL offscreen renderer produces garbled/static frames when +# only ONE worker env is running (reproducible with lerobot-eval --batch_size=1). +# With >=2 workers the renderer is stable. We therefore run the interactive sim +# with a small vec env, drive env 0 with the policy, and ignore the rest. +_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 + reset a RoboCasa AsyncVectorEnv (n_envs=_SIM_N_ENVS), return (env, obs). + + MUST be called BEFORE the policy initialises CUDA in the parent process, so + the forkserver workers don't inherit a CUDA context (which corrupts EGL). + Uses >=2 workers because single-worker EGL rendering is broken on this stack + (garbled frames) — the same reason lerobot-eval renders cleanly only at + batch_size>=2. Only env 0 is driven/displayed. + """ + from lerobot.envs.configs import RoboCasaEnv as RoboCasaEnvConfig # noqa: PLC0415 + + # Higher-res observation cameras => higher-quality display. The policy is + # unaffected: its preprocessor resizes images to 224 and VISUAL norm is + # identity, so only render cost (not behaviour) changes with render_size. + env_cfg = RoboCasaEnvConfig( + task=task, + split=split, + obj_registries=list(obj_registries), + observation_height=render_size, + observation_width=render_size, + ) + # Persistent kitchen: never end/reset on task success, and use a huge horizon + # so the scene doesn't truncate. The user drives it with 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 serving frames from ``get_frame()`` on ``port``. + + Started early (before the ~60s policy load) so the port listens immediately + and browsers get a page instead of connection-refused. ``get_frame`` returns + the latest annotated frame or None (a "waiting" placeholder is shown until + frames arrive). The server thread only reads/encodes frames — no CUDA/EGL — + so it never affects rendering. Returns the server (for shutdown) or None. + """ + 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: + server = ThreadingHTTPServer(("0.0.0.0", port), _Handler) + 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: + """Drive a single RoboCasa gym env from the language runtime. + + Exposes ``observation_provider`` / ``action_executor`` closures matching the + runtime's injected-callable contract, plus ``disconnect`` so the shared + ``_run_autonomous`` cleanup path can close the env (and flush the video). + + The env must be created via :func:`create_sim_env` *before* the policy + touches CUDA (see that function's note on the EGL/CUDA fork hazard). + """ + + 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 + # Camera views to composite into the display frame (order = left→right). + 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._frames: list[np.ndarray] = [] + self._live_counter = 0 + self._latest_frame: np.ndarray | None = None + self._stream_server: Any = None + self._reset_count = 0 + # State getters wired after the runtime exists (bind_runtime), so the + # video overlay can show the live task/subtask/memory. + 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.get("current_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 + # ``task`` feeds the recipe RenderMessagesStep; the PI052 adapter + # overwrites the language tokens with its generated subtask before the + # action forward pass, so this only needs to be present, not exact. + 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() + # Only env 0 is policy-driven; tile its action across all workers so + # env.step gets a full (n_envs, action_dim) batch. The extra workers + # exist only to keep MuJoCo's EGL renderer stable (single-worker + # rendering is broken); their rollouts are ignored. + 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 + if self.record: + self._capture_frame() + # AsyncVectorEnv auto-resets a sub-env after it terminates, so the + # scene continues on its own — no manual reset needed here. + 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 _frontal_obs_image(self) -> np.ndarray | None: + """Return the current front agent-view camera image (H, W, 3) uint8. + + Uses the observation the policy already consumes rather than a separate + ``env.render()`` call: the render path's camera is intermittently + corrupted by the offscreen EGL context, whereas the policy's obs images + come straight through the eval pipeline and stay clean. + """ + pixels = (self._last_obs or {}).get("pixels") + if not isinstance(pixels, dict) or not pixels: + return None + cam = "robot0_agentview_left" if "robot0_agentview_left" in pixels else next(iter(pixels)) + img = np.asarray(pixels[cam]) + if img.ndim == 4: # vec env batches to (1, H, W, C) + img = img[0] + if img.ndim != 3 or img.shape[-1] != 3: + return None + return img.astype(np.uint8) + + def _multiview_frame(self) -> np.ndarray | None: + """Composite the configured camera views (env 0) side by side, labeled. + + Uses the policy's own high-res observation images (env.step already + rendered them), so there's no extra render cost and orientation matches. + """ + 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 = _overlay_text(frame, self._current_task(), subtask, memory) + self._frames.append(annotated) + self._latest_frame = annotated # served by the live MJPEG stream + self._write_live_frame(annotated) + + 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. + """ + if not self.record: + return + 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 not self.record or not self._frames: + return + from datetime import datetime # noqa: PLC0415 + + from lerobot.utils.io_utils import write_video # noqa: PLC0415 + + self.output_dir.mkdir(parents=True, exist_ok=True) + stamp = datetime.now().strftime("%Y%m%d_%H%M%S") + path = self.output_dir / f"sim_{stamp}.mp4" + fps = int((getattr(self.env, "metadata", None) or {}).get("render_fps", 20)) + try: + write_video(str(path), np.stack(self._frames), fps) + logger.info("[sim] wrote video (%d frames) to %s", len(self._frames), path) + print(f"[runtime] sim video saved to {path}", flush=True) + except Exception as exc: # noqa: BLE001 + logger.warning("[sim] write_video 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 (called by _run_autonomous).""" + 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)