armnet-runtime 0.3.2__tar.gz → 0.3.3__tar.gz

{armnet_runtime-0.3.2 → armnet_runtime-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: armnet-runtime
-Version: 0.3.2
+Version: 0.3.3
 Summary: armnet runtime SDK — used by code running inside customer containers on a cell.
 Project-URL: Homepage, https://armnet.dev
 Project-URL: Documentation, https://docs.armnet.dev/packages/runtime/

{armnet_runtime-0.3.2 → armnet_runtime-0.3.3}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
 name = "armnet-runtime"
 # 0.2.0 follows armnet-core: ctx.cell now carries an environment and hands
 # reset/scoring to it, and CompletionStatus moved to armnet-core.
-version = "0.3.2"
+version = "0.3.3"
 description = "armnet runtime SDK \u2014 used by code running inside customer containers on a cell."
 readme = "README.md"
 requires-python = ">=3.12,<3.13"

{armnet_runtime-0.3.2 → armnet_runtime-0.3.3}/src/armnet_runtime/context.py

@@ -69,6 +69,29 @@ _RESET_RAIL_APPROACH_IF_WITHIN = 0.03
 _RESET_RAIL_APPROACH_DISTANCE = 0.04
 
 
+def is_command_safety_violation(exc: BaseException) -> bool:
+    """Whether ``exc`` is the edge rejecting an unsafe commanded motion.
+
+    That lock lasts until :meth:`Cell.reset` is called with ``unlock=True``.
+    A physical stall, an over-current, or a reset timeout is a different trip
+    and stays latched.
+    """
+
+    pending: list[BaseException | None] = [exc]
+    seen: set[int] = set()
+    while pending:
+        current = pending.pop()
+        if current is None or id(current) in seen:
+            continue
+        seen.add(id(current))
+        if "unsafe commanded motion" in str(current):
+            return True
+        pending.append(current.__cause__)
+        if not current.__suppress_context__:
+            pending.append(current.__context__)
+    return False
+
+
 class ResetTimeoutException(RuntimeError):
     """Raised by ``ctx.cell.reset()`` when a manual reset is not actioned in time.
 
@@ -224,20 +247,23 @@ class RuntimeLightbox(_CellDevice):
 class RuntimeCameraMount(_CellDevice):
     """One camera's pan/tilt mount.
 
-    Only cameras with an ``arducam_mount`` block have one; a wrist camera rides
+    Only cameras with a ``servo_mount`` block have one; a wrist camera rides
     on the arm and does not appear here.
     """
 
     camera: str
     default_pan: float = 90.0
     default_tilt: float = 90.0
-    # The cell's mount block verbatim (I2C bus, address, channels, step delay).
-    # Forwarded so the edge can build the mount on demand for a camera it has
-    # not driven yet in this job.
+    # The cell's servo_mount block verbatim: serial port, servo ids and
+    # measured step ranges. Forwarded so the edge can build the mount on
+    # demand for a camera it has not driven yet in this job.
     mount_config: dict[str, Any] = field(default_factory=dict, repr=False)
 
     def set(self, pan: float, tilt: float) -> None:
-        """Aim the mount, in degrees, without disturbing its calibrated pose."""
+        """Aim the mount, without disturbing its calibrated pose.
+
+        ``pan`` and ``tilt`` are encoder steps.
+        """
         self._owner()._set_camera_mount(self, pan, tilt)
 
 
@@ -319,9 +345,8 @@ class Cell:
     camera_mounts: dict[str, RuntimeCameraMount] = field(default_factory=dict)
     """Pan/tilt mounts keyed by camera name.
 
-    Only cameras the cell config gives an ``arducam_mount`` block appear here,
-    so a wrist camera riding on the arm is absent rather than present and
-    immovable.
+    Only cameras the cell config gives a mount block appear here, so a wrist
+    camera riding on the arm is absent rather than present and immovable.
     """
 
     scene_calibration: Optional[SceneCalibration] = None
@@ -582,7 +607,7 @@ class Cell:
             self._instrumented.close()
             self._instrumented = None
 
-    def reset(self, *, confirm: Optional[bool] = None) -> None:
+    def reset(self, *, confirm: Optional[bool] = None, unlock: bool = False) -> None:
         """Restore the workspace to the state this task starts from.
 
         The default asks the cell's environment what that takes. A BusyBox
@@ -594,6 +619,12 @@ class Cell:
         a job's own setup needs checking — or ``confirm=False`` to forbid one on
         an uninstrumented cell.
 
+        Pass ``unlock=True`` after a command-safety trip (the edge rejected an
+        unsafe goal and stopped accepting motion). The lock is released before
+        the arm is returned to rest, so the same job can drive the next
+        rollout. Other trips — a stalled joint, an over-current, a reset
+        timeout — are not cleared this way.
+
         The job's ``disable_automated_reset`` argument forces the confirmation on
         for every reset, so a caller that passes nothing still gets an operator.
 
@@ -602,6 +633,8 @@ class Cell:
         to the position selected for this job's episodes.
         """
 
+        if unlock:
+            self._clear_safety_interlock()
         if self.disable_automated_reset and confirm is None:
             confirm = True
         if self._instrumented is not None:
@@ -866,6 +899,17 @@ class Cell:
                     )
                 time.sleep(_RAIL_POLL_INTERVAL_S)
 
+    def _clear_safety_interlock(self) -> None:
+        """Release a command-safety lock on the robot connector, if there is one."""
+
+        if not self.robot_port or not _looks_like_connector_endpoint(self.robot_port):
+            return
+        response = _connector_request(self.robot_port, {"op": "clear_interlock"})
+        if not response.get("ok"):
+            raise RuntimeError(
+                response.get("error", "could not clear the safety interlock")
+            )
+
     def _reset_cell(self, confirm: bool = True) -> None:
         """Return the robot to rest, then block until the operator confirms.
 
@@ -901,17 +945,26 @@ class Cell:
         if not confirm:
             return
 
-        self.confirm_scene_reset()
+        self._await_operator_reset()
 
     def confirm_scene_reset(self) -> None:
-        """Wait for operator scene staging without moving the robot first.
-
-        Most callers should use :meth:`reset`, which returns the robot to rest
-        before requesting confirmation. A runtime that deliberately starts
-        from a non-rest pose can call ``reset(confirm=False)``, move to that
-        pose, then call this method so the operator stages objects around the
-        robot's actual rollout start configuration.
+        """Ask the operator to stage the scene and confirm, moving nothing.
+
+        :meth:`reset` always returns the arm to rest before any operator
+        confirmation. A job that has already moved the arms to a
+        policy-specific start pose wants the staging prompt alone: returning
+        to rest would undo that pose, and ramping to it only after staging
+        could sweep the arm through the objects the operator just placed.
+        Call :meth:`reset` with ``confirm=False`` first, command the start
+        pose, then call this.
         """
+
+        self._report_progress("Waiting for workspace reset...")
+        self._await_operator_reset()
+
+    def _await_operator_reset(self) -> None:
+        # Operator confirmation on the cell-served operator-call endpoint
+        # (fallback to the dev local-control endpoint).
         request = {"op": "reset", "request": {"kind": "manual"}}
         operator_endpoint = self.operator_call_endpoint or self.local_control_endpoint
         if operator_endpoint:
@@ -1892,13 +1945,21 @@ def _build_runtime_camera_mounts(
         return {}
     mounts: dict[str, RuntimeCameraMount] = {}
     for name, camera in (cell_config.camera_configs or {}).items():
-        mount = camera.get("arducam_mount") if isinstance(camera, dict) else None
+        if not isinstance(camera, dict):
+            continue
+        mount = camera.get("arducam_mount")
+        if not isinstance(mount, dict):
+            mount = camera.get("servo_mount")
         if not isinstance(mount, dict):
             continue
+        # Mid-travel differs by generation: 90 degrees on the hobby mount,
+        # half an encoder turn on a serial one. Only reached by a mount block
+        # that records no aim at all.
+        centre = 2048.0 if mount.get("port") else 90.0
         mounts[name] = RuntimeCameraMount(
             camera=name,
-            default_pan=float(mount.get("pan", 90.0)),
-            default_tilt=float(mount.get("tilt", 90.0)),
+            default_pan=float(mount.get("pan", centre)),
+            default_tilt=float(mount.get("tilt", centre)),
             mount_config=dict(mount),
         )
     return mounts
@@ -1974,12 +2035,24 @@ def _build_camera_config(config: Any) -> Any:
         return config
 
     camera_type = str(config.get("type", "opencv")).lower()
-    # Strip keys that are not OpenCVCameraConfig params (e.g. arducam_mount is
-    # only used by the edge connector for camera mount actuation; id_path is the
-    # USB port used at provisioning time to create the /dev/video_<name> symlink).
-    _NON_CAMERA_KEYS = {"type", "arducam_mount", "id_path"}
+    # Strip keys that are not OpenCVCameraConfig params. Mount blocks are only
+    # used by the edge connector; id_path is the USB port used at provisioning
+    # time to create the /dev/video_<name> symlink. A Pi CSI camera is still
+    # described to LeRobot as an OpenCV camera — the edge swaps in PiCamera
+    # from the cell config's type. use_depth and serial_number tell the edge to
+    # open a RealSense; the job camera is still an OpenCV config, and depth is
+    # read from the remote camera object rather than from these fields.
+    _NON_CAMERA_KEYS = {
+        "type",
+        "arducam_mount",
+        "servo_mount",
+        "id_path",
+        "use_depth",
+        "serial_number",
+    }
+    _OPENCV_SHAPED = {"opencv", "pi", "picam", "picamera", "picamera2", "csi"}
     kwargs = {key: value for key, value in config.items() if key not in _NON_CAMERA_KEYS}
-    if camera_type != "opencv":
+    if camera_type not in _OPENCV_SHAPED:
         return config
 
     try:

armnet_runtime-0.3.3/src/armnet_runtime/depth_view.py

@@ -0,0 +1,280 @@
+"""Colorize a RealSense depth frame so Rerun can show it next to color.
+
+The wire frame is ``uint16`` millimetres, 0 meaning no reading. Rerun and the
+dataset video are 8-bit, so both use the same colorizer as realsense-viewer’s
+default Dynamic preset.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import threading
+from typing import Any
+
+import numpy as np
+
+from armnet_core import env as env_keys
+
+logger = logging.getLogger(__name__)
+
+# librealsense colorizer.cpp, map index 0 ("Jet"). Stops are RGB.
+_JET_STOPS = np.array(
+    [
+        [0, 0, 255],
+        [0, 255, 255],
+        [255, 255, 0],
+        [255, 0, 0],
+        [50, 0, 0],
+    ],
+    dtype=np.float64,
+)
+_JET_STEPS = 4000
+_JET_LUT: np.ndarray | None = None
+_DEPTH_VALUES = 65536
+
+
+def _realsense_jet_lut() -> np.ndarray:
+    """4001-entry RGB cache, matching ``color_map`` with ``steps=4000``."""
+
+    global _JET_LUT
+    if _JET_LUT is None:
+        positions = np.linspace(0.0, 1.0, len(_JET_STOPS))
+        samples = np.linspace(0.0, 1.0, _JET_STEPS + 1)
+        cache = np.empty((samples.size, 3), dtype=np.float64)
+        for channel in range(3):
+            cache[:, channel] = np.interp(samples, positions, _JET_STOPS[:, channel])
+        _JET_LUT = cache.astype(np.uint8)
+    return _JET_LUT
+
+
+def colorize_depth_mm(depth: Any) -> np.ndarray:
+    """Return an RGB ``uint8`` view of a millimetre depth map.
+
+    This is librealsense’s histogram-equalized jet: each depth ranks by how
+    many valid pixels are nearer, then that rank picks the viewer’s jet color.
+    Near is blue, far is red. A zero stays black.
+    """
+
+    plane = np.asarray(depth)
+    if plane.ndim == 3:
+        plane = plane[..., 0]
+    plane = np.ascontiguousarray(plane, dtype=np.uint16)
+    counts = np.bincount(plane.reshape(-1), minlength=_DEPTH_VALUES)
+    # Cumulative count of depths 1..i, matching update_histogram: zeros stay
+    # out of the sum, and hist[i] for i >= 1 is the count of values in 1..i.
+    cumulative = np.cumsum(counts, dtype=np.int64)
+    total = int(cumulative[-1] - counts[0])
+    if total <= 0:
+        return np.zeros((*plane.shape, 3), dtype=np.uint8)
+    # The color depends only on the depth value, so build it once per value
+    # and look every pixel up in that table.
+    rank = (cumulative - int(counts[0])).astype(np.float64) / total
+    lut = _realsense_jet_lut()
+    index = np.clip((rank * (lut.shape[0] - 1)).astype(np.int32), 0, lut.shape[0] - 1)
+    colors = lut[index]
+    colors[0] = 0
+    return np.take(colors, plane, axis=0)
+
+
+def depth_video_features(
+    features: dict[str, Any], camera_names: list[str]
+) -> dict[str, Any]:
+    """Dataset video features for depth, one per color camera.
+
+    Each entry copies the color camera's shape and dtype, under
+    ``observation.images.<name>_depth``. LeRobot then encodes it with the
+    other videos. The frames are the colorized view, because a video feature
+    has to be an 8-bit image.
+    """
+
+    added: dict[str, Any] = {}
+    for name in camera_names:
+        source = features.get(f"observation.images.{name}")
+        if not isinstance(source, dict):
+            continue
+        added[f"observation.images.{name}_depth"] = dict(source)
+    return added
+
+
+def cell_depth_camera_names() -> list[str]:
+    """Cameras whose depth is recorded and shown, from the cell config.
+
+    Known before the robot is built, so an Armnet robot can be told which
+    depth maps to bring back with each observation. The wrist camera stays
+    color-only. Its depth is not fetched, so it is absent from Rerun and from
+    the dataset.
+    """
+
+    raw = os.environ.get(env_keys.CELL_CONFIG)
+    if not raw:
+        return []
+    try:
+        config = json.loads(raw)
+    except json.JSONDecodeError:
+        return []
+    cameras = config.get("camera_configs") if isinstance(config, dict) else None
+    if not isinstance(cameras, dict):
+        return []
+    return [
+        str(name)
+        for name, cfg in cameras.items()
+        if str(name) != "wrist" and isinstance(cfg, dict) and cfg.get("use_depth")
+    ]
+
+
+def depth_cameras(robot: Any) -> list[tuple[str, Any]]:
+    """Cameras on ``robot`` whose depth is recorded and shown.
+
+    For a plain LeRobot robot, whose cameras are read one by one: pair with
+    :func:`read_depth_planes`. An Armnet robot brings depth back with the
+    observation instead, in ``robot.last_depth``. The wrist camera stays
+    color-only. Its depth is not fetched, so it is absent from Rerun and from
+    the dataset.
+    """
+
+    cameras = getattr(robot, "cameras", None) or {}
+    return [
+        (str(name), cam)
+        for name, cam in cameras.items()
+        if str(name) != "wrist" and getattr(cam, "use_depth", False)
+    ]
+
+
+# 40 mm, or 3% of the reading, whichever is larger. Farther pixels are noisier.
+_SPECKLE_LIMIT_MM = np.maximum(
+    40, np.arange(_DEPTH_VALUES, dtype=np.int64) * 3 // 100
+).astype(np.uint16)
+
+
+def smooth_depth_mm(plane: Any, previous: np.ndarray | None = None) -> np.ndarray:
+    """Drop stereo speckle while leaving large holes empty.
+
+    A pixel that disagrees with its 5×5 neighbourhood is replaced by that
+    neighbourhood. Pixels that agree are blended toward it, which takes the
+    grain off a surface. A hole is filled only when the neighbourhood is
+    mostly valid, so the big missing regions stay black. When ``previous`` is
+    the last frame, values that barely moved are blended with it.
+    """
+
+    import cv2
+
+    current = np.asarray(plane)
+    if current.ndim == 3 and current.shape[-1] == 1:
+        current = current[..., 0]
+    current = np.ascontiguousarray(current, dtype=np.uint16)
+    local = cv2.medianBlur(current, 5)
+    agrees = (current > 0) & (
+        cv2.absdiff(current, local) <= np.take(_SPECKLE_LIMIT_MM, current)
+    )
+    # The mean of two uint16 is exact in float32; OpenCV rounds it half to even.
+    blended = cv2.addWeighted(current, 0.5, local, 0.5, 0)
+    # Where the neighbourhood is a hole, keep the pixel. Otherwise a speckle or
+    # a hole takes the neighbourhood, and a pixel that agrees takes the blend.
+    out = np.where(local > 0, np.where(agrees, blended, local), current)
+    if previous is not None and previous.shape == out.shape:
+        previous = np.ascontiguousarray(previous, dtype=np.uint16)
+        close = (out > 0) & (previous > 0) & (cv2.absdiff(out, previous) <= 80)
+        mixed = np.rint(
+            0.4 * out.astype(np.float32) + 0.6 * previous.astype(np.float32)
+        ).astype(np.uint16)
+        out = np.where(close, mixed, out)
+    return out
+
+
+def read_depth_planes(cameras: list[tuple[str, Any]]) -> dict[str, Any]:
+    """Latest raw depth map per camera. Only the read, no processing.
+
+    A remote camera's map stays PNG-encoded; :meth:`DepthView.update` decodes
+    it, on whichever thread does the processing.
+    """
+
+    planes: dict[str, Any] = {}
+    for name, cam in cameras:
+        read = getattr(cam, "read_latest_depth_encoded", None) or cam.read_latest_depth
+        depth = read()
+        if depth is not None:
+            planes[name] = depth
+    return planes
+
+
+class DepthView:
+    """Smooths and colorizes depth maps, and keeps the newest view for Rerun.
+
+    Each camera's last smoothed frame feeds the temporal blend. The viewer's
+    temporal filter does the same job on the raw stream; doing it here covers
+    the picture without another edge restart.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._previous: dict[str, np.ndarray] = {}
+        self._latest: dict[str, np.ndarray] = {}
+
+    def update(self, planes: dict[str, Any]) -> dict[str, np.ndarray]:
+        """Colorized view of each raw plane, keyed by camera name."""
+
+        views: dict[str, np.ndarray] = {}
+        with self._lock:
+            for name, plane in planes.items():
+                if hasattr(plane, "decode"):
+                    plane = plane.decode()
+                smoothed = smooth_depth_mm(plane, self._previous.get(name))
+                self._previous[name] = smoothed
+                views[name] = colorize_depth_mm(smoothed)
+            if views:
+                self._latest = {**self._latest, **views}
+        return views
+
+    def rerun_images(self) -> dict[str, np.ndarray]:
+        """Newest views keyed the way ``log_rerun_data`` expects images."""
+
+        latest = self._latest
+        return {f"images.{name}_depth": view for name, view in latest.items()}
+
+
+class DepthViewWorker:
+    """Runs :meth:`DepthView.update` on its own thread, newest frame only.
+
+    The control loop hands over raw planes and returns at once. A frame that
+    arrives while the previous one is still processing replaces the one
+    waiting, so the view never builds a backlog.
+    """
+
+    def __init__(self, view: DepthView | None = None) -> None:
+        self.view = view if view is not None else DepthView()
+        self._wake = threading.Condition()
+        self._pending: dict[str, Any] | None = None
+        self._closed = False
+        self._thread = threading.Thread(target=self._run, name="depth-view", daemon=True)
+        self._thread.start()
+
+    def submit(self, planes: dict[str, Any]) -> None:
+        if not planes:
+            return
+        with self._wake:
+            self._pending = planes
+            self._wake.notify()
+
+    def rerun_images(self) -> dict[str, np.ndarray]:
+        return self.view.rerun_images()
+
+    def close(self) -> None:
+        with self._wake:
+            self._closed = True
+            self._wake.notify()
+        self._thread.join()
+
+    def _run(self) -> None:
+        while True:
+            with self._wake:
+                while self._pending is None and not self._closed:
+                    self._wake.wait()
+                if self._closed:
+                    return
+                planes, self._pending = self._pending, None
+            try:
+                self.view.update(planes)
+            except Exception:  # noqa: BLE001 - a bad frame must not stop the view
+                logger.exception("depth view update failed")

armnet_runtime-0.3.3/src/armnet_runtime/lerobot/__init__.py

@@ -0,0 +1,52 @@
+"""LeRobot integration helpers for armnet-runtime.
+
+Two ways to drive a cell from LeRobot code:
+
+* Plain LeRobot robots (``SO101Follower``, ``BiSOFollower``, the YAM plugin's
+  ``YAMFollowerRobot``). The import swap replaces their motor buses and
+  cameras with proxies, so unmodified code runs on a cell, at one round trip
+  per device per tick.
+* Armnet robots (``ArmnetSO101Follower``, ``ArmnetBiSO101Follower``,
+  ``ArmnetYAMFollowerRobot``, ``ArmnetBiYAMFollowerRobot``). The same
+  observations and actions, at one round trip to observe and one to act.
+  Armnet's own record, teleop and eval runtimes use these.
+"""
+
+from typing import Any
+
+from armnet_runtime.lerobot.import_swap import install_lerobot_remote_mode
+from armnet_runtime.lerobot.remote_arx5 import RemoteARX5Arm
+from armnet_runtime.lerobot.remote_camera import RemoteOpenCVCamera
+from armnet_runtime.lerobot.remote_feetech import RemoteFeetechMotorsBus
+from armnet_runtime.lerobot.remote_yam import RemoteYAMFollower
+
+# Imported on first use: the SO-101 robots need LeRobot, and a motors-only
+# consumer of this package may not have it installed.
+_ARMNET_ROBOTS = {
+    "ArmnetSO101Follower": "armnet_runtime.lerobot.so101",
+    "ArmnetSO101FollowerConfig": "armnet_runtime.lerobot.so101",
+    "ArmnetBiSO101Follower": "armnet_runtime.lerobot.so101",
+    "ArmnetBiSO101FollowerConfig": "armnet_runtime.lerobot.so101",
+    "ArmnetYAMFollowerRobot": "armnet_runtime.lerobot.yam",
+    "ArmnetYAMFollowerConfig": "armnet_runtime.lerobot.yam",
+    "ArmnetBiYAMFollowerRobot": "armnet_runtime.lerobot.yam",
+    "ArmnetBiYAMFollowerConfig": "armnet_runtime.lerobot.yam",
+}
+
+__all__ = [
+    "RemoteARX5Arm",
+    "RemoteFeetechMotorsBus",
+    "RemoteOpenCVCamera",
+    "RemoteYAMFollower",
+    "install_lerobot_remote_mode",
+    *_ARMNET_ROBOTS,
+]
+
+
+def __getattr__(name: str) -> Any:
+    module_name = _ARMNET_ROBOTS.get(name)
+    if module_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    import importlib
+
+    return getattr(importlib.import_module(module_name), name)