feat(runtime): add interactive language rollouts

tests/envs/test_robocasa_ordering.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+from lerobot.envs.robocasa import RoboCasaEnv, convert_action
+
+
+def test_robocasa_action_uses_openpi_checkpoint_order():
+    action = np.arange(12, dtype=np.float32)
+
+    converted = convert_action(action)
+
+    np.testing.assert_array_equal(converted["action.end_effector_position"], [0, 1, 2])
+    np.testing.assert_array_equal(converted["action.end_effector_rotation"], [3, 4, 5])
+    np.testing.assert_array_equal(converted["action.gripper_close"], [6])
+    np.testing.assert_array_equal(converted["action.base_motion"], [7, 8, 9, 10])
+    np.testing.assert_array_equal(converted["action.control_mode"], [11])
+
+
+def test_robocasa_state_uses_openpi_checkpoint_order():
+    env = object.__new__(RoboCasaEnv)
+    env.obs_type = "pixels_agent_pos"
+    env.camera_name = []
+    raw_observation = {
+        "state.end_effector_position_relative": np.arange(0, 3),
+        "state.end_effector_rotation_relative": np.arange(3, 7),
+        "state.base_position": np.arange(7, 10),
+        "state.base_rotation": np.arange(10, 14),
+        "state.gripper_qpos": np.arange(14, 16),
+    }
+
+    observation = env._format_raw_obs(raw_observation)
+
+    np.testing.assert_array_equal(observation["agent_pos"], np.arange(16, dtype=np.float32))

tests/runtime/test_adapter.py

