diff --git a/src/lerobot/policies/molmoact2/molmoact2_adapter.py b/src/lerobot/policies/molmoact2/molmoact2_adapter.py new file mode 100644 index 000000000..4675ed078 --- /dev/null +++ b/src/lerobot/policies/molmoact2/molmoact2_adapter.py @@ -0,0 +1,58 @@ +# 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. + +"""MolmoAct2 adapter for the language-conditioned runtime. + +MolmoAct2 is a flat VLA: it conditions on a single natural-language ``task`` +string (``"The task is to {task}. ..."``) that its processor packs — together +with the images and discretized state — into model inputs (``input_ids`` / +``pixel_values`` / ...). It has no subtask/memory generation head, so the runtime +just predicts an action chunk from the already-packed observation. + +Run with ``--direct_subtask`` (robot) or ``--sim.direct_subtask`` (sim): what you +type becomes the ``task`` the processor packs, and the runtime does not attempt +subtask/memory generation. The observation provider re-packs on every frame with +the live task (see ``_build_robot_observation_provider`` / the dynamic task +getter in ``runtime.cli``), so typing a new command switches the instruction +immediately. +""" + +from __future__ import annotations + +from typing import Any + +from lerobot.runtime import RuntimeState +from lerobot.runtime.adapter import BaseLanguageAdapter + + +class MolmoAct2PolicyAdapter(BaseLanguageAdapter): + """Runtime bridge for flat MolmoAct2 policies (direct task-text conditioning).""" + + def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: + # The current task/subtask was packed into the model inputs (input_ids, + # pixel_values, ...) by the policy processor, fed the live task by the + # observation provider. ``predict_action_chunk`` resolves the action mode + # from the checkpoint config (``inference_action_mode`` must be set to + # "continuous" or "discrete"). + 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: + # MolmoAct2 has no text-generation head; direct-subtask mode skips this. + return "" diff --git a/src/lerobot/runtime/cli.py b/src/lerobot/runtime/cli.py index eac66c619..ebe513b93 100644 --- a/src/lerobot/runtime/cli.py +++ b/src/lerobot/runtime/cli.py @@ -456,14 +456,32 @@ def _strip_runtime_owned_language_cols(sample: dict) -> None: sample.pop(k, None) +# Model-input keys some policies emit OUTSIDE the ``observation.*`` namespace and +# still need at inference. MolmoAct2's processor packs its prompt + images into +# these top-level keys; PI0-family policies never produce them, so keeping the +# allowlist is a no-op for them. +_MODEL_INPUT_PASSTHROUGH_KEYS = ( + "input_ids", + "attention_mask", + "token_type_ids", + "pixel_values", + "image_token_pooling", + "image_grids", + "image_num_crops", + "pixel_values_videos", + "video_token_pooling", + "video_grids", +) + + def _select_observation_to_device(sample: dict, device: Any) -> dict: - """Filter to ``observation.*`` keys and move tensors to ``device``.""" + """Keep ``observation.*`` (+ model-input passthrough) keys, move tensors to ``device``.""" import torch # noqa: PLC0415 return { k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in sample.items() - if isinstance(k, str) and k.startswith("observation.") + if isinstance(k, str) and (k.startswith("observation.") or k in _MODEL_INPUT_PASSTHROUGH_KEYS) } @@ -871,11 +889,17 @@ def _build_robot_observation_provider( task: str | None, ds_features: dict[str, Any] | None, rerun_log: bool = False, + get_task: Callable[[], str | None] | None = None, ) -> Callable[[], dict | None]: """Closure reading from the robot each call: ``robot.get_observation()`` → ``build_inference_frame`` (state vector + image tensors, batched, on device) → ``EnvTransition``-wrapped preprocessor (rename, normalise) → flat observation batch for ``select_action`` / ``select_message``. + + ``get_task`` (optional) is read every frame so the instruction packed into + the observation tracks the live task/subtask (e.g. MolmoAct2, whose processor + tokenizes the task into ``input_ids`` each frame). Falls back to the static + ``task`` when it returns nothing. """ import torch # noqa: PLC0415 @@ -914,6 +938,9 @@ def _build_robot_observation_provider( target_image_shapes[cam_key] = (int(h), int(w)) def _provider() -> dict | None: + # Live task: re-read every frame so a typed command re-packs the prompt + # (falls back to the static startup task). + cur_task = (get_task() if get_task is not None else None) or task try: raw = robot.get_observation() except Exception as exc: # noqa: BLE001 @@ -933,7 +960,7 @@ def _build_robot_observation_provider( for k, v in raw.items() if isinstance(v, (int, float)) and k not in cam_keys } - rerun_viz.log_robot_frame(raw, cam_keys, state=state, task=task) + rerun_viz.log_robot_frame(raw, cam_keys, state=state, task=cur_task) # The runtime supplies messages itself; strip any language # columns the robot stream may carry through. @@ -997,7 +1024,7 @@ def _build_robot_observation_provider( raw, torch_device, ds_features=ds_features, - task=task, + task=cur_task, robot_type=robot_type, ) else: @@ -1007,7 +1034,7 @@ def _build_robot_observation_provider( obs_tensors = prepare_observation_for_inference( raw, torch_device, - task=task, + task=cur_task, robot_type=robot_type, ) except Exception as exc: # noqa: BLE001 @@ -1602,6 +1629,15 @@ def run( robot_executor: Callable[[Any], None] | None = None robot = None sim_backend = None + # Late-bound handle to the runtime so the robot observation provider can read + # the live task/subtask each frame (the runtime is created further below). + runtime_box: dict[str, Any] = {} + + def _live_task() -> str | None: + rt = runtime_box.get("rt") + if rt is None: + return args.task + return rt.state.get("current_subtask") or rt.state.get("task") or args.task if sim_mode: from lerobot.runtime.sim_robocasa import RoboCasaSimBackend # noqa: PLC0415 @@ -1653,6 +1689,7 @@ def run( task=args.task, ds_features=ds_meta.features if ds_meta is not None else None, rerun_log=bool(args.rerun), + get_task=_live_task, ) robot_executor = _build_robot_action_executor( robot=robot, @@ -1702,6 +1739,8 @@ def run( ctrl_hz=args.ctrl_hz, high_level_hz=args.high_level_hz, ) + # Let the robot observation provider read the live task/subtask each frame. + runtime_box["rt"] = runtime # Apply the startup mode chosen above the task picker. runtime.state["mode"] = startup_mode if args.task: diff --git a/src/lerobot/runtime/registry.py b/src/lerobot/runtime/registry.py index eaf0d7246..a8c9d4f8f 100644 --- a/src/lerobot/runtime/registry.py +++ b/src/lerobot/runtime/registry.py @@ -28,6 +28,7 @@ from typing import Any _ADAPTERS: dict[str, str] = { "pi052": "lerobot.policies.pi052.inference.pi052_adapter:PI052PolicyAdapter", "pi05": "lerobot.policies.pi05.pi05_adapter:PI05PolicyAdapter", + "molmoact2": "lerobot.policies.molmoact2.molmoact2_adapter:MolmoAct2PolicyAdapter", }