From dbb7f5b769e723e91f4172d2e99bd7a3c926440e Mon Sep 17 00:00:00 2001 From: Pepijn Date: Wed, 15 Jul 2026 11:31:19 +0200 Subject: [PATCH] feat(rollout): integrate language runtime --- docs/source/inference.mdx | 49 +++++++++++++- pyproject.toml | 2 +- .../policies/pi052/inference/__init__.py | 2 +- src/lerobot/runtime/__init__.py | 2 +- src/lerobot/runtime/cli.py | 8 +-- src/lerobot/runtime/registry.py | 2 +- .../scripts/lerobot_language_runtime.py | 6 +- src/lerobot/scripts/lerobot_rollout.py | 66 ++++++++++++++++++- .../pi052/test_pi052_runtime_adapter.py | 33 +++++++--- 9 files changed, 148 insertions(+), 22 deletions(-) diff --git a/docs/source/inference.mdx b/docs/source/inference.mdx index 31405b5de..55358568a 100644 --- a/docs/source/inference.mdx +++ b/docs/source/inference.mdx @@ -1,6 +1,6 @@ # Policy Deployment (lerobot-rollout) -`lerobot-rollout` is the single CLI for deploying trained policies on real robots. It supports multiple execution strategies and inference backends, from quick evaluation to continuous recording and human-in-the-loop data collection. +`lerobot-rollout` is the single CLI for deploying trained policies on real robots or in an interactive simulator. It supports multiple execution strategies and inference backends, from quick evaluation to continuous recording, language-driven control, and human-in-the-loop data collection. ## Quick Start @@ -197,6 +197,53 @@ 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. The legacy `lerobot-language-runtime` +command remains as a compatibility alias. + +--- + ## Inference Backends Select a backend with `--inference.type=`. All strategies work with both backends. diff --git a/pyproject.toml b/pyproject.toml index fffd11c6e..70abbf425 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -358,7 +358,7 @@ lerobot-edit-dataset="lerobot.scripts.lerobot_edit_dataset:main" lerobot-setup-can="lerobot.scripts.lerobot_setup_can:main" lerobot-annotate="lerobot.scripts.lerobot_annotate:main" lerobot-rollout="lerobot.scripts.lerobot_rollout:main" -# Interactive high/low-level runtime for language-conditioned policies (pi052, ...). +# Backward-compatible alias; interactive language deployment is part of lerobot-rollout. lerobot-language-runtime="lerobot.scripts.lerobot_language_runtime:main" # ---------------- Tool Configurations ---------------- diff --git a/src/lerobot/policies/pi052/inference/__init__.py b/src/lerobot/policies/pi052/inference/__init__.py index 33a68afb0..15a9b721f 100644 --- a/src/lerobot/policies/pi052/inference/__init__.py +++ b/src/lerobot/policies/pi052/inference/__init__.py @@ -16,7 +16,7 @@ The runtime, REPL, and CLI are policy-agnostic and live in :mod:`lerobot.runtime`. PI052 supplies only :class:`PI052PolicyAdapter`; -the ``lerobot-language-runtime`` entry point wires it into +the ``lerobot-rollout --language`` entry point wires it into :func:`lerobot.runtime.cli.run`. """ diff --git a/src/lerobot/runtime/__init__.py b/src/lerobot/runtime/__init__.py index db38d9a5a..46815e740 100644 --- a/src/lerobot/runtime/__init__.py +++ b/src/lerobot/runtime/__init__.py @@ -17,7 +17,7 @@ The tick loop, REPL, and interactive CLI here are policy-independent. A policy plugs in by subclassing :class:`BaseLanguageAdapter` (or satisfying :class:`LanguageConditionedPolicyAdapter` directly) and registering it in -:mod:`lerobot.runtime.registry`; ``lerobot-language-runtime`` then serves it. +:mod:`lerobot.runtime.registry`; ``lerobot-rollout --language`` then serves it. """ from .adapter import BaseLanguageAdapter, GenerationConfig, LanguageDiagnostics diff --git a/src/lerobot/runtime/cli.py b/src/lerobot/runtime/cli.py index c7e27865d..5997b6e49 100644 --- a/src/lerobot/runtime/cli.py +++ b/src/lerobot/runtime/cli.py @@ -29,7 +29,7 @@ Examples Dry run on a Hub checkpoint, no robot connected — useful for sanity- checking text generation:: - uv run lerobot-language-runtime \\ + uv run lerobot-rollout --language \\ --policy.path= \\ --no_robot \\ --task="please clean the kitchen" @@ -37,7 +37,7 @@ checking text generation:: Same, but feed real frames from an annotated dataset so plan / subtask / memory generation runs against actual video + state:: - uv run lerobot-language-runtime \\ + uv run lerobot-rollout --language \\ --policy.path= \\ --dataset.repo_id= \\ --dataset.episode=0 \\ @@ -46,7 +46,7 @@ Same, but feed real frames from an annotated dataset so plan / subtask With a real robot:: - uv run lerobot-language-runtime \\ + uv run lerobot-rollout --language \\ --policy.path=... \\ --robot.type=so101 --robot.port=/dev/tty.usbmodem... @@ -1561,7 +1561,7 @@ def run( ``adapter_factory`` turns ``(policy, GenerationConfig)`` into a :class:`LanguageConditionedPolicyAdapter` (typically the adapter class). When ``None`` it is resolved from :mod:`lerobot.runtime.registry` by the - loaded policy's type, so a single ``lerobot-language-runtime`` entry + loaded policy's type, so the ``lerobot-rollout`` entry point serves every registered policy. ``panel_label`` defaults to the policy type. """ diff --git a/src/lerobot/runtime/registry.py b/src/lerobot/runtime/registry.py index a8c9d4f8f..088ff4f56 100644 --- a/src/lerobot/runtime/registry.py +++ b/src/lerobot/runtime/registry.py @@ -14,7 +14,7 @@ """Registry mapping a policy type to its language-runtime adapter. -Kept as import strings (resolved lazily) so ``lerobot-language-runtime`` +Kept as import strings (resolved lazily) so ``lerobot-rollout --language`` never imports a policy package until it actually loads that policy — the same pattern as :mod:`lerobot.policies.factory`. """ diff --git a/src/lerobot/scripts/lerobot_language_runtime.py b/src/lerobot/scripts/lerobot_language_runtime.py index b3a5b0089..c197326da 100644 --- a/src/lerobot/scripts/lerobot_language_runtime.py +++ b/src/lerobot/scripts/lerobot_language_runtime.py @@ -13,12 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Entry point for ``lerobot-language-runtime``. +"""Compatibility entry point for ``lerobot-language-runtime``. Policy-agnostic: the runtime resolves the right adapter from the loaded policy's type via :mod:`lerobot.runtime.registry`. A new language-conditioned policy just registers its adapter there — no new -script needed. +script needed. New commands should use ``lerobot-rollout --language`` (or +``lerobot-rollout --sim``); this alias remains so existing scripts do not +break. """ from __future__ import annotations diff --git a/src/lerobot/scripts/lerobot_rollout.py b/src/lerobot/scripts/lerobot_rollout.py index 879070721..06023e6e8 100644 --- a/src/lerobot/scripts/lerobot_rollout.py +++ b/src/lerobot/scripts/lerobot_rollout.py @@ -151,6 +151,7 @@ Usage examples """ import logging +import sys from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401 @@ -241,10 +242,69 @@ def rollout(cfg: RolloutConfig): logger.info("Rollout finished") -def main(): - """CLI entry point for ``lerobot-rollout``.""" +_LANGUAGE_RUNTIME_FLAGS = { + "--language", + "--no_robot", + "--sim", + "--direct_subtask", + "--sim.direct_subtask", + "--disable_memory", + "--fp8", +} +_LANGUAGE_RUNTIME_PREFIXES = ( + "--sim.", + "--chunk_hz", + "--ctrl_hz", + "--high_level_hz", + "--subtask_chunks_per_gen", + "--text_min_new_tokens", + "--text_temperature", + "--text_top_p", +) + + +def _uses_language_runtime(argv: list[str]) -> bool: + """Return whether *argv* selects the interactive language runtime. + + ``--language`` is the explicit selector for real-robot runs whose other + options overlap with the standard rollout CLI. Language-only options also + select it automatically, which keeps the former language-runtime examples + working after replacing their command name with ``lerobot-rollout``. + """ + return any( + arg.split("=", 1)[0] in _LANGUAGE_RUNTIME_FLAGS or arg.startswith(_LANGUAGE_RUNTIME_PREFIXES) + for arg in argv + ) + + +def main(argv: list[str] | None = None): + """CLI entry point for ``lerobot-rollout``. + + Standard policy deployment continues through :class:`RolloutConfig`. + Interactive language-conditioned and RoboCasa runs share this entry point + and are selected with ``--language`` or any language-runtime-only option. + """ register_third_party_plugins() - rollout() + cli_args = list(sys.argv[1:] if argv is None else argv) + if _uses_language_runtime(cli_args): + from lerobot.runtime.cli import run as run_language_runtime + + # ``--language`` is a dispatcher flag, not part of the runtime's own + # argparse surface. All other arguments pass through unchanged. + runtime_args = [arg for arg in cli_args if arg != "--language"] + return run_language_runtime(runtime_args, prog="lerobot-rollout") + + if argv is None: + return rollout() + + # draccus reads sys.argv. Supporting an explicit argv keeps this entry + # point easy to smoke-test and mirrors the language-runtime branch above. + previous_argv = sys.argv + try: + sys.argv = [previous_argv[0], *cli_args] + return rollout() + finally: + sys.argv = previous_argv if __name__ == "__main__": diff --git a/tests/policies/pi052/test_pi052_runtime_adapter.py b/tests/policies/pi052/test_pi052_runtime_adapter.py index ca4a1a109..752660b14 100644 --- a/tests/policies/pi052/test_pi052_runtime_adapter.py +++ b/tests/policies/pi052/test_pi052_runtime_adapter.py @@ -34,21 +34,38 @@ def test_pi052_adapter_strips_say_markers_from_plan_text(): assert adapter.plan_from_text(text) == "Move to the sink." -def test_language_runtime_cli_smoke_does_not_load_model(monkeypatch): - """The general entry resolves the pi052 adapter from the registry by policy type.""" +def test_rollout_language_cli_smoke_does_not_load_model(monkeypatch): + """lerobot-rollout dispatches language flags to the adapter-based runtime.""" from lerobot.runtime import cli - from lerobot.scripts import lerobot_language_runtime + from lerobot.scripts import lerobot_rollout fake_policy = SimpleNamespace(config=SimpleNamespace(device="cpu", type="pi052")) monkeypatch.setattr( cli, "_load_policy_and_preprocessor", - lambda policy_path, dataset_repo_id: (fake_policy, None, None, None), + lambda policy_path, dataset_repo_id, **kwargs: (fake_policy, None, None, None), ) monkeypatch.setattr(cli, "_run_repl", lambda runtime, **kwargs: 0) - assert ( - lerobot_language_runtime.main(["--policy.path=fake", "--no_robot", "--task=clean", "--max_ticks=0"]) - == 0 - ) + assert lerobot_rollout.main(["--policy.path=fake", "--no_robot", "--task=clean", "--max_ticks=0"]) == 0 + + +def test_rollout_language_dispatch_preserves_standard_molmoact2_path(monkeypatch): + """MolmoAct2 only opts into open prompting when a language flag is present.""" + from lerobot.scripts import lerobot_rollout + + standard = [ + "--policy.path=lerobot/MolmoAct2-SO100_101-LeRobot", + "--robot.type=so101_follower", + "--task=pick up the cube", + ] + assert not lerobot_rollout._uses_language_runtime(standard) + assert lerobot_rollout._uses_language_runtime([*standard, "--direct_subtask"]) + assert lerobot_rollout._uses_language_runtime(["--policy.path=lerobot/pi052_robocasa", "--sim"]) + + standard_calls = [] + monkeypatch.setattr(lerobot_rollout, "register_third_party_plugins", lambda: None) + monkeypatch.setattr(lerobot_rollout, "rollout", lambda: standard_calls.append(True)) + lerobot_rollout.main(standard) + assert standard_calls == [True]