@@ -0,0 +1,105 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from lerobot.runtime import RuntimeState
+from lerobot.runtime.adapter import (
+    BaseLanguageAdapter,
+    DirectTaskPolicyAdapter,
+    GenerationConfig,
+)
+
+
+class ScriptedAdapter(BaseLanguageAdapter):
+    """Base adapter whose text generation returns queued strings per kind."""
+
+    def __init__(self, scripts, gen=None):
+        super().__init__(policy=object(), gen=gen)
+        self.scripts = {k: list(v) for k, v in scripts.items()}
+        self.calls = []
+
+    def select_action(self, observation, state):
+        return None
+
+    def generate_text(self, kind, observation, state, user_text=None):
+        self.calls.append(kind)
+        queue = self.scripts.get(kind, [])
+        return queue.pop(0) if queue else ""
+
+
+def test_cascade_sets_subtask_then_memory():
+    adapter = ScriptedAdapter({"subtask": ["pick the red cup"], "memory": ["the cup is grasped"]})
+    state = RuntimeState(task="clean")
+
+    adapter.update_language_state(None, state)
+
+    assert state.language_context["subtask"] == "pick the red cup"
+    assert state.language_context["memory"] == "the cup is grasped"
+    assert adapter.calls == ["subtask", "memory"]
+
+
+def test_nonempty_generation_is_used_verbatim():
+    adapter = ScriptedAdapter({"subtask": [":::: ::"], "memory": ["memory"]})
+    state = RuntimeState(task="clean")
+
+    adapter.update_language_state(None, state)
+
+    assert state.language_context["subtask"] == ":::: ::"
+    assert state.language_context["memory"] == "memory"
+    assert adapter.calls == ["subtask", "memory"]
+
+
+def test_throttle_regenerates_every_n_chunks():
+    adapter = ScriptedAdapter(
+        {
+            "subtask": ["pick the first cup", "pick the second cup"],
+            "memory": ["memory one two three", "memory four five six"],
+        },
+        gen=GenerationConfig(chunks_per_regen=2),
+    )
+    state = RuntimeState(task="clean")
+
+    adapter.update_language_state(None, state)  # generates
+    assert state.language_context["subtask"] == "pick the first cup"
+    adapter.update_language_state(None, state)  # throttled — no generation
+    assert state.language_context["subtask"] == "pick the first cup"
+    adapter.update_language_state(None, state)  # generates again
+    assert state.language_context["subtask"] == "pick the second cup"
+
+
+def test_handle_interjection_sets_plan_and_strips_say():
+    adapter = ScriptedAdapter({"interjection": ["turn to the left now <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

tests/runtime/test_cli.py

@@ -0,0 +1,75 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from lerobot.runtime.cli import _build_rollout_runtime_io, _parse_args
+
+
+def test_parse_args_preserves_rollout_robot_overrides():
+    args = _parse_args(
+        [
+            "--policy.path=checkpoint",
+            "--robot.type=so101_follower",
+            "--robot.calibration_dir=/tmp/calibration",
+        ]
+    )
+
+    assert args.robot_type == "so101_follower"
+    assert "--robot.calibration_dir=/tmp/calibration" in args.raw_argv
+
+
+def test_parse_args_rejects_removed_dataset_replay_flags():
+    with pytest.raises(SystemExit):
+        _parse_args(["--policy.path=checkpoint", "--dataset.repo_id=dataset"])
+
+
+def test_rollout_runtime_io_uses_context_processors():
+    robot = MagicMock()
+    robot.robot_type = "mock_robot"
+    robot.cameras = {}
+    robot.get_observation.return_value = {"joint.pos": 1.5}
+    ctx = SimpleNamespace(
+        hardware=SimpleNamespace(robot_wrapper=robot),
+        runtime=SimpleNamespace(cfg=SimpleNamespace(device="cpu")),
+        processors=SimpleNamespace(
+            robot_observation_processor=lambda observation: observation,
+            robot_action_processor=lambda pair: pair[0],
+        ),
+        policy=SimpleNamespace(
+            preprocessor=lambda observation: observation,
+            postprocessor=lambda action: action,
+        ),
+        data=SimpleNamespace(
+            dataset_features={
+                "observation.state": {
+                    "dtype": "float32",
+                    "shape": (1,),
+                    "names": ["joint.pos"],
+                },
+                "action": {"dtype": "float32", "shape": (1,), "names": ["joint.pos"]},
+            }
+        ),
+    )
+    provider, executor = _build_rollout_runtime_io(ctx, rerun_log=False, get_task=lambda: "move")
+
+    observation = provider()
+    executor(torch.tensor([[2.0]]))
+
+    assert observation["observation.state"].shape == (1, 1)
+    robot.send_action.assert_called_once_with({"joint.pos": 2.0})

tests/runtime/test_language_runtime.py

@@ -0,0 +1,100 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import threading
+import time
+
+from lerobot.runtime import LanguageConditionedRuntime, Tick
+
+
+class FakeAdapter:
+    def __init__(self):
+        self.updated = False
+        self.interjections = []
+
+    def select_action(self, observation, state):
+        assert observation == {"observation.state": 1}
+        assert state.task == "clean"
+        return ["a0", "a1"]
+
+    def update_language_state(self, observation, state):
+        self.updated = True
+        state.set_context("subtask", "pick cup", label="subtask")
+
+    def handle_interjection(self, user_text, observation, state):
+        self.interjections.append(user_text)
+        state.set_context("plan", "new plan", label="plan")
+
+
+def test_runtime_tick_updates_language_enqueues_and_dispatches_action():
+    adapter = FakeAdapter()
+    executed = []
+    runtime = LanguageConditionedRuntime(
+        policy_adapter=adapter,
+        observation_provider=lambda: {"observation.state": 1},
+        action_executor=executed.append,
+    )
+    runtime.set_task("clean")
+
+    logs = runtime.step_once()
+
+    assert adapter.updated
+    assert runtime.state.language_context["subtask"] == "pick cup"
+    assert executed == ["a0"]
+    assert list(runtime.state.action_queue) == ["a1"]
+    assert "  subtask: pick cup" in logs
+
+
+def test_runtime_handles_user_interjection():
+    adapter = FakeAdapter()
+    runtime = LanguageConditionedRuntime(
+        policy_adapter=adapter,
+        observation_provider=lambda: {"observation.state": 1},
+    )
+    runtime.set_task("clean")
+    runtime.state.extra["recent_interjection"] = "please say ok"
+    runtime.state.emit("user_interjection")
+
+    runtime.step_once()
+
+    assert "please say ok" in adapter.interjections
+    assert runtime.state.language_context["plan"] == "new plan"
+
+
+def test_prompt_change_discards_in_flight_action_chunk():
+    started = threading.Event()
+    release = threading.Event()
+
+    class BlockingAdapter(FakeAdapter):
+        def select_action(self, observation, state):
+            started.set()
+            assert release.wait(timeout=2)
+            return ["stale"]
+
+    runtime = LanguageConditionedRuntime(
+        policy_adapter=BlockingAdapter(),
+        observation_provider=lambda: {"observation.state": 1},
+    )
+    runtime.set_task("old task")
+    runtime.state.tick = Tick(index=1, monotonic_seconds=time.monotonic())
+    inference = threading.Thread(target=runtime.maybe_enqueue_action_chunk, kwargs={"force": True})
+    inference.start()
+    assert started.wait(timeout=2)
+
+    runtime.set_task("new task")
+    release.set()
+    inference.join(timeout=2)
+
+    assert not inference.is_alive()
+    assert list(runtime.state.action_queue) == []

tests/runtime/test_sim_robocasa.py

@@ -0,0 +1,82 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+from types import SimpleNamespace
+
+import numpy as np
+
+from lerobot.runtime.sim_robocasa import RoboCasaSimBackend
+from lerobot.utils.video_annotation import annotate_frame
+
+
+def test_overlay_draws_each_label_once(monkeypatch):
+    put_text_calls = []
+    rectangle_calls = []
+
+    def put_text(image, text, origin, font, scale, color, thickness, line_type):
+        put_text_calls.append((text, color, thickness))
+        return image
+
+    def rectangle(image, start, end, color, thickness):
+        rectangle_calls.append((start, end, color, thickness))
+        return image
+
+    def add_weighted(src1, alpha, src2, beta, gamma, *, dst):
+        dst[:] = src1 * alpha + src2 * beta + gamma
+        return dst
+
+    fake_cv2 = SimpleNamespace(
+        FONT_HERSHEY_SIMPLEX=0,
+        LINE_AA=16,
+        getTextSize=lambda text, font, scale, thickness: ((len(text) * 7, 10), 0),
+        putText=put_text,
+        rectangle=rectangle,
+        addWeighted=add_weighted,
+    )
+    monkeypatch.setitem(sys.modules, "cv2", fake_cv2)
+
+    frame = np.full((120, 480, 3), 200, dtype=np.uint8)
+    annotated = annotate_frame(
+        frame,
+        (("Task", "close the fridge"), ("Subtask", "reach for the handle"), ("Memory", None)),
+    )
+
+    assert [call[0] for call in put_text_calls] == [
+        "Task: close the fridge",
+        "Subtask: reach for the handle",
+    ]
+    assert all(color == (255, 255, 255) and thickness == 1 for _, color, thickness in put_text_calls)
+    assert len(rectangle_calls) == 1
+    assert not np.shares_memory(annotated, frame)
+
+
+def test_capture_updates_live_frame_when_recording_is_disabled(monkeypatch):
+    backend = object.__new__(RoboCasaSimBackend)
+    frame = np.full((8, 8, 3), 42, dtype=np.uint8)
+    written = []
+    backend.record = False
+    backend.runtime_state = None
+    backend._multiview_frame = lambda: frame
+    backend._current_task = lambda: "task"
+    backend._subtask_getter = None
+    backend._memory_getter = None
+    backend._latest_frame = None
+    backend._write_live_frame = written.append
+    monkeypatch.setattr("lerobot.runtime.sim_robocasa.annotate_frame", lambda image, labels: image)
+
+    backend._capture_frame()
+
+    assert backend._latest_frame is frame
+    assert written == [frame]