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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (50) hide show
  1. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/PKG-INFO +2 -2
  2. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/pyproject.toml +2 -2
  3. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/__init__.py +6 -0
  4. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/context.py +161 -18
  5. armnet_runtime-0.3.3/src/armnet_runtime/depth_view.py +280 -0
  6. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/leaderboard.py +1 -1
  7. armnet_runtime-0.3.3/src/armnet_runtime/lerobot/__init__.py +52 -0
  8. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/lerobot/_remote_protocol.py +81 -5
  9. armnet_runtime-0.3.3/src/armnet_runtime/lerobot/_robot_link.py +256 -0
  10. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/lerobot/import_swap.py +20 -2
  11. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/lerobot/remote_camera.py +19 -0
  12. armnet_runtime-0.3.3/src/armnet_runtime/lerobot/remote_yam.py +319 -0
  13. armnet_runtime-0.3.3/src/armnet_runtime/lerobot/so101.py +330 -0
  14. armnet_runtime-0.3.3/src/armnet_runtime/lerobot/yam.py +321 -0
  15. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/variation.py +100 -30
  16. armnet_runtime-0.3.3/tests/test_armnet_robots.py +583 -0
  17. armnet_runtime-0.3.3/tests/test_confirm_scene_reset.py +107 -0
  18. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/tests/test_context_arms.py +21 -1
  19. armnet_runtime-0.3.3/tests/test_context_camera_mounts.py +61 -0
  20. armnet_runtime-0.3.3/tests/test_depth_view.py +275 -0
  21. armnet_runtime-0.3.3/tests/test_embodiment.py +166 -0
  22. armnet_runtime-0.3.3/tests/test_persistent_connection.py +213 -0
  23. armnet_runtime-0.3.3/tests/test_remote_yam.py +630 -0
  24. armnet_runtime-0.3.3/tests/test_scene_calibration_context.py +30 -0
  25. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/tests/test_variation.py +120 -35
  26. armnet_runtime-0.3.1/src/armnet_runtime/lerobot/__init__.py +0 -8
  27. armnet_runtime-0.3.1/tests/test_embodiment.py +0 -47
  28. armnet_runtime-0.3.1/tests/test_persistent_connection.py +0 -113
  29. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/.gitignore +0 -0
  30. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/LICENSE +0 -0
  31. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/README.md +0 -0
  32. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/proto/rerun.proto +0 -0
  33. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/scripts/gen_proto.sh +0 -0
  34. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/_rerun_pb2.py +0 -0
  35. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/cli.py +0 -0
  36. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/decorator.py +0 -0
  37. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/env.py +0 -0
  38. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/lerobot/remote_arx5.py +0 -0
  39. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/lerobot/remote_feetech.py +0 -0
  40. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/markers.py +0 -0
  41. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/rerun.py +0 -0
  42. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/src/armnet_runtime/volume_sync.py +0 -0
  43. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/tests/test_cell_instrument.py +0 -0
  44. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/tests/test_completion.py +0 -0
  45. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/tests/test_edge_goal_clamp.py +0 -0
  46. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/tests/test_leaderboard.py +0 -0
  47. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/tests/test_proxy_method_cache.py +0 -0
  48. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/tests/test_rail_recovery.py +0 -0
  49. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/tests/test_rail_reset_restore.py +0 -0
  50. {armnet_runtime-0.3.1 → armnet_runtime-0.3.3}/tests/test_robot_telemetry.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: armnet-runtime
3
- Version: 0.3.1
3
+ Version: 0.3.3
4
4
  Summary: armnet runtime SDK — used by code running inside customer containers on a cell.
5
5
  Project-URL: Homepage, https://armnet.dev
6
6
  Project-URL: Documentation, https://docs.armnet.dev/packages/runtime/
@@ -8,7 +8,7 @@ Author: Ville Kuosmanen
8
8
  License-Expression: Apache-2.0
9
9
  License-File: LICENSE
10
10
  Requires-Python: <3.13,>=3.12
11
- Requires-Dist: armnet-core>=0.3.1
11
+ Requires-Dist: armnet-core>=0.3.2
12
12
  Requires-Dist: protobuf<7,>=5
13
13
  Requires-Dist: pydantic<3,>=2.6
14
14
  Provides-Extra: dev
@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
6
6
  name = "armnet-runtime"
7
7
  # 0.2.0 follows armnet-core: ctx.cell now carries an environment and hands
