refactor(imitation): isolate LeRobot policy runtime

dimos/control/tasks/teleop_ik_task/teleop_ik_task.py

@@ -189,13 +189,13 @@ def _on_controller_pose(
 
     def on_teleop_buttons(self, msg: Buttons, t_now: float) -> bool:
         """Update the all-bound-hands deadman condition."""
-        primary_by_hand = {
-            OperatorHand.LEFT: msg.left_primary,
-            OperatorHand.RIGHT: msg.right_primary,
+        grip_by_hand = {
+            OperatorHand.LEFT: msg.left_grip,
+            OperatorHand.RIGHT: msg.right_grip,
         }
         with self._lock:
             self._last_button_update_time = t_now
-            condition = all(primary_by_hand[hand] for hand in self._bindings)
+            condition = all(grip_by_hand[hand] for hand in self._bindings)
             if self._session_state is _SessionState.ESTOPPED:
                 return True
             if condition and self._session_state is _SessionState.DISENGAGED:

dimos/control/tasks/teleop_ik_task/test_teleop_ik_task.py

@@ -97,11 +97,26 @@ def _buttons(
     right: bool = False,
 ) -> Buttons:
     buttons = Buttons()
-    buttons.left_primary = left
-    buttons.right_primary = right
+    buttons.left_grip = left
+    buttons.right_grip = right
     return buttons
 
 
+def test_face_buttons_do_not_engage_arm_teleop(mocker: MockerFixture) -> None:
+    task = TeleopIKTask(
+        "quest",
+        _config((_binding("right", "right_tool"),)),
+        solver=_solver(mocker),
+    )
+    buttons = Buttons()
+    buttons.right_primary = True
+
+    task.on_teleop_buttons(buttons, 1.0)
+    task.on_right_cartesian_command(_pose(0.5), 1.0)
+
+    assert task.compute(_state()) is None
+
+
 def _pose(x: float) -> PoseStamped:
     return PoseStamped(
         position=Vector3(x, 0.0, 0.0),

dimos/imitation/README.md

@@ -1,12 +1,17 @@
-# Teleop Data Collection → Dataset
+# Imitation Learning
 
-End-to-end: teleoperate an arm, record episodes to a session DB, then convert
+Collect demonstrations, build training datasets, and run trained policies in
+DimOS. Teleoperation records episodes to a session DB, and DataPrep converts
 that DB into a LeRobot or HDF5 dataset for imitation learning.
 
 ```
 teleop (WebXR) ─▶ CollectionRecorder ─▶ session_<robot>_<ts>.db ─▶ dimos dataprep ─▶ dataset
 ```
 
+After training, use the production
+[`LeRobotPolicyModule`](policy/lerobot/README.md) to run a checkpoint against
+live camera and joint-state observations.
+
 ---
 
 ## 1. Record a session

dimos/imitation/policy/lerobot/README.md

@@ -0,0 +1,62 @@
+# LeRobot Policy Module
+
+`LeRobotPolicyModule` runs trained LeRobot policies in a managed Python-native
+subprocess. Its LeRobot, Transformers, Torch, and NumPy versions live in the
+`native/python/lerobot` project and do not change the main DimOS environment.
+
+The host contract subscribes to:
+
+- `color_image: Image`
+- `coordinator_joint_state: JointState`
+- `button_pressed: Buttons`
+- `teleop_buttons: Buttons`
+
+It submits complete, timestamped action chunks to the existing
+`joint_trajectory` task through the control coordinator. The state and action
+vectors use `joint_names` order, including any gripper joint. The policy output
+is already postprocessed into each joint's native absolute coordinate; the
+runtime does not reinterpret gripper values.
+
+```python
+from dimos.imitation.policy.lerobot.module import LeRobotPolicyModule
+
+policy = LeRobotPolicyModule.blueprint(
+    policy_path="outputs/pick/checkpoints/last/pretrained_model",
+    task="pick up the object",
+    joint_names=["arm/joint1", "arm/joint2", "arm/gripper"],
+    fps=30.0,
+    robot_type="my_robot",
+    image_width=640,
+    image_height=480,
+)
+```
+
+The module exposes `preflight_rollout`, `start_rollout`, `stop_rollout`, and
+`rollout_status` RPCs. Preflight loads the checkpoint and processors, validates
+the control task and fresh live observations, and sends no trajectory.
+`start_rollout` refuses to run until preflight passes and rechecks observations
+before starting.
+The runtime rejects missing or stale observations, missing joints, non-finite
+values, incompatible checkpoint features, and malformed action chunks. Pressing the
+configured Quest button (A by default) toggles a preflighted rollout. Pressing
+either middle-finger grip stops rollout. Release both grips before explicitly
+restarting; releasing a grip alone never resumes the policy.
+
+The runtime calls LeRobot's `predict_action_chunk()`, postprocesses the entire
+chunk, clips every action dimension to the checkpoint's recorded data range,
+and executes its first `n_action_steps` at the configured `fps`. Trajectory execution uses the coordinator's existing start-position and velocity handling. Configure `fps` to
+match the action frequency used by the training dataset.
+
+Current limitation: this contract assumes every postprocessed action is an
+absolute target in the connected hardware joint's native coordinate. A generic
+contract for checkpoints that encode grippers in normalized or device-specific
+coordinates remains future work; this runtime does not special-case those
+grippers.
+
+Run isolated runtime checks with:
+
+```bash
+cd native/python/lerobot
+uv run --isolated --locked --group tests --with-editable ../../.. python -m pytest
+uv run --isolated --locked --group tests --with-editable ../../.. python -m mypy
+```

dimos/imitation/policy/lerobot/module.py

@@ -0,0 +1,145 @@
+# Copyright 2026 Dimensional Inc.
+#
+# 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.
+
+"""Host contract for isolated LeRobot policy rollout."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Protocol, TypedDict
+
+from pydantic import Field, field_validator
+
+from dimos.control.tasks.trajectory_task.trajectory_task import (
+    TrajectoryCancellationResult,
+    TrajectoryExecutionResult,
+)
+from dimos.core.core import rpc
+from dimos.core.stream import In
+from dimos.experimental.isolated_python.module import (
+    IsolatedPythonModule,
+    IsolatedPythonModuleConfig,
+)
+from dimos.msgs.sensor_msgs.Image import Image
+from dimos.msgs.sensor_msgs.JointState import JointState
+from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory
+from dimos.spec.utils import Spec
+from dimos.teleop.webxr.controller_types import BUTTON_ALIASES, Buttons
+
+
+class PolicyControlSpec(Spec, Protocol):
+    """Coordinator operations used by policy rollout."""
+
+    def execute_trajectory(
+        self,
+        trajectory: JointTrajectory,
+    ) -> TrajectoryExecutionResult: ...
+
+    def cancel_trajectory(self) -> TrajectoryCancellationResult: ...
+
+    def list_tasks(self) -> list[str]: ...
+
+
+class RolloutStatus(TypedDict):
+    """Operator-facing state of the configured policy rollout."""
+
+    active: bool
+    policy_path: str
+    task: str
+    device: str | None
+    policy_ready: bool
+    observations_ready: bool
+    chunks_accepted: int
+    last_error: str | None
+
+
+class RolloutControlSpec(Spec, Protocol):
+    """The existing policy's operator controls, discoverable by attached clients."""
+
+    def preflight_rollout(self) -> RolloutStatus: ...
+    def start_rollout(self) -> RolloutStatus: ...
+    def stop_rollout(self) -> RolloutStatus: ...
+    def rollout_status(self) -> RolloutStatus: ...
+
+
+class LeRobotPolicyModuleConfig(IsolatedPythonModuleConfig):
+    """Configuration for one checkpoint shared with the isolated runtime."""
+
+    policy_path: str = Field(min_length=1)
+    task: str = Field(min_length=1)
+    device: str | None = None
+    joint_names: list[str] = Field(min_length=1)
+    fps: float = Field(default=30.0, gt=0)
+    robot_type: str = ""
+    image_width: int = Field(default=640, gt=0)
+    image_height: int = Field(default=480, gt=0)
+    max_observation_age_s: float = Field(default=0.5, gt=0)
+    rollout_button: str = "A"
+
+    @field_validator("policy_path")
+    @classmethod
+    def policy_path_must_not_be_blank(cls, policy_path: str) -> str:
+        if not policy_path.strip():
+            raise ValueError("policy_path must not be blank")
+        path = Path(policy_path).expanduser()
+        return str(path.resolve()) if path.exists() else policy_path
+
+    @field_validator("joint_names")
+    @classmethod
+    def joint_names_must_be_unique(cls, joint_names: list[str]) -> list[str]:
+        if len(set(joint_names)) != len(joint_names):
+            raise ValueError("joint_names must not contain duplicates")
+        return joint_names
+
+    @field_validator("rollout_button")
+    @classmethod
+    def rollout_button_must_be_digital(cls, name: str) -> str:
+        if BUTTON_ALIASES.get(name, name) not in Buttons.BITS:
+            raise ValueError(f"unknown Quest button {name!r}")
+        return name
+
+
+class LeRobotPolicyModule(IsolatedPythonModule):
+    """Convert live image and joint-state observations into joint targets."""
+
+    project_dir = "native/python/lerobot"
+    implementation = "dimos_lerobot.runtime:LeRobotPolicyRuntime"
+    config: LeRobotPolicyModuleConfig
+
+    color_image: In[Image]
+    coordinator_joint_state: In[JointState]
+    button_pressed: In[Buttons]
+    teleop_buttons: In[Buttons]
+
+    _control: PolicyControlSpec
+
+    @rpc
+    def preflight_rollout(self) -> RolloutStatus:
+        """Load and validate the policy and live inputs without moving the robot."""
+        raise NotImplementedError
+
+    @rpc
+    def start_rollout(self) -> RolloutStatus:
+        """Start the configured policy until explicitly stopped or it fails."""
+        raise NotImplementedError
+
+    @rpc
+    def stop_rollout(self) -> RolloutStatus:
+        """Stop rollout publication and clear the policy action queue."""
+        raise NotImplementedError
+
+    @rpc
+    def rollout_status(self) -> RolloutStatus:
+        """Return the lifecycle and observation state of the configured policy."""
+        raise NotImplementedError

dimos/imitation/policy/lerobot/test_module.py

@@ -0,0 +1,98 @@
+# Copyright 2026 Dimensional Inc.
+#
+# 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 pathlib import Path
+
+from pydantic import ValidationError
+import pytest
+
+from dimos.experimental.isolated_python.module import contract_rpc_names
+from dimos.imitation.policy.lerobot.module import (
+    LeRobotPolicyModule,
+    LeRobotPolicyModuleConfig,
+)
+from dimos.utils.data import get_project_root
+
+
+def test_contract_imports_without_runtime_dependencies() -> None:
+    assert LeRobotPolicyModule.implementation == "dimos_lerobot.runtime:LeRobotPolicyRuntime"
+    assert contract_rpc_names(LeRobotPolicyModule) == {
+        "preflight_rollout",
+        "rollout_status",
+        "start_rollout",
+        "stop_rollout",
+    }
+
+
+def test_contract_resolves_checkout_runtime_project() -> None:
+    module = LeRobotPolicyModule(
+        policy_path="unused",
+        task="test task",
+        joint_names=["joint"],
+    )
+    try:
+        assert module.runtime_project == get_project_root() / "native/python/lerobot"
+    finally:
+        module.stop()
+
+
+@pytest.mark.parametrize(
+    ("config", "message"),
+    [
+        (
+            {
+                "policy_path": "checkpoint",
+                "task": "test task",
+                "joint_names": ["joint1", "joint1"],
+            },
+            "joint_names must not contain duplicates",
+        ),
+        (
+            {
+                "policy_path": " ",
+                "task": "test task",
+                "joint_names": ["joint1"],
+            },
+            "policy_path must not be blank",
+        ),
+        (
+            {
+                "policy_path": "checkpoint",
+                "task": "test task",
+                "joint_names": ["joint1"],
+                "rollout_button": "NOPE",
+            },
+            "unknown Quest button",
+        ),
+    ],
+)
+def test_config_rejects_ambiguous_names(config: dict[str, object], message: str) -> None:
+    with pytest.raises(ValidationError, match=message):
+        LeRobotPolicyModuleConfig.model_validate(config)
+
+
+def test_existing_relative_checkpoint_is_resolved_before_isolation(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    checkpoint = tmp_path / "checkpoint"
+    checkpoint.mkdir()
+    monkeypatch.chdir(tmp_path)
+
+    config = LeRobotPolicyModuleConfig(
+        policy_path="checkpoint",
+        task="test task",
+        joint_names=["joint1"],
+    )
+
+    assert config.policy_path == str(checkpoint)

dimos/robot/all_blueprints.py

@@ -263,6 +263,7 @@
     "joystick-module": "dimos.robot.unitree.b1.joystick_module.JoystickModule",
     "keyboard-teleop": "dimos.robot.unitree.keyboard_teleop.KeyboardTeleop",
     "keyboard-teleop-module": "dimos.teleop.keyboard.keyboard_teleop_module.KeyboardTeleopModule",
+    "le-robot-policy-module": "dimos.imitation.policy.lerobot.module.LeRobotPolicyModule",
     "lidar-window-relocalization": "dimos.mapping.relocalization.lidar.module.LidarWindowRelocalization",
     "local-map-relocalization": "dimos.mapping.relocalization.lidar.module.LocalMapRelocalization",
     "m20-camera-relay": "dimos.robot.deeprobotics.m20.camera.M20CameraRelay",