8
8
  # reset/scoring to it, and CompletionStatus moved to armnet-core.
9
- version = "0.3.1"
9
+ version = "0.3.3"
10
10
  description = "armnet runtime SDK \u2014 used by code running inside customer containers on a cell."
11
11
  readme = "README.md"
12
12
  requires-python = ">=3.12,<3.13"
@@ -15,7 +15,7 @@ license-files = ["LICENSE"]
15
15
  authors = [{ name = "Ville Kuosmanen" }]
16
16
 
17
17
  dependencies = [
18
- "armnet-core>=0.3.1",
18
+ "armnet-core>=0.3.2",
19
19
  "pydantic>=2.6,<3",
20
20
  # Wire format for Rerun visualization packets (armnet_runtime.rerun).
21
21
  "protobuf>=5,<7",
@@ -37,6 +37,7 @@ from armnet_core import (
37
37
 
38
38
  from armnet_runtime.context import (
39
39
  BIMANUAL_SO101_EMBODIMENT,
40
+ BIMANUAL_YAM_EMBODIMENT,
40
41
  SO101_EMBODIMENT,
41
42
  Cell,
42
43
  CompletionStatus,
@@ -45,6 +46,8 @@ from armnet_runtime.context import (
45
46
  RuntimeCameraMount,
46
47
  RuntimeLightbox,
47
48
  RuntimeRail,
49
+ require_eval_embodiment,
50
+ require_bimanual_yam_embodiment,
48
51
  require_so101_embodiment,
49
52
  )
50
53
  from armnet_runtime.decorator import MainRegistrationError, main
@@ -57,6 +60,7 @@ from armnet_runtime.variation import (
57
60
 
58
61
  __all__ = [
59
62
  "BIMANUAL_SO101_EMBODIMENT",
63
+ "BIMANUAL_YAM_EMBODIMENT",
60
64
  "SO101_EMBODIMENT",
61
65
  "Cell",
62
66
  "CompletionStatus",
@@ -65,6 +69,8 @@ __all__ = [
65
69
  "NothingToVary",
66
70
  "Scenario",
67
71
  "sample_scenario",
72
+ "require_eval_embodiment",
73
+ "require_bimanual_yam_embodiment",
68
74
  "require_so101_embodiment",
69
75
  "Embodiment",
70
76
  "Job",
@@ -22,6 +22,7 @@ from typing import Any, Iterable, Mapping, Optional
22
22
  from urllib.parse import urlparse
23
23
 
24
24
  from armnet_core import (
25
+ BIMANUAL_YAM_EMBODIMENT,
25
26
  CellDeviceConfig,
26
27
  CompletionStatus,
27
28
  Embodiment,
@@ -29,6 +30,7 @@ from armnet_core import (
29
30
  InstrumentedCell,
30
31
  Reading,
31
32
  RobotCellConfig,
33
+ SceneCalibration,
32
34
  Task,
33
35
  environment_for,
34
36
  )
@@ -67,6 +69,29 @@ _RESET_RAIL_APPROACH_IF_WITHIN = 0.03
67
69
  _RESET_RAIL_APPROACH_DISTANCE = 0.04
68
70
 
69
71
 
72
+ def is_command_safety_violation(exc: BaseException) -> bool:
73
+ """Whether ``exc`` is the edge rejecting an unsafe commanded motion.
74
+
75
+ That lock lasts until :meth:`Cell.reset` is called with ``unlock=True``.
76
+ A physical stall, an over-current, or a reset timeout is a different trip
77
+ and stays latched.
78
+ """
79
+
80
+ pending: list[BaseException | None] = [exc]
81
+ seen: set[int] = set()
82
+ while pending:
83
+ current = pending.pop()
84
+ if current is None or id(current) in seen:
85
+ continue
86
+ seen.add(id(current))
87
+ if "unsafe commanded motion" in str(current):
88
+ return True
89
+ pending.append(current.__cause__)
90
+ if not current.__suppress_context__:
91
+ pending.append(current.__context__)
92
+ return False
93
+
94
+
70
95
  class ResetTimeoutException(RuntimeError):
71
96
  """Raised by ``ctx.cell.reset()`` when a manual reset is not actioned in time.
72
97
 
@@ -222,20 +247,23 @@ class RuntimeLightbox(_CellDevice):
222
247
  class RuntimeCameraMount(_CellDevice):
223
248
  """One camera's pan/tilt mount.
224
249
 
225
- Only cameras with an ``arducam_mount`` block have one; a wrist camera rides
250
+ Only cameras with a ``servo_mount`` block have one; a wrist camera rides
226
251
  on the arm and does not appear here.
227
252
  """
228
253
 
229
254
  camera: str
230
255
  default_pan: float = 90.0
231
256
  default_tilt: float = 90.0
232
- # The cell's mount block verbatim (I2C bus, address, channels, step delay).
233
- # Forwarded so the edge can build the mount on demand for a camera it has
234
- # not driven yet in this job.
257
+ # The cell's servo_mount block verbatim: serial port, servo ids and
258
+ # measured step ranges. Forwarded so the edge can build the mount on
259
+ # demand for a camera it has not driven yet in this job.
235
260
  mount_config: dict[str, Any] = field(default_factory=dict, repr=False)
236
261
 
237
262
  def set(self, pan: float, tilt: float) -> None:
238
- """Aim the mount, in degrees, without disturbing its calibrated pose."""
263
+ """Aim the mount, without disturbing its calibrated pose.
264
+
265
+ ``pan`` and ``tilt`` are encoder steps.
266
+ """
239
267
  self._owner()._set_camera_mount(self, pan, tilt)
240
268
 
241
269
 
@@ -317,11 +345,13 @@ class Cell:
317
345
  camera_mounts: dict[str, RuntimeCameraMount] = field(default_factory=dict)
318
346
  """Pan/tilt mounts keyed by camera name.
319
347
 
320
- Only cameras the cell config gives an ``arducam_mount`` block appear here,
321
- so a wrist camera riding on the arm is absent rather than present and
322
- immovable.
348
+ Only cameras the cell config gives a mount block appear here, so a wrist
349
+ camera riding on the arm is absent rather than present and immovable.
323
350
  """
324
351
 
352
+ scene_calibration: Optional[SceneCalibration] = None
353
+ """Manipulation-area snapshot selected when this cell process started."""
354
+
325
355
  device_config: Optional[CellDeviceConfig] = None
326
356
  """What this cell may vary and within what bounds.
327
357
 
@@ -577,7 +607,7 @@ class Cell:
577
607
  self._instrumented.close()
578
608
  self._instrumented = None
579
609
 
580
- def reset(self, *, confirm: Optional[bool] = None) -> None:
610
+ def reset(self, *, confirm: Optional[bool] = None, unlock: bool = False) -> None:
581
611
  """Restore the workspace to the state this task starts from.
582
612
 
583
613
  The default asks the cell's environment what that takes. A BusyBox
@@ -589,6 +619,12 @@ class Cell:
589
619
  a job's own setup needs checking — or ``confirm=False`` to forbid one on
590
620
  an uninstrumented cell.
591
621
 
622
+ Pass ``unlock=True`` after a command-safety trip (the edge rejected an
623
+ unsafe goal and stopped accepting motion). The lock is released before
624
+ the arm is returned to rest, so the same job can drive the next
625
+ rollout. Other trips — a stalled joint, an over-current, a reset
626
+ timeout — are not cleared this way.
627
+
592
628
  The job's ``disable_automated_reset`` argument forces the confirmation on
593
629
  for every reset, so a caller that passes nothing still gets an operator.
594
630
 
@@ -597,6 +633,8 @@ class Cell:
597
633
  to the position selected for this job's episodes.
598
634
  """
599
635
 
636
+ if unlock:
637
+ self._clear_safety_interlock()
600
638
  if self.disable_automated_reset and confirm is None:
601
639
  confirm = True
602
640
  if self._instrumented is not None:
@@ -861,6 +899,17 @@ class Cell:
861
899
  )
862
900
  time.sleep(_RAIL_POLL_INTERVAL_S)
863
901
 
902
+ def _clear_safety_interlock(self) -> None:
903
+ """Release a command-safety lock on the robot connector, if there is one."""
904
+
905
+ if not self.robot_port or not _looks_like_connector_endpoint(self.robot_port):
906
+ return
907
+ response = _connector_request(self.robot_port, {"op": "clear_interlock"})
908
+ if not response.get("ok"):
909
+ raise RuntimeError(
910
+ response.get("error", "could not clear the safety interlock")
911
+ )
912
+
864
913
  def _reset_cell(self, confirm: bool = True) -> None:
865
914
  """Return the robot to rest, then block until the operator confirms.
866
915
 
@@ -896,7 +945,25 @@ class Cell:
896
945
  if not confirm:
897
946
  return
898
947
 
899
- # 2. Operator confirmation on the cell-served operator-call endpoint
948
+ self._await_operator_reset()
949
+
950
+ def confirm_scene_reset(self) -> None:
951
+ """Ask the operator to stage the scene and confirm, moving nothing.
952
+
953
+ :meth:`reset` always returns the arm to rest before any operator
954
+ confirmation. A job that has already moved the arms to a
955
+ policy-specific start pose wants the staging prompt alone: returning
956
+ to rest would undo that pose, and ramping to it only after staging
957
+ could sweep the arm through the objects the operator just placed.
958
+ Call :meth:`reset` with ``confirm=False`` first, command the start
959
+ pose, then call this.
960
+ """
961
+
962
+ self._report_progress("Waiting for workspace reset...")
963
+ self._await_operator_reset()
964
+
965
+ def _await_operator_reset(self) -> None:
966
+ # Operator confirmation on the cell-served operator-call endpoint
900
967
  # (fallback to the dev local-control endpoint).
901
968
  request = {"op": "reset", "request": {"kind": "manual"}}
902
969
  operator_endpoint = self.operator_call_endpoint or self.local_control_endpoint
@@ -1595,6 +1662,13 @@ class Context:
1595
1662
  SO101_EMBODIMENT = "lerobot/so-101"
1596
1663
  BIMANUAL_SO101_EMBODIMENT = "lerobot/bimanual_so101"
1597
1664
 
1665
+ _EVAL_EMBODIMENTS = (
1666
+ SO101_EMBODIMENT,
1667
+ BIMANUAL_SO101_EMBODIMENT,
1668
+ BIMANUAL_YAM_EMBODIMENT,
1669
+ )
1670
+ _BIMANUAL_EVAL_EMBODIMENTS = (BIMANUAL_SO101_EMBODIMENT, BIMANUAL_YAM_EMBODIMENT)
1671
+
1598
1672
 
1599
1673
  def require_so101_embodiment(ctx: "Context", runtime_name: str) -> bool:
1600
1674
  """Validate the job's embodiment is a (single or bimanual) SO-101.
@@ -1624,6 +1698,52 @@ def require_so101_embodiment(ctx: "Context", runtime_name: str) -> bool:
1624
1698
  return expect_bimanual
1625
1699
 
1626
1700
 
1701
+ def require_eval_embodiment(ctx: "Context", runtime_name: str) -> None:
1702
+ """Validate the job's embodiment is one LeRobot eval can drive.
1703
+
1704
+ Accepts single SO-101, bimanual SO-101, and bimanual YAM. Both bimanual
1705
+ types must land on a cell wired with left and right arms.
1706
+
1707
+ Raises :class:`NotImplementedError` for any other embodiment, and
1708
+ :class:`RuntimeError` if the embodiment's arm count disagrees with the
1709
+ cell's actual wiring.
1710
+ """
1711
+ if ctx.embodiment not in _EVAL_EMBODIMENTS:
1712
+ supported = ", ".join(repr(name) for name in _EVAL_EMBODIMENTS)
1713
+ raise NotImplementedError(
1714
+ f"{runtime_name} supports {supported}, got {ctx.embodiment!r}"
1715
+ )
1716
+ expect_bimanual = ctx.embodiment in _BIMANUAL_EVAL_EMBODIMENTS
1717
+ if expect_bimanual != ctx.cell.is_bimanual:
1718
+ raise RuntimeError(
1719
+ f"embodiment {ctx.embodiment!r} expects "
1720
+ f"{'two arms (left+right)' if expect_bimanual else 'a single arm'}, "
1721
+ f"but the cell exposes arms={sorted(ctx.cell.arms)}"
1722
+ )
1723
+
1724
+
1725
+ def require_bimanual_yam_embodiment(ctx: "Context", runtime_name: str) -> None:
1726
+ """Validate the job targets the bimanual YAM cell exactly.
1727
+
1728
+ Requires ``ctx.embodiment`` to be :data:`BIMANUAL_YAM_EMBODIMENT` and the
1729
+ cell to expose exactly ``left`` and ``right`` arms — no other embodiment or
1730
+ arm layout is accepted.
1731
+
1732
+ Raises :class:`NotImplementedError` for any other embodiment, and
1733
+ :class:`RuntimeError` if the cell's arms are not exactly left and right.
1734
+ """
1735
+ if ctx.embodiment != BIMANUAL_YAM_EMBODIMENT:
1736
+ raise NotImplementedError(
1737
+ f"{runtime_name} supports {BIMANUAL_YAM_EMBODIMENT!r}, "
1738
+ f"got {ctx.embodiment!r}"
1739
+ )
1740
+ if set(ctx.cell.arms) != {"left", "right"}:
1741
+ raise RuntimeError(
1742
+ f"embodiment {ctx.embodiment!r} expects arms left and right exactly, "
1743
+ f"but the cell exposes arms={sorted(ctx.cell.arms)}"
1744
+ )
1745
+
1746
+
1627
1747
  def build_context() -> Context:
1628
1748
  """Read the cell-injected env vars and construct a :class:`Context`.
1629
1749
 
@@ -1698,6 +1818,9 @@ def build_context() -> Context:
1698
1818
  rails=_build_runtime_rails(cell_config, args),
1699
1819
  lightbox=_build_runtime_lightbox(cell_config),
1700
1820
  camera_mounts=_build_runtime_camera_mounts(cell_config),
1821
+ scene_calibration=(
1822
+ cell_config.scene_calibration if cell_config else None
1823
+ ),
1701
1824
  device_config=(
1702
1825
  CellDeviceConfig.from_cell_config(cell_config) if cell_config else None
1703
1826
  ),
@@ -1822,13 +1945,21 @@ def _build_runtime_camera_mounts(
1822
1945
  return {}
1823
1946
  mounts: dict[str, RuntimeCameraMount] = {}
1824
1947
  for name, camera in (cell_config.camera_configs or {}).items():
1825
- mount = camera.get("arducam_mount") if isinstance(camera, dict) else None
1948
+ if not isinstance(camera, dict):
1949
+ continue
1950
+ mount = camera.get("arducam_mount")
1951
+ if not isinstance(mount, dict):
1952
+ mount = camera.get("servo_mount")
1826
1953
  if not isinstance(mount, dict):
1827
1954
  continue
1955
+ # Mid-travel differs by generation: 90 degrees on the hobby mount,
1956
+ # half an encoder turn on a serial one. Only reached by a mount block
1957
+ # that records no aim at all.
1958
+ centre = 2048.0 if mount.get("port") else 90.0
1828
1959
  mounts[name] = RuntimeCameraMount(
1829
1960
  camera=name,
1830
- default_pan=float(mount.get("pan", 90.0)),
1831
- default_tilt=float(mount.get("tilt", 90.0)),
1961
+ default_pan=float(mount.get("pan", centre)),
1962
+ default_tilt=float(mount.get("tilt", centre)),
1832
1963
  mount_config=dict(mount),
1833
1964
  )
1834
1965
  return mounts
@@ -1904,12 +2035,24 @@ def _build_camera_config(config: Any) -> Any:
1904
2035
  return config
1905
2036
 
1906
2037
  camera_type = str(config.get("type", "opencv")).lower()
1907
- # Strip keys that are not OpenCVCameraConfig params (e.g. arducam_mount is
1908
- # only used by the edge connector for camera mount actuation; id_path is the
1909
- # USB port used at provisioning time to create the /dev/video_<name> symlink).
1910
- _NON_CAMERA_KEYS = {"type", "arducam_mount", "id_path"}
2038
+ # Strip keys that are not OpenCVCameraConfig params. Mount blocks are only
2039
+ # used by the edge connector; id_path is the USB port used at provisioning
2040
+ # time to create the /dev/video_<name> symlink. A Pi CSI camera is still
2041
+ # described to LeRobot as an OpenCV camera — the edge swaps in PiCamera
2042
+ # from the cell config's type. use_depth and serial_number tell the edge to
2043
+ # open a RealSense; the job camera is still an OpenCV config, and depth is
2044
+ # read from the remote camera object rather than from these fields.
2045
+ _NON_CAMERA_KEYS = {
2046
+ "type",
2047
+ "arducam_mount",
2048
+ "servo_mount",
2049
+ "id_path",
2050
+ "use_depth",
2051
+ "serial_number",
2052
+ }
2053
+ _OPENCV_SHAPED = {"opencv", "pi", "picam", "picamera", "picamera2", "csi"}
1911
2054
  kwargs = {key: value for key, value in config.items() if key not in _NON_CAMERA_KEYS}
1912
- if camera_type != "opencv":
2055
+ if camera_type not in _OPENCV_SHAPED:
1913
2056
  return config
1914
2057
 
1915
2058
  try:
@@ -0,0 +1,280 @@
1
+ """Colorize a RealSense depth frame so Rerun can show it next to color.
2
+
3
+ The wire frame is ``uint16`` millimetres, 0 meaning no reading. Rerun and the
4
+ dataset video are 8-bit, so both use the same colorizer as realsense-viewer’s
5
+ default Dynamic preset.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import json
11
+ import logging
12
+ import os
13
+ import threading
14
+ from typing import Any
15
+
16
+ import numpy as np
17
+
18
+ from armnet_core import env as env_keys
19
+
20
+ logger = logging.getLogger(__name__)
21
+
22
+ # librealsense colorizer.cpp, map index 0 ("Jet"). Stops are RGB.
23
+ _JET_STOPS = np.array(
24
+ [
25
+ [0, 0, 255],
26
+ [0, 255, 255],
27
+ [255, 255, 0],
28
+ [255, 0, 0],
29
+ [50, 0, 0],
30
+ ],
31
+ dtype=np.float64,
32
+ )
33
+ _JET_STEPS = 4000
34
+ _JET_LUT: np.ndarray | None = None
35
+ _DEPTH_VALUES = 65536
36
+
37
+
38
+ def _realsense_jet_lut() -> np.ndarray:
39
+ """4001-entry RGB cache, matching ``color_map`` with ``steps=4000``."""
40
+
41
+ global _JET_LUT
42
+ if _JET_LUT is None:
43
+ positions = np.linspace(0.0, 1.0, len(_JET_STOPS))
44
+ samples = np.linspace(0.0, 1.0, _JET_STEPS + 1)
45
+ cache = np.empty((samples.size, 3), dtype=np.float64)
46
+ for channel in range(3):
47
+ cache[:, channel] = np.interp(samples, positions, _JET_STOPS[:, channel])
48
+ _JET_LUT = cache.astype(np.uint8)
49
+ return _JET_LUT
50
+
51
+
52
+ def colorize_depth_mm(depth: Any) -> np.ndarray:
53
+ """Return an RGB ``uint8`` view of a millimetre depth map.
54
+
55
+ This is librealsense’s histogram-equalized jet: each depth ranks by how
56
+ many valid pixels are nearer, then that rank picks the viewer’s jet color.
57
+ Near is blue, far is red. A zero stays black.
58
+ """
59
+
60
+ plane = np.asarray(depth)
61
+ if plane.ndim == 3:
62
+ plane = plane[..., 0]
63
+ plane = np.ascontiguousarray(plane, dtype=np.uint16)
64
+ counts = np.bincount(plane.reshape(-1), minlength=_DEPTH_VALUES)
65
+ # Cumulative count of depths 1..i, matching update_histogram: zeros stay
66
+ # out of the sum, and hist[i] for i >= 1 is the count of values in 1..i.
67
+ cumulative = np.cumsum(counts, dtype=np.int64)
68
+ total = int(cumulative[-1] - counts[0])
69
+ if total <= 0:
70
+ return np.zeros((*plane.shape, 3), dtype=np.uint8)
71
+ # The color depends only on the depth value, so build it once per value
72
+ # and look every pixel up in that table.
73
+ rank = (cumulative - int(counts[0])).astype(np.float64) / total
74
+ lut = _realsense_jet_lut()
75
+ index = np.clip((rank * (lut.shape[0] - 1)).astype(np.int32), 0, lut.shape[0] - 1)
76
+ colors = lut[index]
77
+ colors[0] = 0
78
+ return np.take(colors, plane, axis=0)
79
+
80
+
81
+ def depth_video_features(
82
+ features: dict[str, Any], camera_names: list[str]
83
+ ) -> dict[str, Any]:
84
+ """Dataset video features for depth, one per color camera.
85
+
86
+ Each entry copies the color camera's shape and dtype, under
87
+ ``observation.images.<name>_depth``. LeRobot then encodes it with the
88
+ other videos. The frames are the colorized view, because a video feature
89
+ has to be an 8-bit image.
90
+ """
91
+
92
+ added: dict[str, Any] = {}
93
+ for name in camera_names:
94
+ source = features.get(f"observation.images.{name}")
95
+ if not isinstance(source, dict):
96
+ continue
97
+ added[f"observation.images.{name}_depth"] = dict(source)
98
+ return added
99
+
100
+
101
+ def cell_depth_camera_names() -> list[str]:
102
+ """Cameras whose depth is recorded and shown, from the cell config.
103
+
104
+ Known before the robot is built, so an Armnet robot can be told which
105
+ depth maps to bring back with each observation. The wrist camera stays
106
+ color-only. Its depth is not fetched, so it is absent from Rerun and from
107
+ the dataset.
108
+ """
109
+
110
+ raw = os.environ.get(env_keys.CELL_CONFIG)
111
+ if not raw:
112
+ return []
113
+ try:
114
+ config = json.loads(raw)
115
+ except json.JSONDecodeError:
116
+ return []
117
+ cameras = config.get("camera_configs") if isinstance(config, dict) else None
118
+ if not isinstance(cameras, dict):
119
+ return []
120
+ return [
121
+ str(name)
122
+ for name, cfg in cameras.items()
123
+ if str(name) != "wrist" and isinstance(cfg, dict) and cfg.get("use_depth")
124
+ ]
125
+
126
+
127
+ def depth_cameras(robot: Any) -> list[tuple[str, Any]]:
128
+ """Cameras on ``robot`` whose depth is recorded and shown.
129
+
130
+ For a plain LeRobot robot, whose cameras are read one by one: pair with
131
+ :func:`read_depth_planes`. An Armnet robot brings depth back with the
132
+ observation instead, in ``robot.last_depth``. The wrist camera stays
133
+ color-only. Its depth is not fetched, so it is absent from Rerun and from
134
+ the dataset.
135
+ """
136
+
137
+ cameras = getattr(robot, "cameras", None) or {}
138
+ return [
139
+ (str(name), cam)
140
+ for name, cam in cameras.items()
141
+ if str(name) != "wrist" and getattr(cam, "use_depth", False)
142
+ ]
143
+
144
+
145
+ # 40 mm, or 3% of the reading, whichever is larger. Farther pixels are noisier.
146
+ _SPECKLE_LIMIT_MM = np.maximum(
147
+ 40, np.arange(_DEPTH_VALUES, dtype=np.int64) * 3 // 100
148
+ ).astype(np.uint16)
149
+
150
+
151
+ def smooth_depth_mm(plane: Any, previous: np.ndarray | None = None) -> np.ndarray:
152
+ """Drop stereo speckle while leaving large holes empty.
153
+
154
+ A pixel that disagrees with its 5×5 neighbourhood is replaced by that
155
+ neighbourhood. Pixels that agree are blended toward it, which takes the
156
+ grain off a surface. A hole is filled only when the neighbourhood is
157
+ mostly valid, so the big missing regions stay black. When ``previous`` is
158
+ the last frame, values that barely moved are blended with it.
159
+ """
160
+
161
+ import cv2
162
+
163
+ current = np.asarray(plane)
164
+ if current.ndim == 3 and current.shape[-1] == 1:
165
+ current = current[..., 0]
166
+ current = np.ascontiguousarray(current, dtype=np.uint16)
167
+ local = cv2.medianBlur(current, 5)
168
+ agrees = (current > 0) & (
169
+ cv2.absdiff(current, local) <= np.take(_SPECKLE_LIMIT_MM, current)
170
+ )
171
+ # The mean of two uint16 is exact in float32; OpenCV rounds it half to even.
172
+ blended = cv2.addWeighted(current, 0.5, local, 0.5, 0)
173
+ # Where the neighbourhood is a hole, keep the pixel. Otherwise a speckle or
174
+ # a hole takes the neighbourhood, and a pixel that agrees takes the blend.
175
+ out = np.where(local > 0, np.where(agrees, blended, local), current)
176
+ if previous is not None and previous.shape == out.shape:
177
+ previous = np.ascontiguousarray(previous, dtype=np.uint16)
178
+ close = (out > 0) & (previous > 0) & (cv2.absdiff(out, previous) <= 80)
179
+ mixed = np.rint(
180
+ 0.4 * out.astype(np.float32) + 0.6 * previous.astype(np.float32)
181
+ ).astype(np.uint16)
182
+ out = np.where(close, mixed, out)
183
+ return out
184
+
185
+
186
+ def read_depth_planes(cameras: list[tuple[str, Any]]) -> dict[str, Any]:
187
+ """Latest raw depth map per camera. Only the read, no processing.
188
+
189
+ A remote camera's map stays PNG-encoded; :meth:`DepthView.update` decodes
190
+ it, on whichever thread does the processing.
191
+ """
192
+
193
+ planes: dict[str, Any] = {}
194
+ for name, cam in cameras:
195
+ read = getattr(cam, "read_latest_depth_encoded", None) or cam.read_latest_depth
196
+ depth = read()
197
+ if depth is not None:
198
+ planes[name] = depth
199
+ return planes
200
+
201
+
202
+ class DepthView:
203
+ """Smooths and colorizes depth maps, and keeps the newest view for Rerun.
204
+
205
+ Each camera's last smoothed frame feeds the temporal blend. The viewer's
206
+ temporal filter does the same job on the raw stream; doing it here covers
207
+ the picture without another edge restart.
208
+ """
209
+
210
+ def __init__(self) -> None:
211
+ self._lock = threading.Lock()
212
+ self._previous: dict[str, np.ndarray] = {}
213
+ self._latest: dict[str, np.ndarray] = {}
214
+
215
+ def update(self, planes: dict[str, Any]) -> dict[str, np.ndarray]:
216
+ """Colorized view of each raw plane, keyed by camera name."""
217
+
218
+ views: dict[str, np.ndarray] = {}
219
+ with self._lock:
220
+ for name, plane in planes.items():
221
+ if hasattr(plane, "decode"):
222
+ plane = plane.decode()
223
+ smoothed = smooth_depth_mm(plane, self._previous.get(name))
224
+ self._previous[name] = smoothed
225
+ views[name] = colorize_depth_mm(smoothed)
226
+ if views:
227
+ self._latest = {**self._latest, **views}
228
+ return views
229
+
230
+ def rerun_images(self) -> dict[str, np.ndarray]:
231
+ """Newest views keyed the way ``log_rerun_data`` expects images."""
232
+
233
+ latest = self._latest
234
+ return {f"images.{name}_depth": view for name, view in latest.items()}
235
+
236
+
237
+ class DepthViewWorker:
238
+ """Runs :meth:`DepthView.update` on its own thread, newest frame only.
239
+
240
+ The control loop hands over raw planes and returns at once. A frame that
241
+ arrives while the previous one is still processing replaces the one
242
+ waiting, so the view never builds a backlog.
243
+ """
244
+
245
+ def __init__(self, view: DepthView | None = None) -> None:
246
+ self.view = view if view is not None else DepthView()
247
+ self._wake = threading.Condition()
248
+ self._pending: dict[str, Any] | None = None
249
+ self._closed = False
250
+ self._thread = threading.Thread(target=self._run, name="depth-view", daemon=True)
251
+ self._thread.start()
252
+
253
+ def submit(self, planes: dict[str, Any]) -> None:
254
+ if not planes:
255
+ return
256
+ with self._wake:
257
+ self._pending = planes
258
+ self._wake.notify()
259
+
260
+ def rerun_images(self) -> dict[str, np.ndarray]:
261
+ return self.view.rerun_images()
262
+
263
+ def close(self) -> None:
264
+ with self._wake:
265
+ self._closed = True
266
+ self._wake.notify()
267
+ self._thread.join()
268
+
269
+ def _run(self) -> None:
270
+ while True:
271
+ with self._wake:
272
+ while self._pending is None and not self._closed:
273
+ self._wake.wait()
274
+ if self._closed:
275
+ return
276
+ planes, self._pending = self._pending, None
277
+ try:
278
+ self.view.update(planes)
279
+ except Exception: # noqa: BLE001 - a bad frame must not stop the view
280
+ logger.exception("depth view update failed")
@@ -31,7 +31,7 @@ from typing import Optional
31
31
  logger = logging.getLogger(__name__)
32
32
 
33
33
  # Shared with demo_space/leaderboard.py — the demo reads from the same dataset.
34
- LEADERBOARD_REPO = "villekuosmanen/armnet-demo-leaderboard"
34
+ LEADERBOARD_REPO = "armnet/armnet-demo-leaderboard"
35
35
  RUNS_FILE = "runs.jsonl"
36
36
 
37
37