diff --git a/docs/source-en/rst_source/development/add_primitive.rst b/docs/source-en/rst_source/development/add_primitive.rst index dcefdd5b..947a38bf 100644 --- a/docs/source-en/rst_source/development/add_primitive.rst +++ b/docs/source-en/rst_source/development/add_primitive.rst @@ -35,7 +35,7 @@ call. They differ only in how the method is implemented. Add a scripted primitive ------------------------ -Adding a scripted primitive usually involves three steps: +Adding a scripted primitive usually involves two steps: 1. **Add a method to the primitives.** Add the method to the current environment's primitives class, such as @@ -43,16 +43,25 @@ Adding a scripted primitive usually involves three steps: the tool-call arguments, performs the work, usually through one or more ``self._env.step(...)`` calls, and returns a small log ``dict``. + Mark the method with :func:`~rpent.tools.toolkit.updatestate` so the + toolkit re-renders state (``get_env_state``) automatically after it runs: + .. code-block:: python + from rpent.tools.toolkit import updatestate + + @updatestate def open_drawer(self, dx: float = 0.15) -> dict: # Move end-effector back by dx while gripper is closed. for _ in range(N): self._env.step(build_open_drawer_chunk(dx)) return {"ok": True, "dx": dx} + Read-only tools (``view_env_state``, ``back_project``, ``segment``, + ...) are simply left unmarked -- the toolkit skips state capture for them. + 2. **Add the tool schema.** Add an entry to ``TOOLS_SPEC`` in - ``toolkit.py``: + ``robots//tools.py``: .. code-block:: python @@ -67,13 +76,10 @@ Adding a scripted primitive usually involves three steps: }, } -3. **Register the tool in the toolkit.** Route it through the toolkit's - ``_step`` helper so that state is re-rendered after execution: - - .. code-block:: python - - self.add_tool("open_drawer", OPEN_DRAWER_SPEC, - lambda **kw: self._step("open_drawer", **kw)) +Once both exist, the toolkit registers the tool automatically: it iterates +``TOOLS_SPEC`` and binds each spec to the matching primitive-driver method +(e.g. ``getattr(self._primitives, name)``); ``@updatestate`` decides whether +state is captured -- no explicit ``add_tool`` call is needed. After these steps, the ``api``, ``claude_code``, and ``codex`` planners can all call the primitive without any other code changes. @@ -167,8 +173,10 @@ Design principles for a new primitive the state dump reflecting the post-action world. Don't let the primitive return before the render finishes. - **Return small dicts.** Tool return values are fed back to the LLM - as text. Store larger content, such as images, depth data, and - ``states.json``, in the state dump instead. + as text. Save larger observations through ``EnvState.save``; ``EnvState`` + automatically records each logical base name in its owned + ``StepRecord.artifacts`` set. Expose images through ``view_env_state`` and + geometry through environment tools rather than returning raw paths. - **Guardrails belong in env_server**, not in the toolkit. The LLM can and will call any tool with any arguments; workspace bounds and safety clamps must be enforced on the server side. diff --git a/docs/source-en/rst_source/development/add_robot.rst b/docs/source-en/rst_source/development/add_robot.rst index 2111a9f7..b3578486 100644 --- a/docs/source-en/rst_source/development/add_robot.rst +++ b/docs/source-en/rst_source/development/add_robot.rst @@ -220,23 +220,29 @@ state needed for the current run. It exposes one method per primitive tool **Tool definitions and handlers** — a module-level ``TOOLS_SPEC`` list of Anthropic-style tool definitions (``name``, ``description``, ``input_schema``), plus any module-level functions referenced by the toolkit (e.g. -``view_driver_state``, ``back_project``, ``finish``). - -**Per-step state dump** — ``dump_state(primitives, output_dir, step_idx, log)`` -serializes whatever state the agent will read back via the ``view_*`` tools -(images, depths, JSON state, camera meta) into ``output_dir``. +``view_env_state``, ``back_project``, ``finish``). + +**Per-step state dump** — ``dump_state(driver, env_state, log)`` opens +``env_state.record_step(...)`` and receives the allocated step index; the +``StepRecord`` is appended and committed immediately. Save large observations +through ``env_state.save(...)`` — inside a ``record_step`` block the ``step`` +argument may be omitted (it defaults to the new step), pass an explicit +``step=`` to target a different step, and ``step=None`` for run-level +artifacts. ``EnvState`` adds every successfully saved base name to the step's +flat ``artifacts`` set automatically. Readers use the canonical artifact +filenames rather than maintaining a parallel observation index. **Toolkit class** — subclass ``rpent.tools.toolkit.Toolkit``: -- build the primitives in ``__init__`` through a custom initialization - helper (named ``init_primitives_clean`` in LIBERO; it wipes stale - ``images/`` etc., constructs the primitives, and dumps step 0), +- build the primitive driver in ``__init__`` through a custom initialization + helper (named ``init_primitives_clean`` in LIBERO; it calls + ``EnvState.reset()``, constructs the primitives, and dumps step 0), - register each tool with ``self.add_tool(name, spec, handler)`` — stateless - readers (``view_driver_state``, ``finish``, …) bind directly to module-level + readers (``view_env_state``, ``finish``, …) bind directly to module-level functions; primitive tools route through ``_step(name, **kwargs)`` which calls ``getattr(self._primitives, name)(**kwargs)`` and re-renders state, -- override ``close()`` to write any remaining agent-side artifacts (e.g. the - LIBERO toolkit saves the agentview MP4 there). +- override ``close()`` to save remaining agent-side artifacts through + ``EnvState`` (for example ``state.save("episode.mp4", frames, step=None)``). ``primitives_kwargs`` (forwarded from ``__init__.py:get_toolkit``) is the dict the toolkit passes verbatim to your primitives' ``__init__`` — typically @@ -246,14 +252,15 @@ Conventions worth keeping ------------------------- - ``output_dir`` is the working directory that the runner creates for each - run. Images, depths, ``states.json``, transcripts, ``episode.mp4``, and other - artifacts go there. + run. Environment observations are owned by ``EnvState``; callers use logical + base names and never construct storage paths. Transcripts and other + run-management outputs share the same run directory. - Tool definitions use the Anthropic format (``name`` / ``description`` / ``input_schema``). Every tool registered with ``self.add_tool(...)`` is exposed to all planners. - Server-side return values must be picklable and torch-free. - Each primitive tool dumps a fresh state snapshot after running so the next - ``view_driver_state`` call reflects the post-action world. + ``view_env_state`` call reflects the post-action world. - Treat ``dump_state`` as the source of truth for what the agent sees — any new modality (e.g. tactile, force) goes through it. diff --git a/docs/source-en/rst_source/development/interfaces.rst b/docs/source-en/rst_source/development/interfaces.rst index bf6a6a59..02d865a0 100644 --- a/docs/source-en/rst_source/development/interfaces.rst +++ b/docs/source-en/rst_source/development/interfaces.rst @@ -90,7 +90,7 @@ Subclass ``Toolkit`` in ``robots//toolkit.py`` and register env tools with ends; optional ``_image_bytes`` (etc.) to return camera images. The base class already registers common file tools; call ``super().__init__()`` then -``add_tool`` for env tools. Per-step state and ``view_driver_state`` are in +``add_tool`` for env tools. Per-step state and ``view_env_state`` are in :doc:`add_primitive`. Inter-process communication diff --git a/docs/source-en/rst_source/quickstart.rst b/docs/source-en/rst_source/quickstart.rst index b9ba30b5..ac49ab68 100644 --- a/docs/source-en/rst_source/quickstart.rst +++ b/docs/source-en/rst_source/quickstart.rst @@ -103,14 +103,11 @@ A successful run: by the elapsed time, token usage, and path to the run record. 3. With the Dashboard enabled, also streams agent output, camera views, the action timeline, and clip replays to the Dashboard. -4. By default, artifacts are saved under - ``logs/__t_s/``. They include - ``transcript_*.json`` (run record), ``states.json`` (one record per - environment step), ``recipe_*.jsonl`` (action sequence), and - ``episode.mp4`` (episode video). - -After the run, inspect the final record in ``states.json``: -``libero_terminated`` set to ``true`` means LIBERO judged the task complete. -You can also open ``episode.mp4`` to review the run. +4. By default, artifacts are saved under ``logs/__t_s/``. They include ``transcript_*.json`` (run record), ``states.json`` (the versioned ``EnvState`` manifest), ``recipe_*.jsonl`` (action sequence), and ``episode.mp4`` (episode video). Step artifact files use flat, zero-padded step prefixes with a minimum width of two digits and are managed internally by ``EnvState``. + +Inspect the final state through the Dashboard or +``view_env_state(step=-1)``. Its top-level ``libero_terminated`` value is the +benchmark outcome. ``states.json`` is internal ``EnvState`` storage and should +not be parsed by callers. You can also open ``episode.mp4`` to review the run. If something goes wrong, inspect the four log files described at the bottom of :doc:`installation`. diff --git a/docs/source-en/rst_source/usage/libero.rst b/docs/source-en/rst_source/usage/libero.rst index 4fd25442..7e90c120 100644 --- a/docs/source-en/rst_source/usage/libero.rst +++ b/docs/source-en/rst_source/usage/libero.rst @@ -116,7 +116,7 @@ What runs where transports (HTTP or socket). It returns only the top compressed PNG mask. - **toolkit** (``robots/libero/toolkit.py``) — defines the tools the LLM can call: ``pi0_pick`` (fed to Pi0.5), ``move_to``, - ``rotate_wrist``, ``back_project``, ``view_driver_state``, + ``rotate_wrist``, ``back_project``, ``view_env_state``, ``finish``, … Tools the planner can call @@ -147,8 +147,10 @@ Physical action tools advance the environment and record new state and images. coordinates. - ``segment(prompt=... / point=..., ...)`` — use SAM3 to segment an existing image with a text or point prompt. -- ``view_driver_state(step=None)`` — read an existing state and image record. -- ``view_camera_meta(camera=..., step=None)`` — read existing camera metadata. +- ``view_env_state(step=-1)`` — read a recorded state and its embedded + observation images. Step ``0`` is initial; ``-1`` is latest. +- ``view_camera_meta(camera=..., step=-1)`` — read camera metadata for a + recorded step. Step ``-1`` is latest. - ``finish(status, summary)`` — end the current run. These tools do not advance the environment. diff --git a/docs/source-zh/rst_source/development/add_primitive.rst b/docs/source-zh/rst_source/development/add_primitive.rst index 909a7e82..59022970 100644 --- a/docs/source-zh/rst_source/development/add_primitive.rst +++ b/docs/source-zh/rst_source/development/add_primitive.rst @@ -33,22 +33,31 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法 添加一个脚本化原语 ------------------ -添加脚本化原语通常需要以下三个步骤： +添加脚本化原语通常需要以下两个步骤： 1. **在 primitives 中添加方法。** 在当前环境的 primitives 类（如 ``LiberoPrimitives``、``MyRobotPrimitives``）中添加 一个方法。该方法接收工具调用的参数，执行一次或多次 ``self._env.step(...)``，并返回一个简短的日志字典。 + 为该方法加上 :func:`~rpent.tools.toolkit.updatestate` 装饰器， + toolkit 会在其执行后自动重新渲染状态（``get_env_state``）： + .. code-block:: python + from rpent.tools.toolkit import updatestate + + @updatestate def open_drawer(self, dx: float = 0.15) -> dict: # 保持夹爪闭合，沿 -x 方向后拉 dx 米。 for _ in range(N): self._env.step(build_open_drawer_chunk(dx)) return {"ok": True, "dx": dx} -2. **添加工具定义。** 在 ``toolkit.py`` 的 ``TOOLS_SPEC`` 中新增一项： + 只读工具（``view_env_state``、``back_project``、``segment`` 等） + 无需装饰——toolkit 会跳过它们的状态捕获。 + +2. **添加工具定义。** 在 ``robots//tools.py`` 的 ``TOOLS_SPEC`` 中新增一项： .. code-block:: python @@ -63,13 +72,9 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法 }, } -3. **在 toolkit 中注册工具。** 通过 toolkit 的 ``_step`` 辅助函数运行 - 该工具，使其在执行结束后自动重新渲染状态： - - .. code-block:: python - - self.add_tool("open_drawer", OPEN_DRAWER_SPEC, - lambda **kw: self._step("open_drawer", **kw)) +两者就位后，toolkit 会自动注册该工具：它遍历 ``TOOLS_SPEC``，把每个定义 +绑定到对应的 primitive driver 方法（如 ``getattr(self._primitives, name)``）， +由 ``@updatestate`` 决定是否捕获状态——无需显式调用 ``add_tool``。 完成以上步骤后，``api``、``claude_code`` 和 ``codex`` 三种 planner 都可以调用该工具，无需修改其他代码。 @@ -152,7 +157,9 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法 - **每个工具执行结束后都要保存新的状态快照。** 下一轮需要读取动作执行后的 环境状态，因此原语不能在渲染完成前返回。 - **工具只返回简短的字典。** 返回值会以文本形式提供给 LLM；图像、深度数据和 - ``states.json`` 等较大的内容则通过状态快照提供。 + 其他大型观测应通过 ``EnvState.save`` 保存；``EnvState`` 会把每个逻辑基础 + 文件名自动加入其持有的 ``StepRecord.artifacts`` 集合。图像通过 + ``view_env_state`` 提供，几何数据通过环境工具访问，不返回原始路径。 - **安全限制由 ``env_server`` 强制执行。** LLM 可能使用任意参数调用工具， 因此工作空间边界和安全限制不能只依赖 toolkit。 diff --git a/docs/source-zh/rst_source/development/add_robot.rst b/docs/source-zh/rst_source/development/add_robot.rst index 995a13e7..e646a37d 100644 --- a/docs/source-zh/rst_source/development/add_robot.rst +++ b/docs/source-zh/rst_source/development/add_robot.rst @@ -207,23 +207,27 @@ toolkit 模块通常包含四部分： **工具定义和处理函数** 包括模块级的 ``TOOLS_SPEC`` 列表（列表元素采用 Anthropic API 的工具定义格式，包含 ``name``、``description`` 和 ``input_schema``），以及 toolkit 引用的模块级函数，例如 -``view_driver_state``、``back_project`` 和 ``finish``。 +``view_env_state``、``back_project`` 和 ``finish``。 -**每步状态 dump** —— ``dump_state(primitives, output_dir, step_idx, log)`` 把 agent -之后会通过 ``view_*`` 工具读回的所有状态 (图像、深度、JSON 状态、camera meta) -序列化到 ``output_dir``。 +**每步状态 dump** —— ``dump_state(driver, env_state, log)`` 通过 +``env_state.record_step(...)`` 创建由 ``EnvState`` 持有的步骤，并取得分配的 +step index；该 ``StepRecord`` 会被立即追加并提交。大型观测通过 +``env_state.save(...)`` 保存——在 ``record_step`` 块内可省略 ``step`` 参数 +（默认指向刚创建的步骤），传显式 ``step=`` 可指定其它步骤，``step=None`` +用于运行级工件。每次保存成功后，``EnvState`` 会自动把基础文件名加入该 +``StepRecord`` 的扁平 ``artifacts`` 集合；读取方直接使用规范化的工件文件名。 **Toolkit 类** 继承 ``rpent.tools.toolkit.Toolkit``： -- 在 ``__init__`` 中通过自定义的初始化辅助方法构建 primitives（LIBERO - 中的方法名为 ``init_primitives_clean``；它会清理过期的 ``images/`` 等目录、 - 构造原语并 dump 第 0 步）, +- 在 ``__init__`` 中通过自定义的初始化辅助方法构建 primitive driver（LIBERO + 中的方法名为 ``init_primitives_clean``；它会调用 ``EnvState.reset()``、构造 + 原语并 dump 第 0 步）, - 用 ``self.add_tool(name, spec, handler)`` 注册每个工具。无状态的读取工具 - （如 ``view_driver_state``、``finish``）直接绑定模块级函数；原语工具通过 + （如 ``view_env_state``、``finish``）直接绑定模块级函数；原语工具通过 ``_step(name, **kwargs)`` 调用。``_step`` 使用 - ``getattr(self._primitives, name)(**kwargs)`` 调用 primitives 方法并重新渲染状态； -- 重写 ``close()``，将 agent 侧生成的文件写入磁盘（例如 LIBERO toolkit - 在这里保存 agentview MP4）。 + ``getattr(self._primitives, name)(**kwargs)`` 调用 driver 方法并重新渲染状态； +- 重写 ``close()``，通过 ``EnvState`` 保存 agent 侧剩余工件（例如 + ``state.save("episode.mp4", frames, step=None)``）。 ``primitives_kwargs`` 由 ``__init__.py:get_toolkit`` 转发给 toolkit，再原样传入 primitives 的 ``__init__``。其中通常包含 @@ -232,14 +236,15 @@ primitives 的 ``__init__``。其中通常包含 建议遵循的约定 -------------- -- ``output_dir`` 是 runner 为单次运行创建的临时目录。图像、深度数据、 - ``states.json``、transcript 和 ``episode.mp4`` 等工件都写入该目录。 +- ``output_dir`` 是 runner 为单次运行创建的工作目录。环境观测由 + ``EnvState`` 管理；调用方只使用逻辑基础文件名，不自行拼接存储路径。 + transcript 等运行管理输出与环境工件共享该目录。 - 工具定义使用 Anthropic API 格式（``name`` / ``description`` / ``input_schema``）。 每个用 ``self.add_tool(...)`` 注册的工具都会暴露给所有 planner。 - 环境侧的返回值必须可 pickle，且不包含 torch 对象。 - 每个原语工具执行后要 dump 一次新的状态快照, 这样下一次 - ``view_driver_state`` 看到的是动作后的世界。 + ``view_env_state`` 看到的是动作后的世界。 - ``dump_state`` 是 Agent 获取环境状态的唯一数据来源；任何新的模态 （例如触觉、力）都通过它提供。 diff --git a/docs/source-zh/rst_source/development/interfaces.rst b/docs/source-zh/rst_source/development/interfaces.rst index d28c3d34..822b02fa 100644 --- a/docs/source-zh/rst_source/development/interfaces.rst +++ b/docs/source-zh/rst_source/development/interfaces.rst @@ -87,7 +87,7 @@ Planner 需要回传相机图时可设 ``_image_bytes`` 等字段。 基类已注册公共文件工具；子类 ``super().__init__()`` 后追加本环境工具即可。逐步状态与 -``view_driver_state`` 见 :doc:`add_primitive`。 +``view_env_state`` 见 :doc:`add_primitive`。 进程间通信 ---------- diff --git a/docs/source-zh/rst_source/quickstart.rst b/docs/source-zh/rst_source/quickstart.rst index b9cf71c7..2b153eb0 100644 --- a/docs/source-zh/rst_source/quickstart.rst +++ b/docs/source-zh/rst_source/quickstart.rst @@ -95,7 +95,9 @@ LIBERO-PRO 仿真资源。下面以 LIBERO-PRO 和 ``claude_code`` planner 1. 终端会先显示 ``env_server``、``vla_server`` 和 ``sam3_server`` 的启动信息。 2. 智能体的逐轮输出和工具调用会显示在终端中；运行结束时还会显示耗时、token 用量和运行记录的路径。 3. 启用 Dashboard 后，智能体的输出、相机视图、动作时间线和片段回放也会实时显示在 Dashboard 中。 -4. 默认输出目录为 ``logs/__t_s/``，其中包含 ``transcript_*.json``\ （运行记录）、``states.json``\ （每个环境步的记录）、``recipe_*.jsonl``\ （动作序列）和 ``episode.mp4``\ （回合录像）。 +4. 默认输出目录为 ``logs/__t_s/``，其中包含 ``transcript_*.json``\ （运行记录）、``states.json``\ （带版本号的 ``EnvState`` 清单）、``recipe_*.jsonl``\ （动作序列）和 ``episode.mp4``\ （回合录像）。每步工件文件采用至少两位、零填充的步骤前缀扁平命名，并由 ``EnvState`` 在内部管理。 -运行结束后，查看 ``states.json`` 的最后一条记录：``libero_terminated`` 为 ``true`` 表示 LIBERO 已判定任务完成；也可以打开 ``episode.mp4`` 复核运行过程。 +通过 Dashboard 或 ``view_env_state(step=-1)`` 查看最终状态；其顶层 +``libero_terminated`` 即为基准任务结果。``states.json`` 是 ``EnvState`` 的内部 +存储，调用方不应直接解析。也可以打开 ``episode.mp4`` 复核运行过程。 出问题时，参考 :doc:`installation` 页底部提到的四份日志文件。 diff --git a/docs/source-zh/rst_source/usage/libero.rst b/docs/source-zh/rst_source/usage/libero.rst index 2675458b..940de6da 100644 --- a/docs/source-zh/rst_source/usage/libero.rst +++ b/docs/source-zh/rst_source/usage/libero.rst @@ -112,7 +112,7 @@ LIBERO-PRO 核心套件一览 排名第一的压缩 PNG mask。 - **toolkit（工具集）** （``robots/libero/toolkit.py``）—— 定义 LLM 能调用的工具：``pi0_pick``（交给 Pi0.5）、``move_to``、``rotate_wrist``、 - ``back_project``、``view_driver_state``、``finish``… + ``back_project``、``view_env_state``、``finish``… Planner 能调用的工具 -------------------- @@ -141,8 +141,10 @@ LIBERO 工具分为物理动作工具和只读工具。 - ``back_project(row, col, ...)`` —— 将图像像素反投影到世界坐标。 - ``segment(prompt=... / point=..., ...)`` —— 通过 SAM3 对已有图像进行文本或 点提示分割。 -- ``view_driver_state(step=None)`` —— 读取已有的状态和图像记录。 -- ``view_camera_meta(camera=..., step=None)`` —— 读取已有的相机元数据。 +- ``view_env_state(step=-1)`` —— 读取已记录的状态和内嵌观测图像；第 0 步为 + 初始状态，``-1`` 表示最新状态。 +- ``view_camera_meta(camera=..., step=-1)`` —— 读取指定步骤的相机元数据； + ``-1`` 表示最新状态。 - ``finish(status, summary)`` —— 结束当前运行。 这些工具不会推进环境。 diff --git a/pyproject.toml b/pyproject.toml index 0200869a..00b1d343 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,8 +21,8 @@ classifiers = [ # 4 - Beta # 5 - Production/Stable "Development Status :: 2 - Pre-Alpha", - "Environment :: GPU :: NVIDIA CUDA :: 12 :: 12.4", "Intended Audience :: Developers", + "Intended Audience :: Robotics Enthusiasts", "Programming Language :: Python :: 3", ] @@ -67,6 +67,10 @@ libero-plus = [ "rpent[libero]", "rlinf-liberoplus", ] +lerobot = [ + "lerobot[feetech,intelrealsense,kinematics]>=0.4.4", + "matplotlib", +] full = [ "rpent[rlinf,openpi,libero-pro,sam3]", ] @@ -79,6 +83,14 @@ sam3 = [ [tool.uv] prerelease = "allow" +# lerobot is not co-installable with the RLinf stack (libero/openpi). +conflicts = [ + [{ extra = "lerobot" }, { extra = "openpi" }], + [{ extra = "lerobot" }, { extra = "libero" }], + [{ extra = "lerobot" }, { extra = "libero-pro" }], + [{ extra = "lerobot" }, { extra = "libero-plus" }], + [{ extra = "lerobot" }, { extra = "full" }], +] [tool.setuptools] include-package-data = true diff --git a/resources/franka/calibration_boards/franka_charuco_7x5_25mm.json b/resources/franka/calibration_boards/franka_charuco_7x5_25mm.json new file mode 100644 index 00000000..d21400ca --- /dev/null +++ b/resources/franka/calibration_boards/franka_charuco_7x5_25mm.json @@ -0,0 +1,18 @@ +{ + "type": "charuco", + "dictionary": "DICT_4X4_50", + "squares_x": 7, + "squares_y": 5, + "square_length_m": 0.025, + "marker_length_m": 0.018, + "margin_m": 0.012, + "dpi": 300, + "image_width_px": 2349, + "image_height_px": 1759, + "print_instructions": [ + "Print the PDF at 100% / actual size.", + "Disable fit-to-page or scaling.", + "Use matte paper and mount it flat to cardboard/foam board.", + "Measure one printed square and update square_length_m if needed." + ] +} \ No newline at end of file diff --git a/resources/franka/calibration_boards/franka_charuco_7x5_25mm.pdf b/resources/franka/calibration_boards/franka_charuco_7x5_25mm.pdf new file mode 100644 index 00000000..d77409bb Binary files /dev/null and b/resources/franka/calibration_boards/franka_charuco_7x5_25mm.pdf differ diff --git a/resources/franka/calibration_boards/franka_charuco_7x5_25mm.png b/resources/franka/calibration_boards/franka_charuco_7x5_25mm.png new file mode 100644 index 00000000..83792409 Binary files /dev/null and b/resources/franka/calibration_boards/franka_charuco_7x5_25mm.png differ diff --git a/resources/lerobot/memory/MEMORY.md b/resources/lerobot/memory/MEMORY.md new file mode 100644 index 00000000..ed1108cc --- /dev/null +++ b/resources/lerobot/memory/MEMORY.md @@ -0,0 +1 @@ +- [For placing in the white plate at x≈0.25,y≈0.105, high vertical standoffs can be unreachable; move low/reachable over the plate and release around z≈-0.03.](plate_place_low_standoff.md) diff --git a/resources/lerobot/memory/plate_place_low_standoff.md b/resources/lerobot/memory/plate_place_low_standoff.md new file mode 100644 index 00000000..b32c986f --- /dev/null +++ b/resources/lerobot/memory/plate_place_low_standoff.md @@ -0,0 +1,7 @@ +--- +hook: "For placing in the white plate at x\u22480.25,y\u22480.105, high vertical standoffs can be unreachable; move low/reachable over the plate and release around z\u2248-0.03." +env: lerobot +updated: 2026-07-03 +--- + +In the green-cube-to-white-plate task, the plate interior back-projected near x=0.224–0.252, y=0.103–0.109, z≈-0.067. Carrying the grasped cube to a high down-oriented standoff at [0.245, 0.105, 0.080] with yaw 90 failed, settling ~45 mm short (final z≈0.037) with a reach note. A lower sequence was reachable and successful: [0.230,0.105,0.035] -> [0.240,0.106,0.015] -> [0.248,0.106,0.005], then lower to about [0.251,0.108,-0.030] and open. The cube remained in the plate after lifting the gripper. For this plate region, avoid insisting on a high vertical standoff; use reachable lower waypoints and verify with the scene/arm images. diff --git a/robots/franka/__init__.py b/robots/franka/__init__.py new file mode 100644 index 00000000..8d3a403b --- /dev/null +++ b/robots/franka/__init__.py @@ -0,0 +1,173 @@ +"""Franka environment extension.""" +from __future__ import annotations + +import argparse +from datetime import datetime +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from robots.franka.prompt import system_prompt, user_prompt +from rpent.dashboard.events import DashboardEventSink +from rpent.envs.env_spec import EnvSpec, RunConfig +from rpent.envs.prompt_bundle import PromptBundle +from rpent.utils.config import get_repo_root +from rpent.utils.logging import get_logger + +if TYPE_CHECKING: + from rpent.utils.daemon import ProcessDaemon + from rpent.utils.rpc import RpcClient + +logger = get_logger("franka") + + +def get_env_spec() -> EnvSpec: + """Return the Franka env identity, prompt bundle, and runner hooks. + + Tool schemas, handlers, and the MCP allowlist live on the toolkit (see + :func:`get_toolkit`). The three runner hooks (:func:`_add_cli_args` / + :func:`_parse_config` / :func:`_init_runtime`) keep ``rpent/cli/main.py`` + env-agnostic, mirroring :mod:`robots.libero`. + """ + return EnvSpec( + name="franka", + prompts=PromptBundle( + system=system_prompt, + user=user_prompt, + ), + add_cli_args=_add_cli_args, + parse_config=_parse_config, + init_runtime=_init_runtime, + ) + + +def get_toolkit( + *, + primitives_kwargs: dict[str, Any], + dashboard_events: DashboardEventSink, +): + """Return the Franka toolkit (common tools + Cartesian primitives). + + ``primitives_kwargs`` is assembled by :func:`_init_runtime` and carries the + env RPC stub (``{"env": FrankaEnvClient(...)}``). + """ + from robots.franka.toolkit import FrankaToolkit + + return FrankaToolkit( + dashboard_events=dashboard_events, + **primitives_kwargs, + ) + + +def _add_cli_args(parser: argparse.ArgumentParser, use_dashboard: bool) -> None: + """Register Franka CLI flags on the shared ``parser``. + + ``use_dashboard`` is unused: the Franka setup is a real robot with no + suite/task/seed, so there is nothing for the (libero-shaped) dashboard + launcher to fill in. + """ + del use_dashboard + parser.add_argument( + "--env-endpoint", default=None, + help="[protocol://]host:port of an existing franka env_server " + "(protocol=http|socket, defaults to http). If unset, it is spawned " + "via run_env_server.sh (RLinf .venv).", + ) + + +def _parse_config(args: argparse.Namespace) -> RunConfig: + """Derive per-run identifiers for a Franka run. + + Real robots have no suite/task/seed, so the run is identified by the env + name. The dashboard is currently libero-shaped, so it is not wired here. + """ + if getattr(args, "dashboard", False): + logger.warning( + "--dashboard is only supported for the libero env; " + "continuing without the live dashboard." + ) + + recipe_tag = "franka" + output_dir = args.output_dir + if output_dir is None: + timestamp = datetime.now().strftime("%Y%m%d-%H:%M:%S") + output_dir = get_repo_root() / "logs" / f"{timestamp}_franka" + output_dir = Path(output_dir) + + return RunConfig( + recipe_tag=recipe_tag, + output_dir=output_dir, + prompt_vars={"env_name": "franka", "recipe_tag": recipe_tag}, + dashboard_state=None, + task_desc={"env": "franka"}, + ) + + +def _parse_endpoint(endpoint: str) -> tuple[str, str, int]: + """Parse ``[protocol://]host:port`` into ``(protocol, host, port)``. + + Protocol defaults to ``http`` when the prefix is omitted. + """ + if "://" in endpoint: + protocol, _, rest = endpoint.partition("://") + else: + protocol, rest = "http", endpoint + host, _, port = rest.partition(":") + if not host or not port: + raise ValueError( + f"--env-endpoint must be [protocol://]host:port, got {endpoint!r}" + ) + return protocol, host, int(port) + + +def _init_runtime( + args: argparse.Namespace, + output_dir: Path, +) -> tuple[list[ProcessDaemon], dict[str, Any]]: + """Spawn (or attach to) the Franka env_server; build primitives_kwargs. + + The Franka driver runs in the RLinf ``.venv`` with the catkin workspace + sourced, so it is spawned via ``run_env_server.sh`` (not this interpreter). + Pass ``--env-endpoint`` to attach to an already-running server. No VLA + server — the Cartesian primitives are scripted. + + Heavy deps are imported lazily so a bare ``import robots.franka`` (for + ``get_env_spec`` / ``get_toolkit``) doesn't drag them in. + """ + from robots.franka.env_client import FrankaEnvClient + from rpent.utils.daemon import ProcessDaemon, pick_free_port + from rpent.utils.http_rpc import HttpRpcClient + from rpent.utils.rpc import wait_for_ready + from rpent.utils.socket_rpc import SocketRpcClient + + daemons: list[ProcessDaemon] = [] + if args.env_endpoint is None: + host, port = "127.0.0.1", pick_free_port() + env_daemon = ProcessDaemon( + name="env_server", + cmd=[ + "bash", + str(get_repo_root() / "robots" / "franka" / "run_env_server.sh"), + "--output-dir", str(output_dir), + "--transport", "http", + "--host", host, + "--port", str(port), + ], + log_path=str(Path(output_dir) / "env_server.log"), + ) + env_daemon.start() + daemons.append(env_daemon) + env_client: RpcClient = HttpRpcClient(f"http://{host}:{port}") + wait_for_ready(env_client) + else: + protocol, host, port = _parse_endpoint(args.env_endpoint) + if protocol == "socket": + env_client = SocketRpcClient(host, port) + elif protocol == "http": + env_client = HttpRpcClient(f"http://{host}:{port}") + else: + raise ValueError( + f"--env-endpoint protocol must be socket or http, got {protocol!r}" + ) + + primitives_kwargs = {"env": FrankaEnvClient(env_client)} + return daemons, primitives_kwargs diff --git a/robots/franka/auto_calibrate_cameras.py b/robots/franka/auto_calibrate_cameras.py new file mode 100644 index 00000000..103c88aa --- /dev/null +++ b/robots/franka/auto_calibrate_cameras.py @@ -0,0 +1,631 @@ +#!/usr/bin/env python3 +"""Automatic Franka RGB-D camera calibration helpers. + +The default mode calibrates the fixed scene camera into the Franka base frame +(``panda_link0``) without markers: + +1. move the TCP through a small, conservative 3-D grid, +2. at each pose, toggle the Franka Hand while the arm is stationary, +3. segment the moving fingers in the scene RGB image and use aligned depth to + get a camera-frame point, +4. pair that point with the live robot TCP position and fit ``T_base_cam`` with + RANSAC Kabsch, +5. save the calibration record to + ``~/.cache/rpent/franka/camera_calibration/.json`` and ask + the running env server to reload it. + +This is the right first calibration for ``back_project`` because the scene +camera is fixed. The wrist camera is eye-in-hand; calibrating it correctly needs +hand-eye calibration (``T_tcp_cam``) using a fixed fiducial/ChArUco/AprilTag or a +scene-calibrated reference target observed from multiple wrist poses. This file +keeps the wrist record format ready, but does not invent an unsafe automatic +wrist calibration from one moving camera alone. + +Run the env server first, then run in the physicalagent env:: + + python robots/franka/auto_calibrate_cameras.py --port 5599 --yes + +Offline math check:: + + python robots/franka/auto_calibrate_cameras.py --self-test +""" +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path +from typing import Any + +import imageio.v2 as imageio +import numpy as np + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from robots.franka import calibration as franka_calib # noqa: E402 +from robots.lerobot import geometry as geom # noqa: E402 +from rpent.utils.socket_rpc import SocketRpcClient # noqa: E402 + +_DEFAULT_GRID_X = (0.46, 0.54, 0.60) +_DEFAULT_GRID_Y = (-0.08, 0.0, 0.08) +_DEFAULT_GRID_Z = (0.20, 0.27) +_DEFAULT_DOWN_EULER = [float(np.pi), 0.0, 0.0] + + +def _self_test() -> int: + """Validate Kabsch/RANSAC and motion-blob detection offline.""" + rng = np.random.default_rng(3) + q, _ = np.linalg.qr(rng.standard_normal((3, 3))) + if np.linalg.det(q) < 0: + q[:, 0] = -q[:, 0] + T_true = np.eye(4) + T_true[:3, :3] = q + T_true[:3, 3] = rng.normal(size=3) + + cam = rng.normal(size=(12, 3)) + base = geom.transform_points(T_true, cam) + rng.normal(scale=0.001, size=(12, 3)) + base[5] += [0.12, -0.08, 0.05] + T_est, rmse, inliers = geom.ransac_kabsch(cam, base, thresh_m=0.015) + fit_ok = np.allclose(T_est, T_true, atol=2e-2) and not bool(inliers[5]) + print( + "ransac_kabsch: " + f"rmse={rmse:.5f}m inliers={int(inliers.sum())}/12 " + f"outlier_excluded={not bool(inliers[5])} -> {'OK' if fit_ok else 'FAIL'}" + ) + + H, W = 480, 640 + rgb_open = np.zeros((H, W, 3), np.uint8) + rgb_closed = rgb_open.copy() + rgb_closed[210:235, 330:365] = 255 + depth = np.full((H, W), 0.55, np.float32) + K = np.array([[600.0, 0.0, 320.0], [0.0, 600.0, 240.0], [0.0, 0.0, 1.0]]) + det = _detect_tip_pixel_by_motion(rgb_open, rgb_closed, depth, K) + det_ok = det is not None and abs(det["pixel"][0] - 222) < 4 and abs(det["pixel"][1] - 347) < 4 + print(f"detect_tip: {det if det else 'None'} -> {'OK' if det_ok else 'FAIL'}") + return 0 if fit_ok and det_ok else 1 + + +def _detect_motion_blob_numpy( + rgb_open, + rgb_closed, + depth_m, + K, + *, + diff_thresh: int, + min_area: int, + max_area: int, +) -> dict | None: + """Locate the largest changed depth-valid blob without OpenCV.""" + a = np.asarray(rgb_open, dtype=np.float32).mean(axis=2) + b = np.asarray(rgb_closed, dtype=np.float32).mean(axis=2) + depth_m = np.asarray(depth_m, dtype=np.float64) + mask = (np.abs(a - b) >= int(diff_thresh)) & np.isfinite(depth_m) & (depth_m > 0) + if not np.any(mask): + return None + + try: + from scipy import ndimage + + labels, num = ndimage.label(mask) + best = None + best_area = 0 + for label in range(1, num + 1): + comp = labels == label + area = int(comp.sum()) + if min_area <= area <= max_area and area > best_area: + best = comp + best_area = area + if best is None: + return None + rows, cols = np.nonzero(best) + depths = depth_m[best] + except Exception: + rows, cols = np.nonzero(mask) + depths = depth_m[mask] + best_area = int(rows.size) + if not (min_area <= best_area <= max_area): + return None + + row = float(np.median(rows)) + col = float(np.median(cols)) + z = float(np.median(depths)) + return { + "pixel": [row, col], + "depth_m": z, + "area": best_area, + "xyz_cam": geom.backproject_pixel(K, col, row, z).tolist(), + } + + +def _save_debug_images( + *, + debug_dir: Path, + pose_idx: int, + camera: str, + rgb_open, + rgb_closed, + depth_m, +) -> dict[str, str]: + """Save open/closed/diff/depth images for diagnosing failed detections.""" + debug_dir.mkdir(parents=True, exist_ok=True) + prefix = debug_dir / f"{pose_idx:02d}_{camera}" + rgb_open = np.asarray(rgb_open, dtype=np.uint8) + rgb_closed = np.asarray(rgb_closed, dtype=np.uint8) + depth_m = np.asarray(depth_m, dtype=np.float32) + + diff = np.abs(rgb_open.astype(np.int16) - rgb_closed.astype(np.int16)).max(axis=2) + diff_img = np.clip(diff, 0, 255).astype(np.uint8) + + valid = np.isfinite(depth_m) & (depth_m > 0) + depth_img = np.zeros(depth_m.shape, dtype=np.uint8) + if np.any(valid): + lo, hi = np.percentile(depth_m[valid], [2, 98]) + if hi > lo: + depth_img[valid] = np.clip((depth_m[valid] - lo) / (hi - lo) * 255, 0, 255) + + paths = { + "open": str(prefix.with_name(prefix.name + "_open.png")), + "closed": str(prefix.with_name(prefix.name + "_closed.png")), + "diff": str(prefix.with_name(prefix.name + "_diff.png")), + "depth": str(prefix.with_name(prefix.name + "_depth.png")), + } + imageio.imwrite(paths["open"], rgb_open) + imageio.imwrite(paths["closed"], rgb_closed) + imageio.imwrite(paths["diff"], diff_img) + imageio.imwrite(paths["depth"], depth_img) + return paths + + +def _manual_pixel_from_terminal( + *, + image_path: str, + depth_m, + K, + patch_radius: int, +) -> dict | None: + """Prompt for a manual pixel and backproject it, or return None to skip.""" + print(f" manual fallback: inspect {image_path}") + print(" enter pixel as row,col (for example 240,320), or press Enter to skip") + text = input(" row,col> ").strip() + if not text: + return None + try: + row_s, col_s = text.replace(" ", "").split(",", 1) + row = int(round(float(row_s))) + col = int(round(float(col_s))) + except Exception: + print(" invalid pixel format; skipped") + return None + + z = geom.sample_depth_patch(depth_m, col, row, radius=patch_radius) + if not np.isfinite(z) or z <= 0: + print(f" no valid depth near ({row},{col}); skipped") + return None + p_cam = geom.backproject_pixel(K, col, row, z) + return { + "pixel": [float(row), float(col)], + "depth_m": float(z), + "area": int((2 * patch_radius + 1) ** 2), + "xyz_cam": p_cam.tolist(), + "manual": True, + } + + +def _detect_tip_pixel_by_motion( + rgb_open, + rgb_closed, + depth_m, + K, + *, + diff_thresh: int = 18, + min_area: int = 40, + max_area: int = 40000, +) -> dict | None: + """Detect gripper motion, preferring OpenCV but falling back gracefully.""" + try: + return geom.detect_tip_pixel_by_motion( + rgb_open, + rgb_closed, + depth_m, + K, + diff_thresh=diff_thresh, + min_area=min_area, + max_area=max_area, + ) + except ModuleNotFoundError as exc: + if exc.name != "cv2": + raise + return _detect_motion_blob_numpy( + rgb_open, + rgb_closed, + depth_m, + K, + diff_thresh=diff_thresh, + min_area=min_area, + max_area=max_area, + ) + + +def _parse_csv_floats(text: str, *, expected: int, name: str) -> tuple[float, ...]: + try: + values = tuple(float(part.strip()) for part in text.split(",") if part.strip()) + except ValueError as exc: + raise argparse.ArgumentTypeError(f"{name} must be comma-separated floats") from exc + if len(values) != expected: + raise argparse.ArgumentTypeError(f"{name} expects {expected} values, got {len(values)}") + return values + + +def _candidate_poses(args: argparse.Namespace) -> list[list[float]]: + xs = args.grid_x + ys = args.grid_y + zs = args.grid_z + poses = [[float(x), float(y), float(z)] for z in zs for y in ys for x in xs] + # Visit the center-ish pose first, then spread out. This makes early aborts + # less likely to leave the robot at a corner of the grid. + center = np.array([np.mean(xs), np.mean(ys), np.mean(zs)], dtype=np.float64) + poses.sort(key=lambda p: float(np.linalg.norm(np.asarray(p) - center))) + return poses + + +def _obs_camera(obs: dict, camera: str) -> tuple[np.ndarray, np.ndarray, dict]: + frames = obs.get("frames") or {} + depths = obs.get("depth") or {} + metas = obs.get("camera_meta") or {} + if camera not in frames: + raise RuntimeError(f"camera {camera!r} missing from observation frames") + if camera not in depths: + raise RuntimeError(f"camera {camera!r} missing from observation depth maps") + if camera not in metas: + raise RuntimeError(f"camera {camera!r} missing from observation metadata") + return ( + np.asarray(frames[camera], dtype=np.uint8), + np.asarray(depths[camera], dtype=np.float32), + dict(metas[camera]), + ) + + +def _tcp_point_from_pose(ee: dict, tcp_offset: tuple[float, float, float]) -> np.ndarray: + xyz = np.asarray(ee["xyz"], dtype=np.float64) + offset = np.asarray(tcp_offset, dtype=np.float64) + if np.allclose(offset, 0.0): + return xyz + from scipy.spatial.transform import Rotation as R + + quat = ee.get("quat_xyzw") + if quat is None: + raise RuntimeError("--tcp-offset requires get_ee_pose to return quat_xyzw") + return xyz + R.from_quat(np.asarray(quat, dtype=np.float64)).as_matrix() @ offset + + +def _detect_correspondence( + *, + client: SocketRpcClient, + camera: str, + tcp_offset: tuple[float, float, float], + diff_thresh: int, + min_area: int, + max_area: int, + settle_s: float, + detect_retries: int, + manual_on_fail: bool, + manual_always: bool, + manual_patch_radius: int, + debug_dir: Path, + pose_idx: int, +) -> dict[str, Any]: + """Toggle gripper once and return one cam/base correspondence.""" + det = None + rgb_closed = None + depth = None + meta = None + debug_paths: dict[str, str] = {} + attempts = max(1, int(detect_retries)) + for attempt in range(1, attempts + 1): + client.call("env.open_gripper", timeout_s=30.0) + time.sleep(settle_s) + obs_open = client.call("env.get_obs", timeout_s=30.0) + rgb_open, _, _ = _obs_camera(obs_open, camera) + + client.call("env.close_gripper", timeout_s=30.0) + time.sleep(settle_s) + obs_closed = client.call("env.get_obs", timeout_s=30.0) + rgb_closed, depth, meta = _obs_camera(obs_closed, camera) + + debug_paths = _save_debug_images( + debug_dir=debug_dir, + pose_idx=pose_idx * 10 + attempt, + camera=camera, + rgb_open=rgb_open, + rgb_closed=rgb_closed, + depth_m=depth, + ) + + if not manual_always: + det = _detect_tip_pixel_by_motion( + rgb_open, + rgb_closed, + depth, + np.asarray(meta["K"], dtype=np.float64), + diff_thresh=diff_thresh, + min_area=min_area, + max_area=max_area, + ) + if det is not None or manual_always: + break + + ee = client.call("env.get_ee_pose", timeout_s=15.0) + client.call("env.open_gripper", timeout_s=30.0) + + if (det is None or manual_always) and manual_on_fail: + assert rgb_closed is not None and depth is not None and meta is not None + det = _manual_pixel_from_terminal( + image_path=debug_paths.get("closed", ""), + depth_m=depth, + K=np.asarray(meta["K"], dtype=np.float64), + patch_radius=manual_patch_radius, + ) + if det is None: + raise RuntimeError( + "could not segment gripper motion in camera image; debug images: " + + json.dumps(debug_paths) + ) + + return { + "xyz_cam": np.asarray(det["xyz_cam"], dtype=np.float64), + "xyz_base": _tcp_point_from_pose(ee, tcp_offset), + "pixel": det["pixel"], + "depth_m": float(det["depth_m"]), + "area": int(det["area"]), + "manual": bool(det.get("manual", False)), + "debug_paths": debug_paths, + "ee": ee, + "camera_meta": meta, + } + + +def _calibrate_scene(args: argparse.Namespace) -> int: + if not args.yes and not args.no_save: + print( + "Refusing to move the robot without --yes. This calibration drives " + "the TCP through a small 3-D grid and toggles the gripper." + ) + return 2 + + client = SocketRpcClient(args.host, args.port) + meta_all = client.call("env.get_camera_meta", timeout_s=15.0) + if args.camera not in meta_all: + print(f"camera {args.camera!r} not available; found {sorted(meta_all)}") + return 2 + camera_meta = meta_all[args.camera] + serial = args.serial or camera_meta.get("serial") + if not serial: + print("could not determine camera serial; pass --serial") + return 2 + + poses = _candidate_poses(args) + print( + f"Calibrating fixed camera {args.camera!r} serial={serial} with up to " + f"{len(poses)} candidate poses; target valid points={args.n_points}." + ) + print("Clear the workspace. The gripper will move and open/close at each pose.") + + cam_pts: list[np.ndarray] = [] + base_pts: list[np.ndarray] = [] + records: list[dict[str, Any]] = [] + + try: + for idx, xyz in enumerate(poses, start=1): + if len(cam_pts) >= args.n_points: + break + print(f"\n[{idx}/{len(poses)}] move_to {np.round(xyz, 3).tolist()}") + move = client.call( + "env.move_to", + args=(xyz,), + kwargs={"euler_xyz": _DEFAULT_DOWN_EULER, "gripper": "open"}, + timeout_s=120.0, + ) + print(f" move: reached={move.get('reached')} err={move.get('pos_error_m')} final={move.get('final_xyz')}") + if move.get("error"): + print(f" skipped: {move['error']}") + continue + + try: + corr = _detect_correspondence( + client=client, + camera=args.camera, + tcp_offset=args.tcp_offset, + diff_thresh=args.diff_thresh, + min_area=args.min_area, + max_area=args.max_area, + settle_s=args.settle_s, + detect_retries=args.detect_retries, + manual_on_fail=args.manual_on_fail, + manual_always=args.manual_always, + manual_patch_radius=args.manual_patch_radius, + debug_dir=Path(args.debug_dir), + pose_idx=idx, + ) + except Exception as exc: + print(f" detection failed: {exc}") + continue + + cam_pts.append(corr["xyz_cam"]) + base_pts.append(corr["xyz_base"]) + records.append( + { + "target_xyz": xyz, + "pixel": corr["pixel"], + "depth_m": round(corr["depth_m"], 4), + "area": corr["area"], + "manual": corr["manual"], + "debug_paths": corr["debug_paths"], + "xyz_cam": np.round(corr["xyz_cam"], 5).tolist(), + "xyz_base": np.round(corr["xyz_base"], 5).tolist(), + "move": move, + } + ) + print( + " captured: " + f"pixel={np.round(corr['pixel'], 1).tolist()} " + f"depth={corr['depth_m']:.3f}m area={corr['area']} " + f"base={np.round(corr['xyz_base'], 3).tolist()} " + f"manual={corr['manual']}" + ) + finally: + try: + client.call("env.open_gripper", timeout_s=30.0) + except Exception: + pass + + if len(cam_pts) < 4: + print(f"Calibration failed: need >=4 correspondences, got {len(cam_pts)}") + print(json.dumps(records, indent=2, default=str)) + return 2 + + cam_arr = np.asarray(cam_pts, dtype=np.float64) + base_arr = np.asarray(base_pts, dtype=np.float64) + T_base_cam, rmse, inliers = geom.ransac_kabsch( + cam_arr, + base_arr, + thresh_m=args.ransac_thresh_m, + iters=args.ransac_iters, + min_inliers=min(4, len(cam_pts)), + seed=args.seed, + ) + n_inliers = int(inliers.sum()) + print( + f"\nFit complete: used={len(cam_pts)} inliers={n_inliers} " + f"RMSE={rmse * 1000:.1f} mm" + ) + if rmse > franka_calib.MAX_ACCEPTABLE_RMSE_M: + print( + "WARNING: RMSE exceeds the loader acceptance gate " + f"({franka_calib.MAX_ACCEPTABLE_RMSE_M * 1000:.0f} mm). " + "The server will reject this calibration unless you rerun with better detections." + ) + + result = { + "camera": args.camera, + "serial": serial, + "n_collected": len(cam_pts), + "n_inliers": n_inliers, + "rmse_m": rmse, + "T_base_cam": T_base_cam, + "records": records, + "inlier_mask": inliers.tolist(), + "tcp_offset": args.tcp_offset, + "method": "franka_markerless_gripper_motion", + } + + if args.no_save: + print("Not saved (--no-save). T_base_cam:") + print(json.dumps(result["T_base_cam"].tolist(), indent=2)) + return 0 + + path = franka_calib.save_scene_extrinsic( + serial, + T_base_cam, + K=np.asarray(camera_meta["K"], dtype=np.float64), + rmse_m=rmse, + num_points=len(cam_pts), + camera=args.camera, + n_inliers=n_inliers, + inlier_mask=inliers.tolist(), + method="franka_markerless_gripper_motion", + tcp_offset=args.tcp_offset, + correspondences=records, + ) + print(f"Saved T_base_cam -> {path}") + + reload_result = client.call( + "env.reload_camera_calibration", + kwargs={"camera": args.camera}, + timeout_s=15.0, + ) + print("Reload result:") + print(json.dumps(reload_result, indent=2, default=str)) + return 0 if rmse <= franka_calib.MAX_ACCEPTABLE_RMSE_M else 1 + + +def _explain_wrist() -> int: + print( + "Wrist camera calibration is hand-eye calibration (T_tcp_cam), not the " + "same fixed-camera problem as the scene camera. A reliable automated " + "script needs either a fixed fiducial/ChArUco/AprilTag board observed " + "from multiple wrist poses, or a scene-calibrated 3-D reference target. " + "This repository now supports loading/saving T_tcp_cam records, and " + "back_project(camera='wrist') will return panda_link0 xyz once such a " + "record exists." + ) + return 0 + + +def _build_argparser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, help="Franka env server transport port.") + parser.add_argument( + "--mode", + choices=["scene-auto", "wrist-info"], + default="scene-auto", + help="scene-auto calibrates fixed scene T_base_cam; wrist-info explains T_tcp_cam requirements.", + ) + parser.add_argument("--camera", default="scene", help="Camera name to calibrate (default: scene).") + parser.add_argument("--serial", default=None, help="Override saved RealSense serial.") + parser.add_argument("--n-points", type=int, default=8, help="Valid correspondences to collect.") + parser.add_argument("--yes", action="store_true", help="Confirm robot motion.") + parser.add_argument("--no-save", action="store_true", help="Fit but do not save/reload calibration.") + parser.add_argument("--self-test", action="store_true", help="Run offline math/detection self-test.") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--grid-x", type=lambda s: _parse_csv_floats(s, expected=3, name="grid-x"), default=_DEFAULT_GRID_X) + parser.add_argument("--grid-y", type=lambda s: _parse_csv_floats(s, expected=3, name="grid-y"), default=_DEFAULT_GRID_Y) + parser.add_argument("--grid-z", type=lambda s: _parse_csv_floats(s, expected=2, name="grid-z"), default=_DEFAULT_GRID_Z) + parser.add_argument( + "--tcp-offset", + type=lambda s: _parse_csv_floats(s, expected=3, name="tcp-offset"), + default=(0.0, 0.0, 0.0), + help="Optional calibration point offset in TCP frame, meters (default: 0,0,0).", + ) + parser.add_argument("--settle-s", type=float, default=0.4, help="Wait after gripper open/close captures.") + parser.add_argument("--diff-thresh", type=int, default=10) + parser.add_argument("--min-area", type=int, default=12) + parser.add_argument("--max-area", type=int, default=40000) + parser.add_argument("--detect-retries", type=int, default=2, + help="Open/close detection attempts per pose before fallback/skip.") + parser.add_argument("--manual-on-fail", action="store_true", + help="When auto detection fails, prompt for row,col on the saved closed image.") + parser.add_argument("--manual-always", action="store_true", + help="Always prompt for row,col instead of using auto detection.") + parser.add_argument("--manual-patch-radius", type=int, default=3, + help="Depth median patch radius for manually clicked pixels.") + parser.add_argument("--debug-dir", default="/tmp/franka_camera_calib_debug", + help="Directory for per-pose open/closed/diff/depth debug images.") + parser.add_argument("--ransac-thresh-m", type=float, default=0.02) + parser.add_argument("--ransac-iters", type=int, default=500) + return parser + + +def main() -> int: + args = _build_argparser().parse_args() + if args.manual_always: + args.manual_on_fail = True + if args.self_test: + return _self_test() + if args.mode == "wrist-info": + return _explain_wrist() + if args.port is None: + raise SystemExit("--port is required unless --self-test or --mode wrist-info") + if args.n_points < 4: + raise SystemExit("--n-points must be >= 4") + return _calibrate_scene(args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/robots/franka/calibrate_charuco_wrist.py b/robots/franka/calibrate_charuco_wrist.py new file mode 100644 index 00000000..93ee472e --- /dev/null +++ b/robots/franka/calibrate_charuco_wrist.py @@ -0,0 +1,448 @@ +#!/usr/bin/env python3 +"""ChArUco-based wrist-camera hand-eye calibration for Franka. + +This estimates ``T_tcp_cam`` for the wrist camera. The ChArUco board must be +fixed in the scene while the wrist camera observes it from multiple robot poses. +The script detects the board in the wrist camera, reads the live TCP pose, then +uses OpenCV hand-eye calibration to solve camera-in-TCP. + +Typical workflow: + +1. Start the Franka env server. +2. Place the printed ChArUco board flat and rigid on the table. +3. Move the wrist camera so the board is visible in the wrist image. +4. Check detection: + + python robots/franka/calibrate_charuco_wrist.py --port 5599 check + +5. Calibrate with a small automatic orbit around the current pose: + + python robots/franka/calibrate_charuco_wrist.py --port 5599 calibrate --yes + +The scene camera is different: a ChArUco board gives ``T_scene_cam_board`` but +not ``T_base_scene_cam`` unless the board pose in ``panda_link0`` is known. Use +``auto_calibrate_cameras.py`` for markerless scene-to-base calibration, or add a +known board pose / touch-corner workflow for scene calibration. +""" +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path +from typing import Any + +import imageio.v2 as imageio +import numpy as np +from scipy.spatial.transform import Rotation as R + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from robots.franka import calibration as franka_calib # noqa: E402 +from robots.lerobot import geometry as geom # noqa: E402 +from rpent.utils.socket_rpc import SocketRpcClient # noqa: E402 + +_DEFAULT_BOARD_SPEC = ( + _REPO_ROOT / "resources" / "franka" / "calibration_boards" / "franka_charuco_7x5_25mm.json" +) +_DEFAULT_OFFSETS = ( + ([0.0, 0.0, 0.0], [0.0, 0.0, 0.0]), + ([0.025, 0.0, 0.0], [0.0, 0.0, 12.0]), + ([-0.025, 0.0, 0.0], [0.0, 0.0, -12.0]), + ([0.0, 0.025, 0.0], [0.0, 10.0, 0.0]), + ([0.0, -0.025, 0.0], [0.0, -10.0, 0.0]), + ([0.0, 0.0, 0.025], [8.0, 0.0, 0.0]), + ([0.0, 0.0, -0.015], [-8.0, 0.0, 0.0]), + ([0.02, 0.02, 0.015], [6.0, -8.0, 8.0]), + ([-0.02, -0.02, 0.015], [-6.0, 8.0, -8.0]), +) + + +def _require_cv2_aruco(): + try: + import cv2 + except ModuleNotFoundError as exc: + raise SystemExit( + "Missing dependency: cv2. Install with `uv pip install -e .[calibration]` " + "or run with an environment that has opencv-contrib-python-headless." + ) from exc + if not hasattr(cv2, "aruco"): + raise SystemExit("Installed cv2 lacks aruco; install opencv-contrib-python-headless.") + return cv2 + + +def _load_board_spec(path: Path) -> dict[str, Any]: + spec = json.loads(path.read_text()) + required = ["dictionary", "squares_x", "squares_y", "square_length_m", "marker_length_m"] + missing = [key for key in required if key not in spec] + if missing: + raise SystemExit(f"board spec missing keys: {missing}") + return spec + + +def _aruco_dictionary(cv2, dictionary_name: str): + aruco = cv2.aruco + key = dictionary_name.upper() + if not key.startswith("DICT_"): + key = f"DICT_{key}" + if not hasattr(aruco, key): + raise SystemExit(f"OpenCV does not know ArUco dictionary {dictionary_name!r}") + return aruco.getPredefinedDictionary(getattr(aruco, key)) + + +def _make_board(cv2, spec: dict[str, Any]): + aruco = cv2.aruco + dictionary = _aruco_dictionary(cv2, spec["dictionary"]) + squares = (int(spec["squares_x"]), int(spec["squares_y"])) + square = float(spec["square_length_m"]) + marker = float(spec["marker_length_m"]) + try: + return aruco.CharucoBoard(squares, square, marker, dictionary) + except Exception: + return aruco.CharucoBoard_create(squares[0], squares[1], square, marker, dictionary) + + +def _pose_to_matrix(xyz, quat_xyzw) -> np.ndarray: + T = np.eye(4, dtype=np.float64) + T[:3, :3] = R.from_quat(np.asarray(quat_xyzw, dtype=np.float64)).as_matrix() + T[:3, 3] = np.asarray(xyz, dtype=np.float64) + return T + + +def _pose_from_rvec_tvec(cv2, rvec, tvec) -> np.ndarray: + T = np.eye(4, dtype=np.float64) + R_cam_board, _ = cv2.Rodrigues(np.asarray(rvec, dtype=np.float64)) + T[:3, :3] = R_cam_board + T[:3, 3] = np.asarray(tvec, dtype=np.float64).reshape(3) + return T + + +def _detect_charuco_pose(cv2, image, K, dist_coeffs, board) -> dict | None: + aruco = cv2.aruco + gray = cv2.cvtColor(np.asarray(image, dtype=np.uint8), cv2.COLOR_RGB2GRAY) + marker_corners = [] + marker_ids = None + + if hasattr(aruco, "CharucoDetector"): + detector = aruco.CharucoDetector(board) + charuco_corners, charuco_ids, marker_corners, marker_ids = detector.detectBoard(gray) + count = 0 if charuco_ids is None else len(charuco_ids) + else: + params = aruco.DetectorParameters() + try: + detector = aruco.ArucoDetector(board.getDictionary(), params) + marker_corners, marker_ids, _ = detector.detectMarkers(gray) + except Exception: + marker_corners, marker_ids, _ = aruco.detectMarkers( + gray, board.getDictionary(), parameters=params + ) + if marker_ids is None or len(marker_ids) < 2: + return None + + try: + count, charuco_corners, charuco_ids = aruco.interpolateCornersCharuco( + marker_corners, + marker_ids, + gray, + board, + cameraMatrix=K, + distCoeffs=dist_coeffs, + ) + except TypeError: + count, charuco_corners, charuco_ids = aruco.interpolateCornersCharuco( + marker_corners, marker_ids, gray, board, K, dist_coeffs + ) + + if charuco_ids is None or int(count) < 6: + return None + + if hasattr(aruco, "estimatePoseCharucoBoard"): + rvec = np.zeros((3, 1), dtype=np.float64) + tvec = np.zeros((3, 1), dtype=np.float64) + ok, rvec, tvec = aruco.estimatePoseCharucoBoard( + charuco_corners, + charuco_ids, + board, + K, + dist_coeffs, + rvec, + tvec, + ) + else: + obj_points, img_points = board.matchImagePoints(charuco_corners, charuco_ids) + ok, rvec, tvec = cv2.solvePnP( + obj_points, + img_points, + K, + dist_coeffs, + flags=cv2.SOLVEPNP_ITERATIVE, + ) + if not ok: + return None + return { + "T_cam_board": _pose_from_rvec_tvec(cv2, rvec, tvec), + "n_markers": 0 if marker_ids is None else int(len(marker_ids)), + "n_corners": int(count), + "rvec": np.asarray(rvec, dtype=np.float64).reshape(3).tolist(), + "tvec": np.asarray(tvec, dtype=np.float64).reshape(3).tolist(), + } + + +def _draw_detection(cv2, image, K, dist_coeffs, board, out_path: Path) -> None: + aruco = cv2.aruco + canvas = np.asarray(image, dtype=np.uint8).copy() + gray = cv2.cvtColor(canvas, cv2.COLOR_RGB2GRAY) + charuco_corners = None + charuco_ids = None + if hasattr(aruco, "CharucoDetector"): + detector = aruco.CharucoDetector(board) + charuco_corners, charuco_ids, marker_corners, marker_ids = detector.detectBoard(gray) + else: + params = aruco.DetectorParameters() + try: + detector = aruco.ArucoDetector(board.getDictionary(), params) + marker_corners, marker_ids, _ = detector.detectMarkers(gray) + except Exception: + marker_corners, marker_ids, _ = aruco.detectMarkers( + gray, board.getDictionary(), parameters=params + ) + if marker_ids is not None and len(marker_ids) > 0: + aruco.drawDetectedMarkers(canvas, marker_corners, marker_ids) + if charuco_ids is not None and len(charuco_ids) > 0: + aruco.drawDetectedCornersCharuco(canvas, charuco_corners, charuco_ids) + out_path.parent.mkdir(parents=True, exist_ok=True) + imageio.imwrite(out_path, canvas) + + +def _capture(client: SocketRpcClient, camera: str) -> tuple[dict, np.ndarray, dict]: + obs = client.call("env.get_obs", timeout_s=30.0) + frames = obs.get("frames") or {} + meta = obs.get("camera_meta") or {} + if camera not in frames: + raise RuntimeError(f"camera {camera!r} missing from observation") + if camera not in meta: + raise RuntimeError(f"camera {camera!r} metadata missing") + ee = client.call("env.get_ee_pose", timeout_s=15.0) + return ee, np.asarray(frames[camera], dtype=np.uint8), dict(meta[camera]) + + +def _check(args: argparse.Namespace) -> int: + cv2 = _require_cv2_aruco() + board = _make_board(cv2, _load_board_spec(args.board_spec)) + client = SocketRpcClient(args.host, args.port) + ee, image, meta = _capture(client, args.camera) + K = np.asarray(meta["K"], dtype=np.float64) + dist = np.asarray(meta.get("dist_coeffs") or np.zeros(5), dtype=np.float64) + det = _detect_charuco_pose(cv2, image, K, dist, board) + out_path = Path(args.debug_dir) / f"{args.camera}_charuco_check.png" + _draw_detection(cv2, image, K, dist, board, out_path) + print(json.dumps({ + "camera": args.camera, + "serial": meta.get("serial"), + "debug_image": str(out_path), + "detected": det is not None, + "n_markers": None if det is None else det["n_markers"], + "n_corners": None if det is None else det["n_corners"], + "tcp_xyz": ee.get("xyz"), + }, indent=2)) + return 0 if det is not None else 2 + + +def _target_pose(start: dict, dxyz, drpy_deg) -> tuple[list[float], list[float]]: + xyz = np.asarray(start["xyz"], dtype=np.float64) + np.asarray(dxyz, dtype=np.float64) + euler = np.asarray(start["euler_xyz"], dtype=np.float64) + np.radians(np.asarray(drpy_deg, dtype=np.float64)) + return xyz.tolist(), euler.tolist() + + +def _collect_samples( + args: argparse.Namespace, +) -> tuple[ + list[np.ndarray], + list[np.ndarray], + list[np.ndarray], + list[np.ndarray], + list[dict], +]: + cv2 = _require_cv2_aruco() + board = _make_board(cv2, _load_board_spec(args.board_spec)) + client = SocketRpcClient(args.host, args.port) + start = client.call("env.get_ee_pose", timeout_s=15.0) + start_xyz = start["xyz"] + start_quat = start["quat_xyzw"] + + R_gripper2base: list[np.ndarray] = [] + t_gripper2base: list[np.ndarray] = [] + R_target2cam: list[np.ndarray] = [] + t_target2cam: list[np.ndarray] = [] + records: list[dict] = [] + + try: + for idx, (dxyz, drpy) in enumerate(_DEFAULT_OFFSETS, start=1): + if len(records) >= args.n_samples: + break + target_xyz, target_euler = _target_pose(start, dxyz, drpy) + print(f"\n[{idx}] target_xyz={np.round(target_xyz, 3).tolist()} drpy={drpy}") + move = client.call( + "env.move_to", + args=(target_xyz,), + kwargs={"euler_xyz": target_euler, "gripper": None}, + timeout_s=120.0, + ) + print(f" move: reached={move.get('reached')} err={move.get('pos_error_m')} final={move.get('final_xyz')}") + time.sleep(args.settle_s) + ee, image, meta = _capture(client, args.camera) + K = np.asarray(meta["K"], dtype=np.float64) + dist = np.asarray(meta.get("dist_coeffs") or np.zeros(5), dtype=np.float64) + det = _detect_charuco_pose(cv2, image, K, dist, board) + debug_path = Path(args.debug_dir) / f"{args.camera}_charuco_{idx:02d}.png" + _draw_detection(cv2, image, K, dist, board, debug_path) + if det is None: + print(f" detection failed (debug: {debug_path})") + continue + + T_base_tcp = _pose_to_matrix(ee["xyz"], ee["quat_xyzw"]) + T_cam_board = det["T_cam_board"] + R_gripper2base.append(T_base_tcp[:3, :3]) + t_gripper2base.append(T_base_tcp[:3, 3].reshape(3, 1)) + R_target2cam.append(T_cam_board[:3, :3]) + t_target2cam.append(T_cam_board[:3, 3].reshape(3, 1)) + records.append({ + "idx": idx, + "target_xyz": target_xyz, + "target_euler": target_euler, + "tcp_xyz": ee["xyz"], + "n_markers": det["n_markers"], + "n_corners": det["n_corners"], + "tvec": det["tvec"], + "debug_image": str(debug_path), + "move": move, + }) + print(f" captured: markers={det['n_markers']} corners={det['n_corners']} tvec={np.round(det['tvec'], 3).tolist()}") + finally: + try: + client.call( + "env.move_to", + args=(start_xyz,), + kwargs={"quat_xyzw": start_quat, "gripper": None}, + timeout_s=120.0, + ) + except Exception as exc: + print(f"warning: failed to return to start pose: {exc}") + + return ( + [np.asarray(r, dtype=np.float64) for r in R_gripper2base], + [np.asarray(t, dtype=np.float64) for t in t_gripper2base], + [np.asarray(r, dtype=np.float64) for r in R_target2cam], + [np.asarray(t, dtype=np.float64) for t in t_target2cam], + records, + ) + + +def _calibrate(args: argparse.Namespace) -> int: + if not args.yes: + print("Refusing to move the robot without --yes.") + return 2 + cv2 = _require_cv2_aruco() + board_spec = _load_board_spec(args.board_spec) + client = SocketRpcClient(args.host, args.port) + meta = client.call("env.get_camera_meta", timeout_s=15.0).get(args.camera) + if not meta: + print(f"camera {args.camera!r} not available") + return 2 + serial = args.serial or meta.get("serial") + if not serial: + print("could not determine wrist camera serial; pass --serial") + return 2 + + R_g2b, t_g2b, R_t2c, t_t2c, records = _collect_samples(args) + if len(records) < 5: + print(f"Need at least 5 valid board detections; got {len(records)}") + return 2 + + R_cam2tcp, t_cam2tcp = cv2.calibrateHandEye( + R_g2b, + t_g2b, + R_t2c, + t_t2c, + method=cv2.CALIB_HAND_EYE_TSAI, + ) + T_tcp_cam = np.eye(4, dtype=np.float64) + T_tcp_cam[:3, :3] = np.asarray(R_cam2tcp, dtype=np.float64) + T_tcp_cam[:3, 3] = np.asarray(t_cam2tcp, dtype=np.float64).reshape(3) + + residuals = [] + target_points_base = [] + for Rb, tb, Rc, tc in zip(R_g2b, t_g2b, R_t2c, t_t2c): + T_base_tcp = np.eye(4) + T_base_tcp[:3, :3] = Rb + T_base_tcp[:3, 3] = tb.reshape(3) + T_cam_board = np.eye(4) + T_cam_board[:3, :3] = Rc + T_cam_board[:3, 3] = tc.reshape(3) + T_base_board = T_base_tcp @ T_tcp_cam @ T_cam_board + target_points_base.append(T_base_board[:3, 3]) + target_points_base = np.asarray(target_points_base) + center = target_points_base.mean(axis=0) + residuals = np.linalg.norm(target_points_base - center, axis=1) + rmse = float(np.sqrt(np.mean(np.square(residuals)))) + + print(f"\nHand-eye fit: samples={len(records)} board-position RMSE={rmse * 1000:.1f} mm") + print("T_tcp_cam:") + print(json.dumps(T_tcp_cam.tolist(), indent=2)) + + if args.no_save: + return 0 + + path = franka_calib.save_wrist_extrinsic( + serial, + T_tcp_cam, + K=np.asarray(meta["K"], dtype=np.float64), + rmse_m=rmse, + num_points=len(records), + camera=args.camera, + method="charuco_hand_eye_tsai", + board_spec=board_spec, + records=records, + board_position_residuals_m=residuals.tolist(), + ) + print(f"Saved T_tcp_cam -> {path}") + reload_result = client.call("env.reload_camera_calibration", kwargs={"camera": args.camera}, timeout_s=15.0) + print("Reload result:") + print(json.dumps(reload_result, indent=2, default=str)) + return 0 if rmse <= franka_calib.MAX_ACCEPTABLE_RMSE_M else 1 + + +def _build_argparser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, required=True) + parser.add_argument("--camera", default="wrist") + parser.add_argument("--serial", default=None) + parser.add_argument("--board-spec", type=Path, default=_DEFAULT_BOARD_SPEC) + parser.add_argument("--debug-dir", type=Path, default=Path("/tmp/franka_charuco_wrist_debug")) + sub = parser.add_subparsers(dest="cmd", required=True) + + sub.add_parser("check", help="Capture one wrist frame and report board detection.") + + calib = sub.add_parser("calibrate", help="Move through a small orbit and solve T_tcp_cam.") + calib.add_argument("--yes", action="store_true") + calib.add_argument("--no-save", action="store_true") + calib.add_argument("--n-samples", type=int, default=8) + calib.add_argument("--settle-s", type=float, default=0.5) + return parser + + +def main() -> int: + args = _build_argparser().parse_args() + if args.cmd == "check": + return _check(args) + if args.cmd == "calibrate": + return _calibrate(args) + raise SystemExit(f"unknown command: {args.cmd}") + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/robots/franka/calibration.py b/robots/franka/calibration.py new file mode 100644 index 00000000..192045c3 --- /dev/null +++ b/robots/franka/calibration.py @@ -0,0 +1,120 @@ +"""Franka camera calibration loading helpers. + +Calibration records are stored per RealSense serial under +``~/.cache/rpent/franka/camera_calibration``. Supported records: + +- fixed scene camera: ``{"T_base_cam": [[...]], ...}`` +- wrist camera: ``{"T_tcp_cam": [[...]], ...}`` + +``T_base_cam`` maps camera-frame points into ``panda_link0``. ``T_tcp_cam`` maps +wrist-camera points into the live TCP frame; the env server composes it with the +current ``T_base_tcp`` for each observation. +""" +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +import numpy as np + +_CALIB_DIR = "~/.cache/rpent/franka/camera_calibration" +MAX_ACCEPTABLE_RMSE_M = 0.02 + + +def calib_path(serial: str) -> Path: + """Return the on-disk calibration path for a RealSense serial.""" + return Path(os.path.expanduser(_CALIB_DIR)) / f"{serial}.json" + + +def load_record(serial: str) -> dict | None: + """Load a camera calibration record, converting known transforms to arrays.""" + path = calib_path(serial) + if not path.is_file(): + return None + with open(path) as f: + data = json.load(f) + for key in ("T_base_cam", "T_tcp_cam"): + if key in data and data[key] is not None: + data[key] = np.asarray(data[key], dtype=np.float64) + data.setdefault("serial", str(serial)) + data["path"] = str(path) + return data + + +def _jsonable(value: Any) -> Any: + """Convert numpy values to JSON-compatible Python values.""" + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, dict): + return {key: _jsonable(val) for key, val in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(val) for val in value] + return value + + +def save_record(serial: str, **fields: Any) -> Path: + """Save a calibration record for ``serial``. + + Known transform fields are ``T_base_cam`` for fixed cameras and + ``T_tcp_cam`` for wrist cameras. Extra diagnostics are preserved. + """ + path = calib_path(serial) + path.parent.mkdir(parents=True, exist_ok=True) + data = {"serial": str(serial), **fields} + with open(path, "w") as f: + json.dump(_jsonable(data), f, indent=2) + return path + + +def save_scene_extrinsic( + serial: str, + T_base_cam, + *, + K=None, + rmse_m: float | None = None, + num_points: int | None = None, + **extra: Any, +) -> Path: + """Persist a fixed-camera ``T_base_cam`` calibration.""" + fields: dict[str, Any] = {"T_base_cam": np.asarray(T_base_cam, dtype=np.float64)} + if K is not None: + fields["K"] = np.asarray(K, dtype=np.float64) + if rmse_m is not None: + fields["rmse_m"] = float(rmse_m) + if num_points is not None: + fields["num_points"] = int(num_points) + fields.update(extra) + return save_record(serial, **fields) + + +def save_wrist_extrinsic( + serial: str, + T_tcp_cam, + *, + K=None, + rmse_m: float | None = None, + num_points: int | None = None, + **extra: Any, +) -> Path: + """Persist a wrist-camera ``T_tcp_cam`` hand-eye calibration.""" + fields: dict[str, Any] = {"T_tcp_cam": np.asarray(T_tcp_cam, dtype=np.float64)} + if K is not None: + fields["K"] = np.asarray(K, dtype=np.float64) + if rmse_m is not None: + fields["rmse_m"] = float(rmse_m) + if num_points is not None: + fields["num_points"] = int(num_points) + fields.update(extra) + return save_record(serial, **fields) + + +def is_accepted(record: dict | None) -> bool: + """Return whether a record exists and passes the saved RMSE gate.""" + if record is None: + return False + rmse = record.get("rmse_m") + return rmse is None or float(rmse) <= MAX_ACCEPTABLE_RMSE_M diff --git a/robots/franka/env_client.py b/robots/franka/env_client.py new file mode 100644 index 00000000..9b0170e3 --- /dev/null +++ b/robots/franka/env_client.py @@ -0,0 +1,113 @@ +"""Franka env client that forwards calls over the driver RPC boundary.""" +from __future__ import annotations + +from typing import Any + +import numpy as np + +from rpent.utils.rpc import RpcClient + + +_TIMEOUT_S = { + "default": 30.0, + "env.reset": 180.0, + "env.get_spec": 30.0, + "env.get_obs": 60.0, + "env.get_ee_pose": 30.0, + "env.get_camera_meta": 30.0, + "env.move_to": 120.0, + "env.move_delta": 90.0, + "env.open_gripper": 30.0, + "env.close_gripper": 30.0, +} + + +class FrankaEnvClient: + """Remote stub for the standalone Franka env server protocol.""" + + def __init__(self, client: RpcClient): + self._client = client + self._spec: dict | None = None + + def reset(self) -> tuple[dict, Any]: + """Clear errors, drive the arm to home, and return ``(obs, info)``.""" + return self._client.call("env.reset", timeout_s=_TIMEOUT_S["env.reset"]) + + def get_spec(self) -> dict: + """Return and cache the env's static self-description.""" + if self._spec is None: + self._spec = self._client.call( + "env.get_spec", timeout_s=_TIMEOUT_S["env.get_spec"] + ) + return self._spec + + def get_obs(self) -> dict: + """Fetch the current observation without moving the arm.""" + return self._client.call("env.get_obs", timeout_s=_TIMEOUT_S["env.get_obs"]) + + def get_ee_pose(self) -> dict: + """Return the current TCP pose in the Franka base frame.""" + return self._client.call( + "env.get_ee_pose", timeout_s=_TIMEOUT_S["env.get_ee_pose"] + ) + + def get_camera_meta(self) -> dict: + """Return live RGB-D camera intrinsics and base-frame extrinsics.""" + return self._client.call( + "env.get_camera_meta", timeout_s=_TIMEOUT_S["env.get_camera_meta"] + ) + + def move_to( + self, + xyz, + *, + yaw_deg: float | None = None, + gripper: str | None = None, + ) -> dict: + """Move to a base-frame Cartesian target. + + The agent-facing API intentionally hides arbitrary quaternions. When + ``yaw_deg`` is provided, the driver receives a down-facing euler target + with the requested yaw; otherwise it preserves the current orientation. + """ + xyz = np.asarray(xyz, dtype=float).reshape(-1)[:3].tolist() + kwargs: dict[str, Any] = {"gripper": gripper} + if yaw_deg is not None: + kwargs["euler_xyz"] = [float(np.pi), 0.0, float(np.radians(yaw_deg))] + return self._client.call( + "env.move_to", + args=(xyz,), + kwargs=kwargs, + timeout_s=_TIMEOUT_S["env.move_to"], + ) + + def move_delta( + self, + *, + dxyz=None, + yaw_delta_deg: float | None = None, + gripper: str | None = None, + ) -> dict: + """Nudge the TCP by relative translation and optional yaw.""" + kwargs: dict[str, Any] = {"gripper": gripper} + if dxyz is not None: + kwargs["dxyz"] = np.asarray(dxyz, dtype=float).reshape(-1)[:3].tolist() + if yaw_delta_deg is not None: + kwargs["drpy_deg"] = [0.0, 0.0, float(yaw_delta_deg)] + return self._client.call( + "env.move_delta", + kwargs=kwargs, + timeout_s=_TIMEOUT_S["env.move_delta"], + ) + + def open_gripper(self) -> dict: + """Open the Franka Hand.""" + return self._client.call( + "env.open_gripper", timeout_s=_TIMEOUT_S["env.open_gripper"] + ) + + def close_gripper(self) -> dict: + """Close/grasp with the Franka Hand.""" + return self._client.call( + "env.close_gripper", timeout_s=_TIMEOUT_S["env.close_gripper"] + ) diff --git a/robots/franka/env_server.py b/robots/franka/env_server.py new file mode 100644 index 00000000..03e4c15c --- /dev/null +++ b/robots/franka/env_server.py @@ -0,0 +1,943 @@ +"""Standalone Franka (FR3) env host for the LLM-in-the-loop agent (env-only). + +Drives a Franka arm through the SERL cartesian-impedance ROS controller and the +``franka_gripper`` action topics, exposing agent-friendly *Cartesian* primitives +(``reset`` / ``get_obs`` / ``get_ee_pose`` / ``move_to`` / ``move_delta`` / +``open_gripper`` / ``close_gripper`` / ``get_spec``) over an RPC server +(:class:`rpent.utils.rpc.RpcFacade`, http transport by default) -- the same +wire protocol the LIBERO and LeRobot drivers use, so the agent side talks to all +three identically. + +Unlike the LIBERO driver, this server does **not** import any +``rlinf.envs.realworld`` module. Importing that package runs node-level ROS setup +side effects at import time (it kills any running ``roscore`` / ``rosmaster``) and +pulls in the Ray-based ``Worker`` stack. This driver instead talks to ROS directly +with ``rospy``, reusing only the *recipe* from +``rlinf.envs.realworld.franka.{franka_controller,franka_env}``: the impedance +controller channel names, the ``roslaunch`` bring-up, the ``franka_gripper`` +action messages, and the safety-box + pose-interpolation logic. + +Run it inside the RLinf ``.venv`` with the ``serl_franka_controllers`` catkin +workspace sourced (see ``robots/franka/run_env_server.sh``):: + + source /home/franka/franka/RLinf/.venv/franka_catkin_ws/devel/setup.bash + /home/franka/franka/RLinf/.venv/bin/python robots/franka/env_server.py \ + --output-dir /tmp/franka_run --robot-ip 172.16.0.2 + +Hardware defaults match the current bench: an FR3 at ``172.16.0.2`` with the +Franka Hand, plus two Intel RealSense D435I cameras. Every default is overridable +from the CLI. +""" +from __future__ import annotations + +import argparse +import os +import signal +import sys +import threading +import time +from pathlib import Path +from typing import Any, Optional + +import numpy as np +from scipy.spatial.transform import Rotation as R +from scipy.spatial.transform import Slerp + +# Make ``rpent`` importable when this file is run from the RLinf .venv +# (which need not have rpent installed) -- the source tree is enough. +_PHYSICALAGENT_ROOT = Path(__file__).resolve().parents[2] +if str(_PHYSICALAGENT_ROOT) not in sys.path: + sys.path.insert(0, str(_PHYSICALAGENT_ROOT)) + +from robots.franka import calibration as camera_calib # noqa: E402 +from rpent.utils.logging import get_logger, init_output_dir # noqa: E402 +from rpent.utils.rpc import RpcFacade # noqa: E402 + +logger = get_logger("franka_driver") + + +# --------------------------------------------------------------------------- +# Bench defaults (override from the CLI) +# --------------------------------------------------------------------------- + +_DEFAULT_ROBOT_IP = os.environ.get("FRANKA_ROBOT_IP", "172.16.0.2") +_DEFAULT_ROS_PKG = "serl_franka_controllers" + +# Two Intel RealSense D435I on the current bench. Order matters: the first camera +# is the primary/overview view (maps to the ToolResult ``_image_bytes`` slot), the +# second maps to ``_image_cam_bytes``. Names are what the agent sees. +_DEFAULT_CAMERAS: tuple[tuple[str, str], ...] = ( + ("scene", "142122070838"), + ("wrist", "141722078696"), +) +_CAM_WIDTH = 640 +_CAM_HEIGHT = 480 +_CAM_FPS = 30 + +# Conservative tabletop workspace box in the base (panda_link0) frame, meters. +# ``move_to`` clips targets to this so an LLM-supplied coordinate cannot drive the +# arm into the table / out of reach. Tuned around the collect-data reset pose +# ([0.5, 0, 0.1] with the gripper pointing down); widen/lower on your bench once +# you have confirmed the table height in the base frame. +_WORKSPACE_MIN = np.array([0.30, -0.50, 0.00], dtype=np.float64) +_WORKSPACE_MAX = np.array([1.10, 0.50, 0.50], dtype=np.float64) + +# Default "home" pose the agent's reset() drives to: above the table, gripper +# pointing straight down. Orientation is euler xyz (radians); rx = -pi points the +# Franka Hand down (matches the RLinf realworld collect configs). +_RESET_XYZ = np.array([0.50, 0.0, 0.25], dtype=np.float64) +_RESET_EULER = np.array([np.pi, 0.0, 0.0], dtype=np.float64) + +# The nominal orientation the arm holds; move_to keeps the current orientation +# unless the caller overrides it, and clips any requested orientation to a window +# around this so the wrist cannot flip into a self-collision. +_TARGET_EULER = _RESET_EULER.copy() +_EULER_WINDOW = np.array([0.6, 0.6, np.pi], dtype=np.float64) # +/- rad per axis + +# Motion smoothness / safety. The impedance controller tracks a streamed sequence +# of equilibrium poses; capping the per-step Cartesian and angular deltas and +# pacing at ``step_frequency`` keeps motions slow and gentle. +_STEP_FREQUENCY = 10.0 # Hz (equilibrium-pose publish rate) +_MAX_STEP_M = 0.01 # max Cartesian move per streamed setpoint (=> ~0.1 m/s) +_MAX_STEP_DEG = 5.0 # max orientation change per streamed setpoint +_MAX_MOVE_M = 0.60 # hard cap on a single move_to path length (safety) +_REACHED_TOL_M = 0.003 # move_to "reached" tolerance +_SETTLE_TIMEOUT_S = 2.0 # max time to hold the final setpoint while settling + +# franka_gripper widths (meters). +_GRIPPER_OPEN_WIDTH = 0.09 +_GRIPPER_GRASP_WIDTH = 0.01 +_GRIPPER_GRASP_FORCE = 130.0 +_GRIPPER_SPEED = 0.3 + + +# --------------------------------------------------------------------------- +# small pose helpers (scipy uses scalar-last quaternions: [x, y, z, w]) +# --------------------------------------------------------------------------- + + +def _euler_to_quat(euler_xyz) -> np.ndarray: + return R.from_euler("xyz", np.asarray(euler_xyz, dtype=np.float64)).as_quat() + + +def _quat_to_euler(quat_xyzw) -> np.ndarray: + return R.from_quat(np.asarray(quat_xyzw, dtype=np.float64)).as_euler("xyz") + + +def _pose_to_matrix(pose7: np.ndarray) -> np.ndarray: + """Convert ``[x, y, z, qx, qy, qz, qw]`` to ``T_base_tcp``.""" + pose7 = np.asarray(pose7, dtype=np.float64).reshape(-1)[:7] + T = np.eye(4, dtype=np.float64) + T[:3, :3] = R.from_quat(pose7[3:]).as_matrix() + T[:3, 3] = pose7[:3] + return T + + +def _clip_euler_window(euler_xyz: np.ndarray) -> np.ndarray: + """Clip an euler orientation to +/-``_EULER_WINDOW`` around ``_TARGET_EULER``. + + Wraps each axis into ``[-pi, pi]`` relative to the target first so the clip is + on the shortest angular distance, not the raw value. + """ + euler = np.asarray(euler_xyz, dtype=np.float64).copy() + delta = (euler - _TARGET_EULER + np.pi) % (2 * np.pi) - np.pi + delta = np.clip(delta, -_EULER_WINDOW, _EULER_WINDOW) + return _TARGET_EULER + delta + + +# --------------------------------------------------------------------------- +# RealSense color + depth camera +# --------------------------------------------------------------------------- + + +class RealSenseDepthCamera: + """Minimal RealSense RGB-D grabber with depth aligned to color.""" + + def __init__(self, name: str, serial: str, *, width: int, height: int, fps: int): + import pyrealsense2 as rs + + self.name = name + self.serial = str(serial) + self._rs = rs + self._pipeline = rs.pipeline() + cfg = rs.config() + cfg.enable_device(self.serial) + cfg.enable_stream(rs.stream.color, width, height, rs.format.rgb8, fps) + cfg.enable_stream(rs.stream.depth, width, height, rs.format.z16, fps) + self._profile = self._pipeline.start(cfg) + self._align = rs.align(rs.stream.color) + + depth_sensor = self._profile.get_device().first_depth_sensor() + self._depth_scale = float(depth_sensor.get_depth_scale()) + + color_stream = self._profile.get_stream( + rs.stream.color + ).as_video_stream_profile() + intr = color_stream.get_intrinsics() + self._K = np.array( + [[intr.fx, 0.0, intr.ppx], [0.0, intr.fy, intr.ppy], [0.0, 0.0, 1.0]], + dtype=np.float64, + ) + self._dist_coeffs = np.asarray(intr.coeffs, dtype=np.float64) + self._distortion_model = str(intr.model) + self._width = int(width) + self._height = int(height) + self._calib = camera_calib.load_record(self.serial) + self._calib_accepted = camera_calib.is_accepted(self._calib) + + # Drop the first few frames so auto-exposure settles. + for _ in range(5): + try: + self._pipeline.wait_for_frames(2000) + except Exception: + break + calib_kind = "uncalibrated" + if self._calib_accepted: + if self._calib and "T_base_cam" in self._calib: + calib_kind = "T_base_cam" + elif self._calib and "T_tcp_cam" in self._calib: + calib_kind = "T_tcp_cam" + elif self._calib is not None: + calib_kind = "rejected" + logger.info( + "camera '%s' (serial %s) RGB-D started; calibration=%s", + name, + self.serial, + calib_kind, + ) + + def read(self) -> tuple[np.ndarray, np.ndarray]: + """Return ``(rgb_uint8, depth_m_float32)`` with depth aligned to RGB.""" + frames = self._pipeline.wait_for_frames(2000) + frames = self._align.process(frames) + color = frames.get_color_frame() + depth = frames.get_depth_frame() + if not color or not depth: + raise RuntimeError(f"camera '{self.name}' incomplete frameset") + rgb = np.ascontiguousarray(np.asanyarray(color.get_data()), dtype=np.uint8) + depth_raw = np.asanyarray(depth.get_data()) + depth_m = np.ascontiguousarray( + depth_raw.astype(np.float32) * self._depth_scale + ) + return rgb, depth_m + + def meta(self, *, T_base_tcp: np.ndarray | None = None) -> dict: + """Return JSON-able camera metadata and any base-frame extrinsic.""" + T_base_cam = None + calibration_kind = None + if self._calib_accepted and self._calib is not None: + if "T_base_cam" in self._calib: + T_base_cam = self._calib["T_base_cam"] + calibration_kind = "T_base_cam" + elif "T_tcp_cam" in self._calib: + calibration_kind = "T_tcp_cam" + if T_base_tcp is not None: + T_base_cam = np.asarray(T_base_tcp, dtype=np.float64) @ self._calib[ + "T_tcp_cam" + ] + + return { + "name": self.name, + "serial": self.serial, + "frame": f"{self.name}_camera", + "width": self._width, + "height": self._height, + "K": self._K.tolist(), + "dist_coeffs": self._dist_coeffs.tolist(), + "distortion_model": self._distortion_model, + "depth_scale": self._depth_scale, + "calibrated": T_base_cam is not None, + "calibration_kind": calibration_kind, + "calibration_rmse_m": ( + None if self._calib is None else self._calib.get("rmse_m") + ), + "calibration_path": ( + None if self._calib is None else self._calib.get("path") + ), + "T_base_cam": None if T_base_cam is None else T_base_cam.tolist(), + } + + def reload_calibration(self) -> dict: + """Reload this camera's calibration record from disk.""" + self._calib = camera_calib.load_record(self.serial) + self._calib_accepted = camera_calib.is_accepted(self._calib) + meta = self.meta() + logger.info( + "camera '%s' calibration reloaded: calibrated=%s kind=%s", + self.name, + meta["calibrated"], + meta["calibration_kind"], + ) + return meta + + def close(self) -> None: + try: + self._pipeline.stop() + except Exception: + pass + + +# --------------------------------------------------------------------------- +# Franka ROS backend (plain rospy; no rlinf, no Ray) +# --------------------------------------------------------------------------- + + +class FrankaRobotBackend: + """Low-level Franka arm + gripper over ROS. + + Reuses the channel names, impedance ``roslaunch`` bring-up, and message + parsing from ``rlinf.envs.realworld.franka.franka_controller`` / + ``.common.gripper.franka_gripper`` -- reimplemented on plain ``rospy`` so this + driver stays free of the Ray ``Worker`` stack and the destructive + ``rlinf.envs.realworld`` import-time side effects. + """ + + _ARM_EQUILIBRIUM = "/cartesian_impedance_controller/equilibrium_pose" + _ARM_STATE = "/franka_state_controller/franka_states" + _ARM_RESET = "/franka_control/error_recovery/goal" + _GRIPPER_MOVE = "/franka_gripper/move/goal" + _GRIPPER_GRASP = "/franka_gripper/grasp/goal" + _GRIPPER_STATE = "/franka_gripper/joint_states" + + def __init__( + self, + *, + robot_ip: str, + ros_pkg: str = _DEFAULT_ROS_PKG, + load_gripper: bool = True, + node_name: str = "franka_agent_driver", + state_timeout_s: float = 4.0, + launch_timeout_s: float = 40.0, + ): + # Lazy ROS imports so the module can be imported (e.g. for --help) on a + # host without ROS on the path. + import geometry_msgs.msg as geom_msg + import rospy + from franka_gripper.msg import GraspActionGoal, MoveActionGoal + from franka_msgs.msg import ErrorRecoveryActionGoal, FrankaState + from sensor_msgs.msg import JointState + + self._rospy = rospy + self._geom_msg = geom_msg + self._FrankaState = FrankaState + self._ErrorRecoveryActionGoal = ErrorRecoveryActionGoal + self._MoveActionGoal = MoveActionGoal + self._GraspActionGoal = GraspActionGoal + self._JointState = JointState + + self._robot_ip = robot_ip + self._ros_pkg = ros_pkg + self._load_gripper = load_gripper + self._impedance_proc = None # only set if we launch it ourselves + + # Live state (updated by subscriber callbacks). + self._tcp_pose = np.zeros(7, dtype=np.float64) + self._tcp_pose[6] = 1.0 # unit quaternion + self._tcp_force = np.zeros(3, dtype=np.float64) + self._tcp_torque = np.zeros(3, dtype=np.float64) + self._state_seen = threading.Event() + self._gripper_width = _GRIPPER_OPEN_WIDTH + self._gripper_open = True + self._gripper_seen = threading.Event() + + self._ensure_master() + rospy.init_node(node_name, anonymous=True, disable_signals=True) + + # Publishers. + self._pub_equilibrium = rospy.Publisher( + self._ARM_EQUILIBRIUM, geom_msg.PoseStamped, queue_size=10 + ) + self._pub_reset = rospy.Publisher( + self._ARM_RESET, ErrorRecoveryActionGoal, queue_size=1 + ) + self._pub_gripper_move = rospy.Publisher( + self._GRIPPER_MOVE, MoveActionGoal, queue_size=1 + ) + self._pub_gripper_grasp = rospy.Publisher( + self._GRIPPER_GRASP, GraspActionGoal, queue_size=1 + ) + + # Subscribers. + rospy.Subscriber(self._ARM_STATE, FrankaState, self._on_state_msg) + rospy.Subscriber(self._GRIPPER_STATE, JointState, self._on_gripper_msg) + + self._ensure_impedance(state_timeout_s, launch_timeout_s) + + # -- lifecycle helpers ------------------------------------------------- + + def _ensure_master(self) -> None: + """Make sure a ROS master is reachable; launch ``roscore`` if not. + + Mirrors ``rlinf...common.ros.ros_controller.ROSController``: reuse a + running master, otherwise start one. ``rosmaster --core`` (what + ``roscore`` spawns) also counts as a running master. + """ + import psutil + + try: + import rosgraph + + if rosgraph.is_master_online(): + logger.info("ROS master already online at %s", os.environ.get( + "ROS_MASTER_URI", "http://localhost:11311")) + return + except Exception: + pass + + for proc in psutil.process_iter(["name"]): + if proc.info.get("name") in ("roscore", "rosmaster"): + logger.info("found running %s (pid %s)", proc.info["name"], proc.pid) + return + + logger.info("no ROS master found; starting `roscore`") + self._roscore_proc = psutil.Popen( + ["roscore"], stdout=sys.stdout, stderr=sys.stdout + ) + time.sleep(2.0) + + def _impedance_up(self, timeout_s: float) -> bool: + """Return True once a franka_states message arrives within ``timeout_s``.""" + return self._state_seen.wait(timeout=timeout_s) + + def _ensure_impedance(self, state_timeout_s: float, launch_timeout_s: float) -> None: + """Attach to a running cartesian-impedance controller, or launch one. + + We never launch a second controller: if ``franka_states`` is already + publishing we simply attach. Only when no state arrives do we + ``roslaunch serl_franka_controllers impedance.launch``. + """ + if self._impedance_up(state_timeout_s): + logger.info("attached to live cartesian-impedance controller") + return + + import psutil + + load_gripper = "true" if self._load_gripper else "false" + cmd = [ + "roslaunch", + self._ros_pkg, + "impedance.launch", + f"robot_ip:={self._robot_ip}", + f"load_gripper:={load_gripper}", + ] + logger.info("no live controller detected; starting `%s`", " ".join(cmd)) + self._impedance_proc = psutil.Popen(cmd, stdout=sys.stdout, stderr=sys.stdout) + + deadline = time.time() + launch_timeout_s + while time.time() < deadline: + if self._impedance_proc.poll() is not None: + raise RuntimeError( + "impedance roslaunch exited before becoming ready; run " + f"`{' '.join(cmd)}` manually to see the error (is " + f"{self._ros_pkg} on the ROS package path?)." + ) + if self._impedance_up(1.0): + logger.info("cartesian-impedance controller is up") + return + raise RuntimeError( + f"cartesian-impedance controller not ready after {launch_timeout_s}s" + ) + + # -- ROS callbacks ----------------------------------------------------- + + def _on_state_msg(self, msg) -> None: + tmatrix = np.array(list(msg.O_T_EE)).reshape(4, 4).T + quat = R.from_matrix(tmatrix[:3, :3].copy()).as_quat() + self._tcp_pose = np.concatenate([tmatrix[:3, 3], quat]) + self._tcp_force = np.array(list(msg.K_F_ext_hat_K)[:3]) + self._tcp_torque = np.array(list(msg.K_F_ext_hat_K)[3:]) + self._state_seen.set() + + def _on_gripper_msg(self, msg) -> None: + # joint_states reports both finger joints; their sum is the opening width. + self._gripper_width = float(np.sum(msg.position)) + self._gripper_open = self._gripper_width > 0.06 + self._gripper_seen.set() + + # -- arm --------------------------------------------------------------- + + def get_tcp_pose(self) -> np.ndarray: + """Current TCP pose ``[x, y, z, qx, qy, qz, qw]`` in the base frame.""" + return self._tcp_pose.copy() + + def get_state(self) -> dict: + return { + "tcp_pose": self._tcp_pose.copy(), + "tcp_force": self._tcp_force.copy(), + "tcp_torque": self._tcp_torque.copy(), + "gripper_width": self._gripper_width, + "gripper_open": self._gripper_open, + } + + def move_arm(self, pose7: np.ndarray) -> None: + """Publish one equilibrium pose ``[x, y, z, qx, qy, qz, qw]``.""" + pose7 = np.asarray(pose7, dtype=np.float64).reshape(-1) + assert pose7.shape[0] == 7, f"expected 7-D pose, got {pose7.shape}" + msg = self._geom_msg.PoseStamped() + msg.header.frame_id = "0" + msg.header.stamp = self._rospy.Time.now() + msg.pose.position = self._geom_msg.Point(pose7[0], pose7[1], pose7[2]) + msg.pose.orientation = self._geom_msg.Quaternion( + pose7[3], pose7[4], pose7[5], pose7[6] + ) + self._pub_equilibrium.publish(msg) + + def clear_errors(self) -> None: + self._pub_reset.publish(self._ErrorRecoveryActionGoal()) + + # -- gripper ----------------------------------------------------------- + + def open_gripper(self, speed: float = _GRIPPER_SPEED) -> None: + msg = self._MoveActionGoal() + msg.goal.width = _GRIPPER_OPEN_WIDTH + msg.goal.speed = speed + self._pub_gripper_move.publish(msg) + self._gripper_open = True + + def close_gripper( + self, speed: float = _GRIPPER_SPEED, force: float = _GRIPPER_GRASP_FORCE + ) -> None: + msg = self._GraspActionGoal() + msg.goal.width = _GRIPPER_GRASP_WIDTH + msg.goal.speed = speed + msg.goal.epsilon.inner = 1.0 + msg.goal.epsilon.outer = 1.0 + msg.goal.force = force + self._pub_gripper_grasp.publish(msg) + self._gripper_open = False + + def shutdown(self) -> None: + """Terminate only the controller we launched (never a pre-existing one).""" + if self._impedance_proc is not None and self._impedance_proc.poll() is None: + logger.info("terminating impedance controller we launched") + self._impedance_proc.terminate() + try: + self._impedance_proc.wait(timeout=10) + except Exception: + self._impedance_proc.kill() + + +# --------------------------------------------------------------------------- +# Agent-facing env facade +# --------------------------------------------------------------------------- + + +class FrankaAgentEnv: + """Cartesian-primitive facade the agent RPCs into. + + Observation:: + + {"state": {"tcp_xyz": (3,) float, # meters, base frame + "tcp_quat": (4,) float, # [x, y, z, w] + "tcp_euler": (3,) float, # radians xyz + "gripper_width": float, # meters + "gripper_open": bool}, + "frames": {: (H, W, 3) uint8, ...}} + + All values are plain numpy / python scalars so they pickle across the RPC + wire (the agent process does not import torch or ROS). + """ + + def __init__( + self, + backend: FrankaRobotBackend, + cameras: list[RealSenseDepthCamera], + *, + workspace_min: np.ndarray = _WORKSPACE_MIN, + workspace_max: np.ndarray = _WORKSPACE_MAX, + reset_xyz: np.ndarray = _RESET_XYZ, + reset_euler: np.ndarray = _RESET_EULER, + step_frequency: float = _STEP_FREQUENCY, + max_step_m: float = _MAX_STEP_M, + max_step_deg: float = _MAX_STEP_DEG, + max_move_m: float = _MAX_MOVE_M, + settle_timeout_s: float = _SETTLE_TIMEOUT_S, + ): + self._backend = backend + self._cameras = cameras + self._workspace_min = np.asarray(workspace_min, dtype=np.float64) + self._workspace_max = np.asarray(workspace_max, dtype=np.float64) + self._reset_xyz = np.asarray(reset_xyz, dtype=np.float64) + self._reset_quat = _euler_to_quat(reset_euler) + self._step_frequency = float(step_frequency) + self._max_step_m = float(max_step_m) + self._max_step_deg = float(max_step_deg) + self._max_move_m = float(max_move_m) + self._settle_timeout_s = float(settle_timeout_s) + + # -- observation ------------------------------------------------------- + + def _frames(self) -> dict[str, np.ndarray]: + out: dict[str, np.ndarray] = {} + for cam in self._cameras: + try: + rgb, _ = cam.read() + out[cam.name] = rgb + except Exception as e: + logger.warning("camera '%s' frame grab failed: %s", cam.name, e) + return out + + def _camera_observation( + self, T_base_tcp: np.ndarray + ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray], dict[str, dict]]: + frames: dict[str, np.ndarray] = {} + depths: dict[str, np.ndarray] = {} + meta: dict[str, dict] = {} + for cam in self._cameras: + try: + rgb, depth = cam.read() + frames[cam.name] = rgb + depths[cam.name] = depth + meta[cam.name] = cam.meta(T_base_tcp=T_base_tcp) + except Exception as e: + logger.warning("camera '%s' RGB-D grab failed: %s", cam.name, e) + return frames, depths, meta + + def _make_obs(self) -> dict: + pose = self._backend.get_tcp_pose() + st = self._backend.get_state() + frames, depths, camera_meta = self._camera_observation(_pose_to_matrix(pose)) + return { + "state": { + "tcp_xyz": pose[:3].astype(np.float64), + "tcp_quat": pose[3:].astype(np.float64), + "tcp_euler": _quat_to_euler(pose[3:]).astype(np.float64), + "gripper_width": float(st["gripper_width"]), + "gripper_open": bool(st["gripper_open"]), + }, + "frames": frames, + "depth": depths, + "camera_meta": camera_meta, + } + + def get_obs(self) -> dict: + """Return the current observation without moving the arm.""" + return self._make_obs() + + def get_ee_pose(self) -> dict: + """Return the current end-effector pose in the base frame.""" + pose = self._backend.get_tcp_pose() + return { + "xyz": [round(float(v), 4) for v in pose[:3]], + "quat_xyzw": [round(float(v), 4) for v in pose[3:]], + "euler_xyz": [round(float(v), 4) for v in _quat_to_euler(pose[3:])], + "frame": "panda_link0", + } + + def get_spec(self) -> dict: + """Static self-description for the agent side.""" + return { + "world_frame": "panda_link0", + "control": "cartesian_impedance", + "position_unit": "meters", + "orientation": "euler_xyz_radians (also quat xyzw)", + "workspace_min": [round(float(v), 3) for v in self._workspace_min], + "workspace_max": [round(float(v), 3) for v in self._workspace_max], + "camera_names": [c.name for c in self._cameras], + "depth_camera_names": [c.name for c in self._cameras], + "preferred_backproject_camera": "wrist", + "gripper": "binary (open_gripper / close_gripper); width in meters", + "reset_xyz": [round(float(v), 3) for v in self._reset_xyz], + } + + def get_camera_meta(self) -> dict: + """Return live per-camera intrinsics and base-frame extrinsics if calibrated.""" + T_base_tcp = _pose_to_matrix(self._backend.get_tcp_pose()) + return {cam.name: cam.meta(T_base_tcp=T_base_tcp) for cam in self._cameras} + + def reload_camera_calibration(self, camera: str | None = None) -> dict: + """Reload camera calibration records from disk.""" + requested = None if camera is None else str(camera) + reloaded: dict[str, dict] = {} + for cam in self._cameras: + if requested is not None and cam.name != requested: + continue + reloaded[cam.name] = cam.reload_calibration() + if requested is not None and requested not in reloaded: + return { + "error": f"unknown camera {requested!r}", + "available_cameras": [cam.name for cam in self._cameras], + } + return reloaded + + # -- motion ------------------------------------------------------------ + + def _clip_xyz(self, xyz: np.ndarray) -> tuple[np.ndarray, bool]: + clipped = np.clip(xyz, self._workspace_min, self._workspace_max) + return clipped, bool(np.any(clipped != xyz)) + + def _stream_to(self, target_xyz: np.ndarray, target_quat: np.ndarray) -> None: + """Stream interpolated equilibrium poses from the current pose to the + target, capping per-step Cartesian and angular deltas and pacing at + ``step_frequency`` so the impedance controller tracks a slow, smooth path. + """ + cur = self._backend.get_tcp_pose() + cur_xyz, cur_quat = cur[:3], cur[3:] + + dist = float(np.linalg.norm(target_xyz - cur_xyz)) + ang = float( + (R.from_quat(cur_quat).inv() * R.from_quat(target_quat)).magnitude() + ) + n_pos = int(np.ceil(dist / self._max_step_m)) if dist > 0 else 0 + n_ang = int(np.ceil(np.degrees(ang) / self._max_step_deg)) if ang > 0 else 0 + n = max(1, n_pos, n_ang) + + slerp = Slerp([0.0, 1.0], R.concatenate([R.from_quat(cur_quat), R.from_quat(target_quat)])) + period = 1.0 / self._step_frequency + for i in range(1, n + 1): + t = i / n + pos = cur_xyz + (target_xyz - cur_xyz) * t + quat = slerp([t])[0].as_quat() + self._backend.move_arm(np.concatenate([pos, quat])) + time.sleep(period) + + def _wait_until_reached( + self, target_xyz: np.ndarray, target_quat: np.ndarray + ) -> tuple[np.ndarray, float, bool]: + """Hold the final setpoint until the measured TCP pose is close enough.""" + target_pose = np.concatenate([target_xyz, target_quat]) + deadline = time.time() + self._settle_timeout_s + period = 1.0 / self._step_frequency + while True: + final = self._backend.get_tcp_pose() + pos_err = float(np.linalg.norm(final[:3] - target_xyz)) + if pos_err <= _REACHED_TOL_M or time.time() >= deadline: + return final, pos_err, pos_err <= _REACHED_TOL_M + self._backend.move_arm(target_pose) + time.sleep(period) + + def _apply_gripper(self, gripper: Optional[str]) -> None: + if gripper is None: + return + g = str(gripper).lower() + if g in ("open", "release"): + self._backend.open_gripper() + time.sleep(0.6) + elif g in ("close", "grasp"): + self._backend.close_gripper() + time.sleep(0.6) + else: + raise ValueError(f"gripper must be 'open' or 'close', got {gripper!r}") + + def move_to( + self, + xyz, + *, + euler_xyz=None, + quat_xyzw=None, + gripper: Optional[str] = None, + ) -> dict: + """Move the TCP to a base-frame ``xyz`` (meters), holding the current + orientation unless ``euler_xyz`` / ``quat_xyzw`` is given. + + The target is clipped to the workspace box and the orientation to a safe + window; the path is streamed as slow capped setpoints. Optionally set the + gripper ("open"/"close") first. Returns a log dict. + """ + self._backend.clear_errors() + cur = self._backend.get_tcp_pose() + + target_xyz = np.asarray(xyz, dtype=np.float64).reshape(-1)[:3] + target_xyz, clipped = self._clip_xyz(target_xyz) + + path_len = float(np.linalg.norm(target_xyz - cur[:3])) + if path_len > self._max_move_m: + return { + "reached": False, + "error": ( + f"requested move of {path_len:.3f} m exceeds the {self._max_move_m} m " + "single-move safety cap; issue smaller moves." + ), + "current_xyz": [round(float(v), 4) for v in cur[:3]], + } + + if quat_xyzw is not None: + target_quat = np.asarray(quat_xyzw, dtype=np.float64).reshape(-1)[:4] + target_quat = _euler_to_quat(_clip_euler_window(_quat_to_euler(target_quat))) + elif euler_xyz is not None: + target_quat = _euler_to_quat(_clip_euler_window(euler_xyz)) + else: + target_quat = cur[3:] + + self._apply_gripper(gripper) + self._stream_to(target_xyz, target_quat) + final, pos_err, reached = self._wait_until_reached(target_xyz, target_quat) + return { + "reached": reached, + "pos_error_m": round(pos_err, 4), + "clipped_to_workspace": clipped, + "target_xyz": [round(float(v), 4) for v in target_xyz], + "final_xyz": [round(float(v), 4) for v in final[:3]], + "final_euler": [round(float(v), 4) for v in _quat_to_euler(final[3:])], + "gripper_open": bool(self._backend.get_state()["gripper_open"]), + } + + def move_delta( + self, + *, + dxyz=None, + drpy_deg=None, + gripper: Optional[str] = None, + ) -> dict: + """Nudge the TCP by a relative ``dxyz`` (meters) and/or ``drpy_deg`` + (degrees, applied in the base frame), for fine alignment. + """ + cur = self._backend.get_tcp_pose() + target_xyz = cur[:3].copy() + if dxyz is not None: + target_xyz = target_xyz + np.asarray(dxyz, dtype=np.float64).reshape(-1)[:3] + + euler = _quat_to_euler(cur[3:]) + if drpy_deg is not None: + euler = euler + np.radians(np.asarray(drpy_deg, dtype=np.float64).reshape(-1)[:3]) + + return self.move_to( + target_xyz, euler_xyz=euler, gripper=gripper + ) + + def open_gripper(self) -> dict: + self._backend.open_gripper() + time.sleep(0.6) + return {"gripper_open": True, "gripper_width": self._backend.get_state()["gripper_width"]} + + def close_gripper(self) -> dict: + self._backend.close_gripper() + time.sleep(0.6) + return {"gripper_open": False, "gripper_width": self._backend.get_state()["gripper_width"]} + + # -- reset / teardown -------------------------------------------------- + + def reset(self) -> tuple[dict, dict]: + """Clear errors and drive the arm to its home pose; return ``(obs, {})``.""" + self._backend.clear_errors() + self._stream_to(self._reset_xyz, self._reset_quat) + time.sleep(0.5) + return self._make_obs(), {} + + def close(self) -> None: + for cam in self._cameras: + cam.close() + self._backend.shutdown() + + +# --------------------------------------------------------------------------- +# RPC dispatcher + parent watchdog (mirrors the LIBERO / LeRobot drivers) +# --------------------------------------------------------------------------- + + +class FrankaEnvFacade(RpcFacade): + """Serve :class:`FrankaAgentEnv` over the RPC boundary. + + ``shutdown`` and ``healthz`` are handled by :class:`RpcFacade` and the + transport; this only routes ``env.*`` methods to the env instance. + """ + + def __init__(self, env: FrankaAgentEnv): + super().__init__() + self._env = env + + def _dispatch(self, method: str, args: tuple, kwargs: dict): + if method.startswith("env."): + attr = method[len("env."):] + try: + return getattr(self._env, attr)(*args, **kwargs) + except Exception as e: + logger.warning("env method %s failed: %s", method, e) + raise + raise ValueError(f"unknown RPC method: {method!r}") + + def request_shutdown(self) -> None: + """Signal the serve loop to exit (used by the signal handlers).""" + self._shutdown_event.set() + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def _parse_cameras(specs: list[str] | None) -> list[tuple[str, str]]: + if not specs: + return list(_DEFAULT_CAMERAS) + out: list[tuple[str, str]] = [] + for spec in specs: + name, _, serial = spec.partition(":") + if not name or not serial: + raise ValueError(f"--camera expects name:serial, got {spec!r}") + out.append((name, serial)) + return out + + +def _build_cameras(specs: list[tuple[str, str]]) -> list[RealSenseDepthCamera]: + cams: list[RealSenseDepthCamera] = [] + for name, serial in specs: + try: + cams.append( + RealSenseDepthCamera( + name, serial, width=_CAM_WIDTH, height=_CAM_HEIGHT, fps=_CAM_FPS + ) + ) + except Exception as e: + logger.warning("could not open camera '%s' (serial %s): %s", name, serial, e) + return cams + + +def main() -> int: + p = argparse.ArgumentParser(description="Standalone Franka env server") + p.add_argument("--output-dir", type=str, required=True) + p.add_argument("--robot-ip", type=str, default=_DEFAULT_ROBOT_IP) + p.add_argument("--ros-pkg", type=str, default=_DEFAULT_ROS_PKG) + p.add_argument("--no-gripper", action="store_true", help="Do not load the Franka Hand.") + p.add_argument( + "--camera", action="append", default=None, + help="Camera as name:serial (repeatable). Defaults to the two bench D435I.", + ) + p.add_argument("--transport", choices=["socket", "http"], default="http", + help="RPC transport (default http). Socket (pickle) is also " + "available for the numpy obs payloads.") + p.add_argument("--host", type=str, default="127.0.0.1") + p.add_argument("--port", type=int, default=0, + help="RPC port. 0 asks the OS for a free port.") + args = p.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + init_output_dir(args.output_dir) + logger.info( + "starting Franka env server: robot_ip=%s output_dir=%s", + args.robot_ip, args.output_dir, + ) + + backend = FrankaRobotBackend( + robot_ip=args.robot_ip, + ros_pkg=args.ros_pkg, + load_gripper=not args.no_gripper, + ) + cameras = _build_cameras(_parse_cameras(args.camera)) + env = FrankaAgentEnv(backend, cameras) + + facade = FrankaEnvFacade(env) + + # Release the robot + cameras cleanly on SIGTERM / SIGINT (ProcessDaemon + # stop, Ctrl-C). The handler only flags shutdown; ``env.close()`` runs in + # the ``finally`` below. RpcFacade also exits on parent death (stdin EOF) + # and on the ``shutdown`` RPC. + def _handle_signal(signum, _frame): + logger.warning( + "received %s; releasing robot and shutting down", + signal.Signals(signum).name, + ) + facade.request_shutdown() + + for _sig in (signal.SIGTERM, signal.SIGINT): + signal.signal(_sig, _handle_signal) + + try: + facade.serve(transport=args.transport, host=args.host, port=args.port) + finally: + env.close() + logger.info("driver exited cleanly") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/robots/franka/generate_fiducial_board.py b/robots/franka/generate_fiducial_board.py new file mode 100644 index 00000000..63f72ac8 --- /dev/null +++ b/robots/franka/generate_fiducial_board.py @@ -0,0 +1,190 @@ +#!/usr/bin/env python3 +"""Generate a printable ChArUco fiducial board for Franka camera calibration. + +The output files are: + +- ``.png``: high-resolution board image with DPI metadata. +- ``.pdf``: printable single-page PDF at the same physical scale. +- ``.json``: board spec consumed by future calibration scripts. + +Print the PDF at 100% / actual size, not "fit to page". After printing, measure +one square and use the measured value if it differs from the requested size. + +Dependency: + uv pip install opencv-contrib-python-headless +""" +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_DEFAULT_OUT_DIR = _REPO_ROOT / "resources" / "franka" / "calibration_boards" + + +def _require_cv2_aruco(): + try: + import cv2 + except ModuleNotFoundError as exc: + raise SystemExit( + "Missing dependency: cv2. Install with:\n" + " uv pip install opencv-contrib-python-headless\n" + "or:\n" + " pip install opencv-contrib-python-headless" + ) from exc + if not hasattr(cv2, "aruco"): + raise SystemExit( + "Installed cv2 lacks the aruco module. Install opencv-contrib, not plain opencv:\n" + " uv pip install opencv-contrib-python-headless" + ) + return cv2 + + +def _aruco_dictionary(cv2, name: str): + aruco = cv2.aruco + key = f"DICT_{name.upper()}" if not name.upper().startswith("DICT_") else name.upper() + if not hasattr(aruco, key): + available = sorted(attr[5:] for attr in dir(aruco) if attr.startswith("DICT_")) + raise SystemExit(f"Unknown ArUco dictionary {name!r}. Available examples: {available[:12]}") + return aruco.getPredefinedDictionary(getattr(aruco, key)), key + + +def _make_charuco_board( + cv2, + *, + squares_x: int, + squares_y: int, + square_px: int, + marker_px: int, + dictionary, +): + aruco = cv2.aruco + try: + return aruco.CharucoBoard( + (int(squares_x), int(squares_y)), + float(square_px), + float(marker_px), + dictionary, + ) + except Exception: + return aruco.CharucoBoard_create( + int(squares_x), + int(squares_y), + float(square_px), + float(marker_px), + dictionary, + ) + + +def _draw_board(board, image_size: tuple[int, int], margin_px: int) -> np.ndarray: + """Draw a ChArUco board across OpenCV API versions.""" + width, height = image_size + if hasattr(board, "generateImage"): + img = board.generateImage((width, height), marginSize=int(margin_px), borderBits=1) + else: + img = board.draw((width, height), marginSize=int(margin_px), borderBits=1) + img = np.asarray(img, dtype=np.uint8) + if img.ndim == 3: + img = img[:, :, 0] + return img + + +def _save_outputs( + image: np.ndarray, + *, + out_dir: Path, + name: str, + dpi: int, + spec: dict[str, Any], +) -> dict[str, str]: + out_dir.mkdir(parents=True, exist_ok=True) + pil = Image.fromarray(image, mode="L") + png_path = out_dir / f"{name}.png" + pdf_path = out_dir / f"{name}.pdf" + json_path = out_dir / f"{name}.json" + pil.save(png_path, dpi=(dpi, dpi)) + pil.save(pdf_path, "PDF", resolution=float(dpi)) + json_path.write_text(json.dumps(spec, indent=2)) + return {"png": str(png_path), "pdf": str(pdf_path), "json": str(json_path)} + + +def _build_argparser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--out-dir", type=Path, default=_DEFAULT_OUT_DIR) + parser.add_argument("--name", default="franka_charuco_7x5_25mm") + parser.add_argument("--squares-x", type=int, default=7) + parser.add_argument("--squares-y", type=int, default=5) + parser.add_argument("--square-mm", type=float, default=25.0) + parser.add_argument("--marker-mm", type=float, default=18.0) + parser.add_argument("--margin-mm", type=float, default=12.0) + parser.add_argument("--dpi", type=int, default=300) + parser.add_argument("--dictionary", default="4X4_50", help="OpenCV ArUco dictionary suffix, e.g. 4X4_50 or APRILTAG_36h11.") + parser.add_argument("--check-deps", action="store_true", help="Only check for cv2.aruco and exit.") + return parser + + +def main() -> int: + args = _build_argparser().parse_args() + cv2 = _require_cv2_aruco() + if args.check_deps: + print(f"cv2 {cv2.__version__} with aruco: OK") + return 0 + + if args.squares_x < 2 or args.squares_y < 2: + raise SystemExit("--squares-x and --squares-y must be >= 2") + if not (0 < args.marker_mm < args.square_mm): + raise SystemExit("--marker-mm must be >0 and < --square-mm") + + px_per_mm = float(args.dpi) / 25.4 + square_px = int(round(args.square_mm * px_per_mm)) + marker_px = int(round(args.marker_mm * px_per_mm)) + margin_px = int(round(args.margin_mm * px_per_mm)) + board_width_px = args.squares_x * square_px + board_height_px = args.squares_y * square_px + image_size = (board_width_px + 2 * margin_px, board_height_px + 2 * margin_px) + + dictionary, dictionary_name = _aruco_dictionary(cv2, args.dictionary) + board = _make_charuco_board( + cv2, + squares_x=args.squares_x, + squares_y=args.squares_y, + square_px=square_px, + marker_px=marker_px, + dictionary=dictionary, + ) + image = _draw_board(board, image_size, margin_px) + + spec = { + "type": "charuco", + "dictionary": dictionary_name, + "squares_x": args.squares_x, + "squares_y": args.squares_y, + "square_length_m": args.square_mm / 1000.0, + "marker_length_m": args.marker_mm / 1000.0, + "margin_m": args.margin_mm / 1000.0, + "dpi": args.dpi, + "image_width_px": int(image_size[0]), + "image_height_px": int(image_size[1]), + "print_instructions": [ + "Print the PDF at 100% / actual size.", + "Disable fit-to-page or scaling.", + "Use matte paper and mount it flat to cardboard/foam board.", + "Measure one printed square and update square_length_m if needed.", + ], + } + paths = _save_outputs(image, out_dir=args.out_dir, name=args.name, dpi=args.dpi, spec=spec) + + physical_w_mm = image_size[0] / px_per_mm + physical_h_mm = image_size[1] / px_per_mm + print(json.dumps({"paths": paths, "physical_size_mm": [round(physical_w_mm, 1), round(physical_h_mm, 1)], "spec": spec}, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/robots/franka/prompt.py b/robots/franka/prompt.py new file mode 100644 index 00000000..3f2ef9c1 --- /dev/null +++ b/robots/franka/prompt.py @@ -0,0 +1,106 @@ +"""Franka prompt fragments and assembly.""" +from __future__ import annotations + +from rpent.context.prompt_utils import BulletList, Numbered, PromptNode +from rpent.context.prompts import prompt as base_prompt + +PREAMBLE = """ +You are a physical agent controlling a Franka arm through tools. Observe the scene through camera images and robot state, reason in the robot base frame, and command small safe Cartesian motions. Under Claude Code / Codex the tools may appear namespaced as mcp__rpent__; call the names shown in your tool list. +""" + +GOAL = """ +Accomplish the user's manipulation task on the real Franka setup. +""" + +ENVIRONMENT = BulletList([ + """ + Robot: Franka arm with Franka Hand. World frame is panda_link0, units are meters. Positive x is front (relative to the robot base), y is left, z is up. Use get_robot_spec for exact workspace bounds and camera names. + """, + """ + Cameras: scene is the fixed overview RGB-D camera; wrist is the hand-mounted RGB-D camera. view_env_state and observe return both images. back_project maps a pixel + depth to 3D; it returns robot-base xyz in panda_link0 only when that camera is calibrated for the selected step. + """, + """ + Use the wrist camera as the default for back_project because it has been calibrated and sees close-range manipulation targets with better depth accuracy. Use the scene camera for overview/context or if you explicitly need it, and only trust any camera's base-frame xyz if back_project reports calibrated=true. Note that the scene camera is placed opposite the robot and therefore has a mirrored view of the robot and table. + """, + """ + Motion tools: move_to sends an absolute TCP xyz in panda_link0; move_delta sends a bounded relative dxyz. Each action returns reached, pos_error_m, final_xyz, and clipping information. Treat images, back_project diagnostics, and returned errors as ground truth. + """, + """ + Gripper: use open_gripper and close_gripper for explicit grasp/release, or set gripper to 'open' or 'close' on move_to/move_delta when that is exactly what you want before the move. + """, + """ + Tools: view_env_state, observe, back_project, get_camera_meta, get_ee_pose, get_robot_spec, move_to, move_delta, rotate_wrist_yaw, rotate_gripper, open_gripper, close_gripper, finish, plus common file and memory tools. + """, +]) + +RULES = BulletList([ + """ + Observe before acting. Call read_memory first, then view_env_state or observe to see the current setup. + """, + """ + Be discreet when moving around. Prefer move_delta for visual servoing and approach/lift motions. + """, + """ + Do not repeat a failed move blindly. If reached is false or pos_error_m is large, inspect the latest images and choose a smaller or different motion. + """, + """ + When grabbing an object, compare the gripper/object z position and the wrist camera view to ensure the object is actually grasped. The gripper z position should be close to the center of the object. + """, + """ + If the env server returns an error, stop and report it instead of continuing blindly. + """, +]) + +WORKFLOW = Numbered([ + """ + Read memory: call read_memory with no arguments, then read any relevant entry. + """, + """ + Observe: call view_env_state or observe and inspect scene plus wrist images and the TCP pose. + """, + """ + Localize with back_project on the wrist camera first. Use get_camera_meta if you need to check which cameras are calibrated to panda_link0, and use the scene camera mainly for overview/context. + """, + """ + Plan conservative motions in panda_link0. Use get_robot_spec if you need bounds and get_ee_pose if you need the live TCP pose. + """, + """ + Use move_delta for local corrections, move_to for known absolute targets, rotate_gripper or rotate_wrist_yaw only for jaw alignment, and explicit gripper tools for grasp/release. + """, + """ + Verify after each step from both the returned result and images. Re-observe if anything may have moved or settled. + """, + """ + Record a durable lesson with write_memory only if this run teaches a non-obvious verified offset, gotcha, or recovery strategy. + """, + """ + Finish only after verifying success or determining the task is unrecoverable. + """, +]) + +USER_CONTEXT = { + "Task": """ + Pick up the blue cube among the colored cubes. You succeed when the blue cube is grasped and lifted above the table. + """, + # Pick up the purple hexagonal prism and insert it into the matching hole in the green block. The prism is on the table in front of the robot, and the block is fixed to the table. + "Run": """ + - output_dir: {{output_dir}} + """, +} + + +def system_prompt() -> dict[str, PromptNode]: + """Return the system prompt tree.""" + return { + "Intro": PREAMBLE, + "Goal": GOAL, + "Rules": RULES, + "Workflow": WORKFLOW, + "Environment": ENVIRONMENT, + "Output": base_prompt.OUTPUT, + } + + +def user_prompt() -> dict[str, PromptNode]: + """Return the first user message tree.""" + return dict(USER_CONTEXT) diff --git a/robots/franka/run_env_server.sh b/robots/franka/run_env_server.sh new file mode 100755 index 00000000..247cc17f --- /dev/null +++ b/robots/franka/run_env_server.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# Launch the standalone Franka env server in the RLinf .venv with the +# serl_franka_controllers catkin workspace sourced. +# +# The agent (agent venv) connects to this over TCP. Two ways to use it: +# +# 1. Fixed port, then attach the agent via --env-endpoint: +# bash robots/franka/run_env_server.sh --output-dir /tmp/franka_run \ +# --transport socket --host 127.0.0.1 --port 5599 +# # in the agent venv: +# python rpent/cli/main.py --env franka \ +# --env-endpoint socket://127.0.0.1:5599 ... +# +# 2. Let the agent CLI spawn it (it invokes this script automatically): +# python rpent/cli/main.py --env franka ... +# +# Override the machine-specific bits via env vars: +# FRANKA_CATKIN_SETUP catkin devel setup.bash (default: RLinf .venv workspace) +# RLINF_VENV_PYTHON python in the RLinf .venv +# FRANKA_ROBOT_IP robot IP (default 172.16.0.2) +set -euo pipefail + +FRANKA_CATKIN_SETUP="${FRANKA_CATKIN_SETUP:-/home/franka/franka/RLinf/.venv/franka_catkin_ws/devel/setup.bash}" +RLINF_VENV_PYTHON="${RLINF_VENV_PYTHON:-/home/franka/franka/RLinf/.venv/bin/python}" +FRANKA_ROBOT_IP="${FRANKA_ROBOT_IP:-172.16.0.2}" + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ENV_SERVER="${SCRIPT_DIR}/env_server.py" + +if [[ ! -f "${FRANKA_CATKIN_SETUP}" ]]; then + echo "!! catkin setup not found: ${FRANKA_CATKIN_SETUP}" >&2 + echo " set FRANKA_CATKIN_SETUP to your serl_franka_controllers workspace." >&2 + exit 1 +fi +if [[ ! -x "${RLINF_VENV_PYTHON}" ]]; then + echo "!! RLinf venv python not found: ${RLINF_VENV_PYTHON}" >&2 + echo " set RLINF_VENV_PYTHON to the RLinf .venv python." >&2 + exit 1 +fi + +# shellcheck disable=SC1090 +source "${FRANKA_CATKIN_SETUP}" + +exec "${RLINF_VENV_PYTHON}" "${ENV_SERVER}" --robot-ip "${FRANKA_ROBOT_IP}" "$@" diff --git a/robots/franka/toolkit.py b/robots/franka/toolkit.py new file mode 100644 index 00000000..2080f768 --- /dev/null +++ b/robots/franka/toolkit.py @@ -0,0 +1,88 @@ +"""Franka toolkit: common tools plus conservative Cartesian primitives.""" +from __future__ import annotations + +from functools import partial +from typing import Any + +from robots.franka import tools as franka_tools +from rpent.dashboard.events import DashboardEventSink +from rpent.tools.state import EnvState +from rpent.tools.toolkit import Toolkit +from rpent.utils.logging import get_output_dir + + +class FrankaToolkit(Toolkit): + """Toolkit for the standalone Franka environment.""" + + # view_env_state image slots: primary scene -> _image_bytes, wrist -> _image_cam_bytes. + _VIEW_IMAGE_SLOTS = { + "_image_bytes": "scene.png", + "_image_cam_bytes": "wrist.png", + } + # Per-env artifact names for the dashboard's live frame images. + _FRAME_ARTIFACTS = { + "camera": "scene.png", + "wrist": "wrist.png", + } + + def __init__( + self, + *, + env: Any, + dashboard_events: DashboardEventSink, + ) -> None: + # EnvState owns the trace + counter for this run (explicit output_dir, + # no process-global). The runner will own its lifecycle in a later cut; + # for now the toolkit constructs it from get_output_dir(). + state = EnvState(get_output_dir()) + super().__init__(dashboard_events=dashboard_events, state=state) + self.init_driver_clean(env=env) + self._register_tools() + + def _register_tools(self) -> None: + # Read-only tools whose handlers aren't driver methods (they need the + # run's EnvState bound in). Every other spec binds to its primitive- + # driver method; @updatestate on the method decides state capture. + state_handlers = { + "view_env_state": partial( + self._state.view, image_slots=self._VIEW_IMAGE_SLOTS + ), + "back_project": partial(franka_tools.back_project, state=self._state), + } + for spec in franka_tools.TOOLS_SPEC: + name = spec["name"] + if name in state_handlers: + handler = state_handlers[name] + else: + handler = getattr(self._driver, name, None) + if handler is None: + continue # spec without a backing driver method + self.add_tool(name, spec, handler) + + def get_env_state( + self, + *, + command: dict[str, Any], + result: dict[str, Any], + elapsed_s: float, + ) -> dict[str, Any]: + self._driver._refresh(increment_step=command["action"] != "observe") + record = franka_tools.dump_state( + self._driver, + self._state, + log={"command": command, "result": result, "elapsed_s": elapsed_s}, + ) + out = self._state.view(record.step_idx, image_slots=self._VIEW_IMAGE_SLOTS) + out["agent_elapsed_s"] = elapsed_s + return out + + def init_driver_clean(self, *, env: Any) -> None: + self._state.reset() + driver = franka_tools.FrankaPrimitives(env=env) + driver.reset() + record = franka_tools.dump_state(driver, self._state, log=None) + self._driver = driver + self._publish_step(record) + + def close(self) -> None: + return None diff --git a/robots/franka/tools.py b/robots/franka/tools.py new file mode 100644 index 00000000..c906ae69 --- /dev/null +++ b/robots/franka/tools.py @@ -0,0 +1,500 @@ +"""Franka tool implementation for the agent-side toolkit.""" +from __future__ import annotations + +import time +from typing import Any + +import numpy as np + +from robots.franka.env_client import FrankaEnvClient +from rpent.tools.common import robust_surface_centroid +from rpent.tools.state import EnvState, StepRecord +from rpent.tools.toolkit import updatestate +from rpent.utils.logging import get_logger + +logger = get_logger("franka") + +_MAX_DELTA_M = 0.05 +_MAX_YAW_DELTA_DEG = 30.0 +_MAX_OBSERVE_DELAY_S = 5.0 +_BACKPROJECT_RADIUS = 6 + + +def _to_list(value) -> list: + """Coerce numpy arrays / scalars into a compact JSON-friendly list.""" + if value is None: + return [] + arr = np.asarray(value, dtype=np.float64).reshape(-1) + return [round(float(v), 4) for v in arr] + + +def _to_scalar(value) -> Any: + if hasattr(value, "item"): + try: + return value.item() + except Exception: + pass + return value + + +# State-trace I/O, image/depth layout, and the robust back-projection math now +# live in :class:`rpent.tools.state.EnvState` (one owner per run). The thin +# wrappers below delegate to it. + + +class FrankaPrimitives: + """Primitive driver owned by :class:`FrankaToolkit`.""" + + def __init__(self, env: FrankaEnvClient): + self.env = env + self._last_obs: dict | None = None + self._spec: dict | None = None + self._num_steps = 0 + + def reset(self) -> tuple[dict, Any]: + """Reset the arm and cache the initial observation.""" + self._spec = self.env.get_spec() + obs, info = self.env.reset() + self._last_obs = obs + self._num_steps = 0 + return obs, info + + @updatestate + def observe(self, delay_s: float = 0.0) -> dict: + """Wait before the toolkit captures a fresh observation.""" + delay = float(np.clip(delay_s, 0.0, _MAX_OBSERVE_DELAY_S)) + if delay > 0: + time.sleep(delay) + return {"delay_s": delay} + + def get_robot_spec(self) -> dict: + """Return the driver self-description.""" + if self._spec is None: + self._spec = self.env.get_spec() + return self._spec + + def get_ee_pose(self) -> dict: + """Return the live TCP pose in the Franka base frame.""" + return self.env.get_ee_pose() + + def get_camera_meta(self) -> dict: + """Return live camera intrinsics/extrinsics metadata.""" + return self.env.get_camera_meta() + + @updatestate + def move_to( + self, + xyz, + *, + yaw_deg: float | None = None, + gripper: str | None = None, + ) -> dict: + """Move to an absolute base-frame Cartesian target.""" + return self.env.move_to(xyz, yaw_deg=yaw_deg, gripper=gripper) + + @updatestate + def move_delta( + self, + dxyz, + *, + gripper: str | None = None, + ) -> dict: + """Nudge the TCP by a bounded relative translation.""" + requested = np.asarray(dxyz, dtype=np.float64).reshape(-1)[:3] + clipped = np.clip(requested, -_MAX_DELTA_M, _MAX_DELTA_M) + result = self.env.move_delta(dxyz=clipped, gripper=gripper) + if np.any(clipped != requested): + result = dict(result) + result["requested_dxyz"] = _to_list(requested) + result["clipped_dxyz"] = _to_list(clipped) + return result + + @updatestate + def rotate_wrist_yaw(self, delta_deg: float) -> dict: + """Rotate the wrist yaw relatively, capped for safety.""" + requested = float(delta_deg) + clipped = float(np.clip(requested, -_MAX_YAW_DELTA_DEG, _MAX_YAW_DELTA_DEG)) + result = self.env.move_delta(yaw_delta_deg=clipped) + if clipped != requested: + result = dict(result) + result["requested_delta_deg"] = round(requested, 3) + result["clipped_delta_deg"] = round(clipped, 3) + return result + + @updatestate + def rotate_gripper(self, delta_deg: float) -> dict: + """Rotate the gripper jaw heading relatively, capped for safety.""" + return self.rotate_wrist_yaw(delta_deg) + + @updatestate + def open_gripper(self) -> dict: + return self.env.open_gripper() + + @updatestate + def close_gripper(self) -> dict: + return self.env.close_gripper() + + def get_state(self) -> dict: + """Return compact proprioception from the latest observation.""" + obs = self._last_obs or {} + state = obs.get("state", {}) if isinstance(obs, dict) else {} + out = { + "tcp_xyz": _to_list(state.get("tcp_xyz")), + "tcp_quat": _to_list(state.get("tcp_quat")), + "tcp_euler": _to_list(state.get("tcp_euler")), + "gripper_width": round(float(_to_scalar(state.get("gripper_width", 0.0))), 4), + "gripper_open": bool(_to_scalar(state.get("gripper_open", False))), + "num_steps": self._num_steps, + } + if self._spec is not None: + out["workspace_min"] = self._spec.get("workspace_min") + out["workspace_max"] = self._spec.get("workspace_max") + out["frame"] = self._spec.get("world_frame") + return out + + def latest_frames(self) -> dict: + """Return the camera frames from the latest observation.""" + if self._last_obs is None: + return {} + return dict(self._last_obs.get("frames", {})) + + def latest_depths(self) -> dict: + """Return metric depth maps from the latest observation.""" + if self._last_obs is None: + return {} + return dict(self._last_obs.get("depth", {})) + + def latest_camera_meta(self) -> dict: + """Return camera metadata from the latest observation.""" + if self._last_obs is None: + return {} + return dict(self._last_obs.get("camera_meta", {})) + + def _refresh(self, *, increment_step: bool = True) -> None: + """Refresh the cached observation for toolkit state capture.""" + try: + self._last_obs = self.env.get_obs() + if increment_step: + self._num_steps += 1 + except Exception as exc: + logger.warning("obs refresh failed: %s", exc) + + +def dump_state( + driver: FrankaPrimitives, + state: EnvState, + log: dict | None = None, +) -> StepRecord: + """Dump camera artifacts and proprioceptive state through ``EnvState``.""" + log = log or {} + with state.record_step( + state=driver.get_state(), + command=log.get("command"), + result=log.get("result"), + elapsed_s=log.get("elapsed_s"), + ) as step_idx: + for camera, frame in driver.latest_frames().items(): + state.save( + f"{camera}.png", + frame, + step=step_idx, + ) + + for camera, depth in driver.latest_depths().items(): + state.save( + f"{camera}_depth.npy", + depth, + step=step_idx, + ) + + for camera, camera_meta in driver.latest_camera_meta().items(): + state.save( + f"{camera}_metadata.json", + camera_meta, + step=step_idx, + ) + + return state.get(step_idx) + + +# view_env_state now lives on EnvState.view (bound by the toolkit with the +# scene/wrist image slots); nothing module-level is needed here. + + +def back_project( + row: int, + col: int, + step: int = -1, + camera: str = "wrist", + radius: int | None = _BACKPROJECT_RADIUS, + *, + state: EnvState, +) -> dict: + """Back-project a saved RGB-D pixel into camera and robot-base coordinates. + + Uses :func:`rpent.tools.common.robust_surface_centroid` (median of the + dominant surface in a window around the pixel) so oblique-view depth noise + does not become ~cm lateral error. Returns robot-base ``xyz`` in + ``panda_link0`` when the camera has ``T_base_cam`` for the step; otherwise + ``xyz_cam`` plus a warning. + """ + try: + record = state.get(step) + except Exception as exc: + return {"error": f"state step not available: {exc}"} + nn = record.step_idx + + camera = str(camera or "wrist") + metadata_name = f"{camera}_metadata.json" + depth_name = f"{camera}_depth.npy" + if metadata_name not in record.artifacts: + return { + "error": f"camera {camera!r} has no metadata at step {nn}", + "available_cameras": sorted( + name.removesuffix("_metadata.json") + for name in record.artifacts + if name.endswith("_metadata.json") + ), + } + try: + meta = state.load(metadata_name, step=nn) + if depth_name not in record.artifacts: + raise FileNotFoundError(depth_name) + depth = state.load(depth_name, step=nn) + except Exception as exc: + return {"error": f"depth for camera {camera!r} step {nn} not found: {exc}"} + + res = robust_surface_centroid( + depth, + meta["K"], + meta.get("T_base_cam"), + row, + col, + radius=_BACKPROJECT_RADIUS if radius is None else radius, + ) + if "error" in res: + return res + + res["step"] = nn + res["camera"] = camera + res["camera_frame"] = meta.get("frame", f"{camera}_camera") + res["calibrated"] = bool(meta.get("calibrated")) + res["calibration_kind"] = meta.get("calibration_kind") + if meta.get("T_base_cam") is not None: + res["frame"] = "panda_link0" + else: + res["frame"] = meta.get("frame", f"{camera}_camera") + res["note"] = ( + f"camera {camera!r} is not calibrated to panda_link0 for this step; " + "do not use xyz_cam as a robot target. Add T_base_cam for a fixed " + "camera or T_tcp_cam for a wrist camera." + ) + return res + + +TOOLS_SPEC: list[dict[str, Any]] = [ + { + "name": "view_env_state", + "description": ( + "Read one recorded state and its observation artifacts. Step -1 " + "selects the latest entry. Embeds scene and wrist images." + ), + "input_schema": { + "type": "object", + "properties": { + "step": { + "type": "integer", + "default": -1, + "description": "Step number; 0 = initial, -1 = latest.", + } + }, + }, + }, + { + "name": "back_project", + "description": ( + "Backproject a pixel from a saved RGB-D camera image to a 3D point. " + "Defaults to camera='wrist', the preferred camera for close-range " + "Franka manipulation after wrist calibration. Returns robot-base " + "`xyz` in panda_link0 only when that camera has calibration for " + "the selected step; otherwise returns " + "xyz_cam plus a warning. Pick row/col on the image returned by " + "view_env_state." + ), + "input_schema": { + "type": "object", + "properties": { + "row": {"type": "integer", "description": "Pixel row (y)."}, + "col": {"type": "integer", "description": "Pixel column (x)."}, + "step": { + "type": "integer", + "default": -1, + "description": "Step whose saved depth to use; -1 = latest.", + }, + "camera": { + "type": "string", + "enum": ["scene", "wrist"], + "description": "RGB-D camera to use. Default and preferred after calibration: wrist.", + }, + "radius": { + "type": ["integer", "null"], + "description": "Half-size of depth window in pixels; null = default 6.", + }, + }, + "required": ["row", "col"], + }, + }, + { + "name": "observe", + "description": ( + "Refresh the live observation without moving the arm, dump a new " + "step, and return the updated state/images. Use this when the scene " + "may have changed or after waiting for settling." + ), + "input_schema": { + "type": "object", + "properties": { + "delay_s": { + "type": ["number", "null"], + "description": "Optional wait before observing; clipped to 0..5 seconds.", + } + }, + }, + }, + { + "name": "get_ee_pose", + "description": ( + "Live TCP pose in the Franka base frame panda_link0. Returns xyz " + "meters, quat_xyzw, euler_xyz radians, and frame." + ), + "input_schema": {"type": "object", "properties": {}}, + }, + { + "name": "get_robot_spec", + "description": ( + "Static robot/environment description: workspace bounds, frame, " + "camera names, control mode, gripper mode, and reset pose." + ), + "input_schema": {"type": "object", "properties": {}}, + }, + { + "name": "get_camera_meta", + "description": ( + "Live per-camera intrinsics, depth scale, and calibration status. " + "For base-frame back_project, wrist needs T_tcp_cam composed with " + "the live TCP pose; scene needs T_base_cam." + ), + "input_schema": {"type": "object", "properties": {}}, + }, + { + "name": "move_to", + "description": ( + "Move the TCP to absolute [x, y, z] in panda_link0, meters. The " + "driver clips targets to the safe workspace and returns reached, " + "pos_error_m, final_xyz, and clipping info. yaw_deg optionally sets " + "a down-facing grasp yaw; null keeps current orientation. gripper " + "may be null, 'open', or 'close'." + ), + "input_schema": { + "type": "object", + "properties": { + "xyz": { + "type": "array", + "description": "Absolute target [x, y, z] in meters, panda_link0 frame.", + "items": {"type": "number"}, + "minItems": 3, + "maxItems": 3, + }, + "yaw_deg": { + "type": ["number", "null"], + "description": "Optional down-facing wrist yaw in degrees; null keeps current orientation.", + }, + "gripper": { + "type": ["string", "null"], + "enum": ["open", "close", None], + "description": "Optional gripper action to execute before the move.", + }, + }, + "required": ["xyz"], + }, + }, + { + "name": "move_delta", + "description": ( + "Nudge the TCP by relative [dx, dy, dz] meters in panda_link0. Each " + "axis is clipped to +/-0.05 m per call. Use for visual servoing and " + "small approach/lift motions. gripper may be null, 'open', or 'close'." + ), + "input_schema": { + "type": "object", + "properties": { + "dxyz": { + "type": "array", + "description": "Relative TCP translation [dx, dy, dz] in meters.", + "items": {"type": "number", "minimum": -0.05, "maximum": 0.05}, + "minItems": 3, + "maxItems": 3, + }, + "gripper": { + "type": ["string", "null"], + "enum": ["open", "close", None], + "description": "Optional gripper action to execute before the nudge.", + }, + }, + "required": ["dxyz"], + }, + }, + { + "name": "rotate_wrist_yaw", + "description": ( + "Rotate the wrist yaw relatively without translating the TCP. The " + "requested delta is clipped to +/-30 degrees per call. Use only for " + "jaw alignment after the arm is already near the target." + ), + "input_schema": { + "type": "object", + "properties": { + "delta_deg": { + "type": "number", + "minimum": -30, + "maximum": 30, + "description": "Relative yaw change in degrees.", + } + }, + "required": ["delta_deg"], + }, + }, + { + "name": "rotate_gripper", + "description": ( + "Rotate the gripper jaw heading relatively without translating the " + "TCP. The requested delta is clipped to +/-30 degrees per call. " + "Use when the fingers need a different angle before grasping." + ), + "input_schema": { + "type": "object", + "properties": { + "delta_deg": { + "type": "number", + "minimum": -30, + "maximum": 30, + "description": "Relative gripper yaw change in degrees.", + } + }, + "required": ["delta_deg"], + }, + }, + { + "name": "open_gripper", + "description": "Open the Franka Hand, refresh observation, and return updated state/images.", + "input_schema": {"type": "object", "properties": {}}, + }, + { + "name": "close_gripper", + "description": ( + "Close/grasp with the Franka Hand, refresh observation, and return " + "updated state/images." + ), + "input_schema": {"type": "object", "properties": {}}, + }, +] diff --git a/robots/lerobot/__init__.py b/robots/lerobot/__init__.py new file mode 100644 index 00000000..d10c95ea --- /dev/null +++ b/robots/lerobot/__init__.py @@ -0,0 +1,187 @@ +"""LeRobot SO101 environment extension. + +Entry point for the env registry: :func:`get_env_spec` and +:func:`get_toolkit` are discovered by +:func:`rpent.envs.base._resolve_env` via +``importlib.import_module("robots.lerobot")`` — dropping this +package on disk is the entire registration step. +""" +from __future__ import annotations + +import argparse +import sys +from datetime import datetime +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from robots.lerobot.prompt import ( + system_prompt, + user_prompt, +) +from rpent.dashboard.events import DashboardEventSink +from rpent.envs.env_spec import EnvSpec, RunConfig +from rpent.envs.prompt_bundle import PromptBundle +from rpent.utils.config import get_repo_root +from rpent.utils.logging import get_logger + +if TYPE_CHECKING: + from rpent.utils.daemon import ProcessDaemon + from rpent.utils.rpc import RpcClient + +logger = get_logger("lerobot") + + +def get_env_spec() -> EnvSpec: + """Return the SO101 env identity, prompt bundle, and runner hooks. + + Tool schemas, handlers, and the MCP allowlist live on the toolkit (see + :func:`get_toolkit`). The three runner hooks (:func:`_add_cli_args` / + :func:`_parse_config` / :func:`_init_runtime`) keep ``rpent/cli/main.py`` + env-agnostic, mirroring :mod:`robots.libero`. + """ + return EnvSpec( + name="lerobot", + prompts=PromptBundle( + system=system_prompt, + user=user_prompt, + ), + add_cli_args=_add_cli_args, + parse_config=_parse_config, + init_runtime=_init_runtime, + ) + + +def get_toolkit( + *, + primitives_kwargs: dict[str, Any], + dashboard_events: DashboardEventSink, +): + """Return the SO101 toolkit (common tools + SO101 primitives). + + ``primitives_kwargs`` is assembled by :func:`_init_runtime` and carries the + env RPC stub (``{"env": LerobotEnvClient(...)}``, optionally plus + ``"model"``). + """ + from robots.lerobot.toolkit import LerobotToolkit + + return LerobotToolkit( + dashboard_events=dashboard_events, + **primitives_kwargs, + ) + + +def _add_cli_args(parser: argparse.ArgumentParser, use_dashboard: bool) -> None: + """Register SO101 CLI flags on the shared ``parser``. + + ``use_dashboard`` is unused: the SO101 is a real robot with no + suite/task/seed, so there is nothing for the (libero-shaped) dashboard + launcher to fill in. + """ + del use_dashboard + parser.add_argument("--max-episode-steps", type=int, default=200) + parser.add_argument( + "--env-endpoint", default=None, + help="[protocol://]host:port of an existing lerobot env_server " + "(protocol=http|socket, defaults to http). If unset, a local " + "env_server is spawned (requires the 'lerobot' extra in this venv).", + ) + + +def _parse_config(args: argparse.Namespace) -> RunConfig: + """Derive per-run identifiers for a SO101 run. + + Real robots have no suite/task/seed, so the run is identified by the env + name. The dashboard is currently libero-shaped, so it is not wired here. + """ + if getattr(args, "dashboard", False): + logger.warning( + "--dashboard is only supported for the libero env; " + "continuing without the live dashboard." + ) + + recipe_tag = "lerobot" + output_dir = args.output_dir + if output_dir is None: + timestamp = datetime.now().strftime("%Y%m%d-%H:%M:%S") + output_dir = get_repo_root() / "logs" / f"{timestamp}_lerobot" + output_dir = Path(output_dir) + + return RunConfig( + recipe_tag=recipe_tag, + output_dir=output_dir, + prompt_vars={"env_name": "lerobot", "recipe_tag": recipe_tag}, + dashboard_state=None, + task_desc={"env": "lerobot"}, + ) + + +def _parse_endpoint(endpoint: str) -> tuple[str, str, int]: + """Parse ``[protocol://]host:port`` into ``(protocol, host, port)``. + + Protocol defaults to ``http`` when the prefix is omitted. + """ + if "://" in endpoint: + protocol, _, rest = endpoint.partition("://") + else: + protocol, rest = "http", endpoint + host, _, port = rest.partition(":") + if not host or not port: + raise ValueError( + f"--env-endpoint must be [protocol://]host:port, got {endpoint!r}" + ) + return protocol, host, int(port) + + +def _init_runtime( + args: argparse.Namespace, + output_dir: Path, +) -> tuple[list[ProcessDaemon], dict[str, Any]]: + """Spawn (or attach to) the SO101 env_server; build primitives_kwargs. + + The SO101 driver needs the ``lerobot`` stack. It is spawned with this + interpreter (works when the venv has the ``lerobot`` extra); otherwise run + the driver in its own ``lerobot`` conda env and attach via + ``--env-endpoint``. No VLA server — the SO101 primitives are scripted. + + Heavy deps are imported lazily so a bare ``import robots.lerobot`` (for + ``get_env_spec`` / ``get_toolkit``) doesn't drag them in. + """ + from robots.lerobot.env_client import LerobotEnvClient + from rpent.utils.daemon import ProcessDaemon, pick_free_port + from rpent.utils.http_rpc import HttpRpcClient + from rpent.utils.rpc import wait_for_ready + from rpent.utils.socket_rpc import SocketRpcClient + + daemons: list[ProcessDaemon] = [] + if args.env_endpoint is None: + host, port = "127.0.0.1", pick_free_port() + env_daemon = ProcessDaemon( + name="env_server", + cmd=[ + sys.executable, + str(get_repo_root() / "robots" / "lerobot" / "env_server.py"), + "--output-dir", str(output_dir), + "--max-episode-steps", str(args.max_episode_steps), + "--transport", "http", + "--host", host, + "--port", str(port), + ], + log_path=str(Path(output_dir) / "env_server.log"), + ) + env_daemon.start() + daemons.append(env_daemon) + env_client: RpcClient = HttpRpcClient(f"http://{host}:{port}") + wait_for_ready(env_client) + else: + protocol, host, port = _parse_endpoint(args.env_endpoint) + if protocol == "socket": + env_client = SocketRpcClient(host, port) + elif protocol == "http": + env_client = HttpRpcClient(f"http://{host}:{port}") + else: + raise ValueError( + f"--env-endpoint protocol must be socket or http, got {protocol!r}" + ) + + primitives_kwargs = {"env": LerobotEnvClient(env_client)} + return daemons, primitives_kwargs diff --git a/robots/lerobot/auto_calibrate_scene_cam.py b/robots/lerobot/auto_calibrate_scene_cam.py new file mode 100644 index 00000000..4453f258 --- /dev/null +++ b/robots/lerobot/auto_calibrate_scene_cam.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python3 +"""Automatic, markerless scene-camera -> base calibration for the SO101. + +Triggers the env server's :meth:`auto_calibrate_scene_camera` routine, which: + +1. drives the gripper to a wide, non-coplanar grid of base-frame positions at + varied wrist orientations (``move_to`` is pure base-frame IK, so it needs no + extrinsic), +2. at each pose toggles the gripper with the arm frozen and segments the motion + in the temporally median-filtered scene image to locate the moving jaw (blob + centroid + depth at that centroid -> camera point); FK gives the tip pose, +3. jointly fits ``T_base_cam`` and the constant centroid-vs-tip offset (RANSAC), + so the offset no longer inflates the residual, and saves it (hot-loaded by + the server, so back_project returns world coords immediately). + +No human input, no markers. Start the env server first, then run:: + + conda activate lerobot + python robots/lerobot/auto_calibrate_scene_cam.py --port 53101 + +WARNING: this moves the arm through many poses. Clear the workspace first. + +Offline math check (no hardware):: + + python robots/lerobot/auto_calibrate_scene_cam.py --self-test +""" +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from robots.lerobot import geometry as geom # noqa: E402 +from rpent.utils.socket_rpc import SocketRpcClient # noqa: E402 + + +def _self_test() -> int: + """Validate RANSAC Kabsch + motion segmentation offline (no hardware).""" + rng = np.random.default_rng(0) + Q, _ = np.linalg.qr(rng.standard_normal((3, 3))) + if np.linalg.det(Q) < 0: + Q[:, 0] = -Q[:, 0] + T_true = np.eye(4) + T_true[:3, :3] = Q + T_true[:3, 3] = rng.standard_normal(3) + cam = rng.standard_normal((10, 3)) + base = geom.transform_points(T_true, cam) + rng.standard_normal((10, 3)) * 1e-3 + base[3] += [0.2, -0.15, 0.1] # inject an outlier + T_est, rmse, inliers = geom.ransac_kabsch(cam, base, thresh_m=0.02) + ok_fit = np.allclose(T_est, T_true, atol=2e-2) and not inliers[3] + print(f"ransac_kabsch: rmse={rmse:.5f}m inliers={int(inliers.sum())}/10 " + f"outlier_excluded={not inliers[3]} -> {'OK' if ok_fit else 'FAIL'}") + + # Synthetic two-frame motion: a blob that appears in the 'closed' frame. + H, W = 480, 640 + rgb_open = np.zeros((H, W, 3), np.uint8) + rgb_closed = rgb_open.copy() + rgb_closed[300:330, 400:430] = 255 # 'fingers' light up at (row~315,col~415) + depth = np.full((H, W), 0.4, np.float32) + K = np.array([[600, 0, 320], [0, 600, 240], [0, 0, 1]], float) + det = geom.detect_tip_pixel_by_motion(rgb_open, rgb_closed, depth, K) + ok_det = det is not None and abs(det["pixel"][0] - 314.5) < 3 and abs(det["pixel"][1] - 414.5) < 3 + print(f"detect_tip: {det if det else 'None'} -> {'OK' if ok_det else 'FAIL'}") + + # Joint extrinsic + tip-offset recovery (solve_extrinsic_with_offset): with + # varied per-pose rotations the solver should recover T_base_cam AND the + # constant local offset, and beat the offset-blind fit, despite an outlier. + rng3 = np.random.default_rng(2) + Qc, _ = np.linalg.qr(rng3.standard_normal((3, 3))) + if np.linalg.det(Qc) < 0: + Qc[:, 0] = -Qc[:, 0] + T_cam = np.eye(4) + T_cam[:3, :3] = Qc + T_cam[:3, 3] = rng3.standard_normal(3) + d_true = np.array([0.015, -0.010, 0.020]) # constant gripper-local offset + n = 24 + origins = rng3.uniform([-0.1, -0.2, 0.0], [0.35, 0.2, 0.25], size=(n, 3)) + rots = np.empty((n, 3, 3)) + for i in range(n): + Qr, _ = np.linalg.qr(rng3.standard_normal((3, 3))) + if np.linalg.det(Qr) < 0: + Qr[:, 0] = -Qr[:, 0] + rots[i] = Qr + feat_base = origins + np.einsum("nij,j->ni", rots, d_true) + cam_j = geom.transform_points(geom.invert_transform(T_cam), feat_base) + cam_j += rng3.standard_normal((n, 3)) * 1e-3 # 1 mm detection noise + cam_j[5] += np.array([0.18, -0.12, 0.15]) # inject an outlier + T_est_j, rmse_j, inl, d_est = geom.solve_extrinsic_with_offset( + cam_j, origins, rots, thresh_m=0.01, + ) + _, rmse0, _ = geom.ransac_kabsch(cam_j, origins, thresh_m=0.01) + ok_solver = ( + np.allclose(T_est_j[:3, :3], T_cam[:3, :3], atol=5e-3) + and np.allclose(T_est_j[:3, 3], T_cam[:3, 3], atol=5e-3) + and np.allclose(d_est, d_true, atol=5e-3) + and not bool(inl[5]) + ) + print(f"solve_offset: rmse={rmse_j * 1000:.2f}mm " + f"(offset-blind {rmse0 * 1000:.1f}mm) " + f"d_err={np.linalg.norm(d_est - d_true) * 1000:.2f}mm " + f"outlier_excluded={not bool(inl[5])} -> {'OK' if ok_solver else 'FAIL'}") + return 0 if (ok_fit and ok_det and ok_solver) else 1 + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--host", default="127.0.0.1") + ap.add_argument("--port", type=int, help="env server transport port.") + ap.add_argument("--n-points", type=int, default=24, + help="Target number of valid correspondences to collect.") + ap.add_argument("--no-save", action="store_true", + help="Compute T_base_cam but do not write it to disk.") + ap.add_argument("--self-test", action="store_true", + help="Run the offline math check and exit.") + args = ap.parse_args() + + if args.self_test: + return _self_test() + if args.port is None: + ap.error("--port is required (the env server's transport port).") + + print("This moves the arm through a grid of poses. Ensure the workspace is " + "clear. Starting...") + client = SocketRpcClient(args.host, args.port) + result = client.call( + "env.auto_calibrate_scene_camera", + kwargs={"n_points": args.n_points, "save": not args.no_save}, + timeout_s=600.0, + ) + + if "error" in result: + print(f"Calibration failed: {result['error']}") + if "poses" in result: + print(json.dumps(result["poses"], indent=2)) + return 2 + + print(f"\nUsed {result['n_used']} poses " + f"({result['n_inliers']} inliers), RMSE = {result['rmse_m'] * 1000:.1f} mm") + off = result.get("tip_offset_local_m") + if off: + print("Estimated tip-detector offset (gripper frame): " + f"[{off[0] * 1000:.1f}, {off[1] * 1000:.1f}, {off[2] * 1000:.1f}] mm") + if result["rmse_m"] > 0.02: + print("WARNING: RMSE > 2 cm — check lighting / gripper visibility and rerun.") + if result.get("saved"): + print(f"Saved T_base_cam -> {result['path']}") + print("The server hot-loaded it; back_project now returns world coords.") + else: + print("Not saved (--no-save). T_base_cam:") + print(json.dumps(result["T_base_cam"], indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/robots/lerobot/calibrate_scene_cam.py b/robots/lerobot/calibrate_scene_cam.py new file mode 100644 index 00000000..95f68633 --- /dev/null +++ b/robots/lerobot/calibrate_scene_cam.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +"""Touch / point-correspondence calibration of the scene camera -> arm base. + +Computes the fixed extrinsic ``T_base_cam`` that maps scene-camera points into +the SO101 ``base_link`` world frame, using only the arm's own FK + the scene +camera's aligned depth (no marker / no extra hardware). The result is saved via +:mod:`robots.lerobot.calibration` and auto-loaded by the env server, after +which ``back_project`` returns world coordinates. + +Procedure (per correspondence): + +1. Free-drive the arm by hand so the gripper tip (``gripper_frame_link``, + roughly the point between the fingertips) rests at a distinct location that + is clearly visible to the scene camera. +2. The script reads the tip position in the base frame from FK + (``env.get_ee_pose``) and grabs the scene color + aligned depth + (``env.get_scene_frame``). +3. You click that same tip point in the color image; the script backprojects + the clicked pixel (median depth over a small patch) into the camera frame. + +After N>=4 non-coplanar points, a rigid Kabsch fit gives ``T_base_cam`` and the +fit RMSE (lower is better; aim for < ~1 cm). + +Run the env server first with a fixed ``--port``, then:: + + conda activate lerobot + python robots/lerobot/calibrate_scene_cam.py --port 53101 --num-points 6 + +Offline self-test of the math (no hardware):: + + python robots/lerobot/calibrate_scene_cam.py --self-test +""" +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import numpy as np + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from robots.lerobot import calibration as scene_calib # noqa: E402 +from robots.lerobot import geometry as geom # noqa: E402 +from rpent.utils.socket_rpc import SocketRpcClient # noqa: E402 + + +def _self_test() -> int: + """Validate the Kabsch pipeline on synthetic correspondences (no hardware).""" + rng = np.random.default_rng(0) + A = rng.standard_normal((3, 3)) + Q, _ = np.linalg.qr(A) + if np.linalg.det(Q) < 0: + Q[:, 0] = -Q[:, 0] + T_true = np.eye(4) + T_true[:3, :3] = Q + T_true[:3, 3] = rng.standard_normal(3) + + cam_pts = rng.standard_normal((8, 3)) + base_pts = geom.transform_points(T_true, cam_pts) + rng.standard_normal((8, 3)) * 1e-3 + T_est, rmse = geom.kabsch_umeyama(cam_pts, base_pts) + ok = np.allclose(T_est, T_true, atol=1e-2) + print(f"self-test: rmse={rmse:.5f}m recovered={'OK' if ok else 'FAIL'}") + return 0 if ok else 1 + + +def _click_pixel(color: np.ndarray, idx: int, total: int) -> tuple[float, float] | None: + """Show the color frame and return the clicked (col, row), or None if skipped.""" + import matplotlib.pyplot as plt + + fig = plt.figure(figsize=(8, 6)) + plt.imshow(color) + plt.title(f"[{idx}/{total}] Click the gripper TIP, then close is automatic. " + "Close window to skip.") + plt.tight_layout() + pts = plt.ginput(1, timeout=0) + plt.close(fig) + if not pts: + return None + return float(pts[0][0]), float(pts[0][1]) + + +def _collect(client: SocketRpcClient, num_points: int, patch_radius: int) -> tuple[np.ndarray, np.ndarray]: + """Collect (cam_point, base_point) correspondences interactively.""" + cam_pts: list[list[float]] = [] + base_pts: list[list[float]] = [] + + i = 0 + while len(cam_pts) < num_points: + i += 1 + input( + f"\n[{len(cam_pts) + 1}/{num_points}] Move the gripper tip to a distinct " + "scene point (vary x/y/z), hold it, then press Enter to capture..." + ) + ee = client.call("env.get_ee_pose", timeout_s=15) + if "error" in ee: + print(f" get_ee_pose failed: {ee['error']}") + continue + frame = client.call("env.get_scene_frame", timeout_s=15) + if "error" in frame: + print(f" get_scene_frame failed: {frame['error']}") + continue + + color = np.asarray(frame["color"], dtype=np.uint8) + depth = np.asarray(frame["depth"], dtype=np.float32) + K = np.asarray(frame["K"], dtype=np.float64) + + click = _click_pixel(color, len(cam_pts) + 1, num_points) + if click is None: + print(" skipped (no pixel clicked).") + continue + col, row = click + z = geom.sample_depth_patch(depth, int(round(col)), int(round(row)), radius=patch_radius) + if not np.isfinite(z) or z <= 0: + print(f" no valid depth at ({int(row)},{int(col)}); try another point/angle.") + continue + + p_cam = geom.backproject_pixel(K, col, row, z) + p_base = np.asarray(ee["xyz"], dtype=np.float64) + cam_pts.append(p_cam.tolist()) + base_pts.append(p_base.tolist()) + print(f" captured: cam={np.round(p_cam, 3).tolist()} " + f"base={np.round(p_base, 3).tolist()} depth={z:.3f}m") + + return np.asarray(cam_pts), np.asarray(base_pts) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--host", default="127.0.0.1", help="env server host.") + ap.add_argument("--port", type=int, help="env server transport port.") + ap.add_argument("--num-points", type=int, default=6, + help="Number of correspondences (>=4; more is better).") + ap.add_argument("--patch-radius", type=int, default=2, + help="Depth median patch radius (pixels) around each click.") + ap.add_argument("--serial", default=None, + help="Override scene serial for the saved file " + "(default: from env.get_scene_camera_meta).") + ap.add_argument("--self-test", action="store_true", + help="Run the offline Kabsch math check and exit.") + args = ap.parse_args() + + if args.self_test: + return _self_test() + if args.port is None: + ap.error("--port is required (the env server's transport port).") + if args.num_points < 4: + ap.error("need at least 4 correspondences for a stable fit.") + + client = SocketRpcClient(args.host, args.port) + meta = client.call("env.get_scene_camera_meta", timeout_s=15) + if "error" in meta: + print(f"scene camera not available: {meta['error']}") + return 2 + serial = args.serial or meta.get("serial") + if not serial: + print("could not determine scene camera serial; pass --serial.") + return 2 + print(f"Calibrating scene camera serial={serial}") + + # Free-drive so the operator can position the tip by hand. + client.call("env.set_torque", args=(False,), timeout_s=15) + print("Arm torque DISABLED — you can move it by hand. (Re-enabled at the end.)") + try: + cam_pts, base_pts = _collect(client, args.num_points, args.patch_radius) + finally: + client.call("env.set_torque", args=(True,), timeout_s=15) + print("Arm torque re-enabled.") + + T_base_cam, rmse = geom.kabsch_umeyama(cam_pts, base_pts) + print(f"\nFit complete: {len(cam_pts)} points, RMSE = {rmse * 1000:.1f} mm") + if rmse > 0.02: + print("WARNING: RMSE > 2 cm — consider recollecting with more spread / " + "better tip-pixel clicks.") + + path = scene_calib.save_extrinsic( + serial, T_base_cam, K=np.asarray(meta["K"]), rmse_m=rmse, num_points=len(cam_pts) + ) + print(f"Saved T_base_cam -> {path}") + print("Restart the env server (or it will pick this up next launch) so " + "back_project returns world coordinates.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/robots/lerobot/calibration.py b/robots/lerobot/calibration.py new file mode 100644 index 00000000..761580c3 --- /dev/null +++ b/robots/lerobot/calibration.py @@ -0,0 +1,78 @@ +"""Load / save the scene-camera → base extrinsic ``T_base_cam``. + +``T_base_cam`` is the fixed rigid transform that maps a point in the scene +camera frame into the arm ``base_link`` world frame. It is produced once by the +touch/Kabsch calibration (``toolkits/lerobot/calibrate_scene_cam.py``) and +loaded by the env server so ``back_project`` can return world coordinates. + +Stored per camera serial under the LeRobot cache so it survives across runs. +""" +from __future__ import annotations + +import json +import os +from pathlib import Path + +import numpy as np + +_CALIB_DIR = "~/.cache/huggingface/lerobot/calibration/scene_cam" + +# A scene-cam -> base fit whose RMSE (meters) exceeds this is treated as a +# FAILED calibration: ``back_project`` would map pixels to badly wrong world +# coordinates (the arm then chases unreachable targets), so the driver refuses +# to save it or trust it on load. ~2 cm matches the auto-calibrator's existing +# "rerun" warning; a good touch/Kabsch fit is typically a few mm. +MAX_ACCEPTABLE_RMSE_M = 0.02 + + +def calib_path(serial: str) -> Path: + """Return the on-disk path of the extrinsic file for a camera ``serial``.""" + return Path(os.path.expanduser(_CALIB_DIR)) / f"{serial}.json" + + +def save_extrinsic( + serial: str, + T_base_cam, + *, + K=None, + rmse_m: float | None = None, + num_points: int | None = None, +) -> Path: + """Persist ``T_base_cam`` (and calibration diagnostics) for ``serial``.""" + path = calib_path(serial) + path.parent.mkdir(parents=True, exist_ok=True) + data: dict = { + "serial": serial, + "T_base_cam": np.asarray(T_base_cam, dtype=np.float64).tolist(), + } + if K is not None: + data["K"] = np.asarray(K, dtype=np.float64).tolist() + if rmse_m is not None: + data["rmse_m"] = float(rmse_m) + if num_points is not None: + data["num_points"] = int(num_points) + with open(path, "w") as f: + json.dump(data, f, indent=2) + return path + + +def load_extrinsic_record(serial: str) -> dict | None: + """Load the full saved calibration record for ``serial`` (or ``None``). + + Returns the parsed JSON with ``T_base_cam`` as a ``(4, 4)`` ndarray plus any + saved diagnostics (``rmse_m``, ``num_points``, ``K``). Use this when you need + to judge fit quality, not just apply the transform. + """ + path = calib_path(serial) + if not path.is_file(): + return None + with open(path) as f: + data = json.load(f) + data["T_base_cam"] = np.asarray(data["T_base_cam"], dtype=np.float64) + return data + + +def load_extrinsic(serial: str) -> np.ndarray | None: + """Load ``T_base_cam`` (4x4) for ``serial``, or ``None`` if not calibrated.""" + record = load_extrinsic_record(serial) + return None if record is None else record["T_base_cam"] diff --git a/robots/lerobot/diagnose_motors.py b/robots/lerobot/diagnose_motors.py new file mode 100644 index 00000000..fba158b6 --- /dev/null +++ b/robots/lerobot/diagnose_motors.py @@ -0,0 +1,173 @@ +"""Motor-health diagnostic for the SO101 follower arm. + +Reads each servo's temperature, voltage, current, load, hardware-error status, +and torque state WITHOUT moving the arm. Use it when the arm stops reaching +commanded positions (drifts / presses the table / grasps miss even though the +scene localization and IK are correct) to check whether a joint is overheated, +under-volted, faulted, or straining. + +Optionally (``--tracking-test``) it nudges each arm joint a few degrees around +its CURRENT pose and measures how far the achieved angle is from the command, +to pinpoint a joint that no longer tracks (slipped horn, weak motor, or +calibration drift). That part MOVES the arm, so put the arm in a safe, raised +pose first. + +Run with the env server STOPPED (it needs exclusive access to the motor bus):: + + conda activate lerobot + python robots/lerobot/diagnose_motors.py # health only, no motion + python robots/lerobot/diagnose_motors.py --tracking-test # also nudges each joint +""" +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +# Make ``rpent`` importable if ever needed; also keeps paths consistent +# with the driver when run from the repo root. +_ROOT = Path(__file__).resolve().parents[2] +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +_ARM_JOINTS = ("shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll") +_ALL_MOTORS = _ARM_JOINTS + ("gripper",) + + +def _read(bus, name: str, motor: str): + """Read one raw register value, returning a short error string on failure.""" + try: + return bus.read(name, motor, normalize=False) + except Exception as e: # noqa: BLE001 - diagnostic must never crash on one bad read + return f"err:{type(e).__name__}" + + +def _fmt(v, width: int) -> str: + return f"{v:>{width}}" if not isinstance(v, float) else f"{v:>{width}.1f}" + + +def _set_p_coefficient(robot, p: int) -> None: + """Set the position gain on the arm joints (to test tracking stiffness).""" + try: + with robot.bus.torque_disabled(): + for m in _ARM_JOINTS: + robot.bus.write("P_Coefficient", m, int(p)) + print(f"[set P_Coefficient={int(p)} on arm joints for this test]\n") + except Exception as e: # noqa: BLE001 + print(f"[could not set P_Coefficient: {e}]\n") + + +def _health(robot) -> None: + bus = robot.bus + try: + obs = robot.get_observation() + except Exception: + obs = {} + print("\n=== motor health (no motion) ===") + print( + f"{'motor':16}{'pos_deg':>9}{'temp_C':>8}{'volt_V':>8}" + f"{'load':>8}{'current':>9}{'status':>8}{'torque':>8}{'Pgain':>7}" + ) + for m in _ALL_MOTORS: + temp = _read(bus, "Present_Temperature", m) # deg C + volt = _read(bus, "Present_Voltage", m) # 0.1 V units + load = _read(bus, "Present_Load", m) + curr = _read(bus, "Present_Current", m) + stat = _read(bus, "Status", m) # 0 = OK; nonzero = HW error flags + torq = _read(bus, "Torque_Enable", m) + pgain = _read(bus, "P_Coefficient", m) + pos = obs.get(f"{m}.pos") + volt_v = round(volt / 10.0, 1) if isinstance(volt, (int, float)) else volt + pos_s = f"{pos:9.1f}" if isinstance(pos, (int, float)) else f"{str(pos):>9}" + print( + f"{m:16}{pos_s}{_fmt(temp, 8)}{_fmt(volt_v, 8)}" + f"{_fmt(load, 8)}{_fmt(curr, 9)}{_fmt(stat, 8)}{_fmt(torq, 8)}{_fmt(pgain, 7)}" + ) + print( + "\nInterpretation:\n" + " temp_C > ~55 -> overheating; Feetech servos derate torque and under-reach.\n" + " volt_V -> should match your supply and be steady; sag => weak torque.\n" + " status != 0 -> a hardware-error flag latched (overload/overheat/voltage).\n" + " load/current -> high while merely holding a light pose => a straining joint.\n" + ) + + +def _tracking_test(robot, nudge_deg: float) -> None: + base = robot.get_observation() + q0 = {m: float(base[f"{m}.pos"]) for m in _ARM_JOINTS} + g0 = float(base.get("gripper.pos", 50.0)) + + def send(qd: dict) -> None: + act = {f"{m}.pos": float(qd[m]) for m in _ARM_JOINTS} + act["gripper.pos"] = g0 + robot.send_action(act) + + print("=== per-joint tracking test (MOVES the arm ±%.0f deg around the current pose) ===" % nudge_deg) + worst = 0.0 + for j in _ARM_JOINTS: + for sgn in (+1.0, -1.0): + qd = dict(q0) + qd[j] = q0[j] + sgn * nudge_deg + send(qd) + time.sleep(0.9) + ach = float(robot.get_observation()[f"{j}.pos"]) + err = ach - qd[j] + worst = max(worst, abs(err)) + flag = " <-- POOR TRACKING" if abs(err) > 3.0 else "" + print(f" {j:16} cmd={qd[j]:7.1f} achieved={ach:7.1f} err={err:+6.1f} deg{flag}") + send(q0) + time.sleep(0.9) + print( + f"\nworst tracking error: {worst:.1f} deg\n" + " > ~3 deg on a free-space nudge = that joint is not tracking " + "(slipped horn / weak or faulted motor / calibration drift). Recalibrate " + "the follower or inspect that motor before running the agent again.\n" + ) + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--port", default="/dev/ttyACM1") + p.add_argument("--calibration-id", default="my_awesome_follower_arm") + p.add_argument( + "--tracking-test", action="store_true", + help="Also nudge each arm joint a few degrees and report the " + "command-vs-readback error. MOVES THE ARM — place it in a safe, " + "raised pose first.", + ) + p.add_argument("--nudge-deg", type=float, default=6.0) + p.add_argument( + "--p-coefficient", type=int, default=None, + help="Set the arm servos' P_Coefficient (position gain) before testing " + "(LeRobot uses 16; factory default 32). Sweep e.g. 32 then 48 to " + "see if tracking tightens.", + ) + args = p.parse_args() + + from lerobot.robots.so_follower import SO101Follower + from lerobot.robots.so_follower.config_so_follower import SO101FollowerConfig + + robot = SO101Follower( + SO101FollowerConfig( + port=args.port, + id=args.calibration_id, + use_degrees=True, + disable_torque_on_disconnect=False, + cameras={}, + ) + ) + robot.connect(calibrate=False) + try: + if args.p_coefficient is not None: + _set_p_coefficient(robot, args.p_coefficient) + _health(robot) + if args.tracking_test: + _tracking_test(robot, args.nudge_deg) + finally: + robot.disconnect() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/robots/lerobot/env_client.py b/robots/lerobot/env_client.py new file mode 100644 index 00000000..c061ae6c --- /dev/null +++ b/robots/lerobot/env_client.py @@ -0,0 +1,141 @@ +"""LeRobot SO101 env client that forwards calls over a driver RPC client. + +Mirrors the RPC surface exposed by ``robots/lerobot/env_server.py`` +(:class:`SO101LeRobotEnv`): a minimal gym-style ``reset`` / ``step`` plus a +``get_spec`` self-description. Each method turns one agent-side call into one +RPC against the driver process via :class:`RpcClient`. + +The agent process does not import torch; the driver returns plain numpy / +floats, so values cross the wire unchanged. +""" +from __future__ import annotations + +from typing import Any + +import numpy as np + +from rpent.utils.rpc import RpcClient + + +# Per-method RPC timeouts (seconds). ``reset`` moves the arm to its rest pose +# (server sleeps ~1s); ``step`` advances one control tick. +_TIMEOUT_S = { + "default": 30.0, + "env.reset": 120.0, + "env.step": 60.0, + "env.get_spec": 30.0, + "env.get_ee_pose": 30.0, + "env.get_scene_camera_meta": 30.0, + "env.move_to": 120.0, + "env.move_joints_delta": 60.0, + "env.get_obs": 60.0, +} + + +class LerobotEnvClient: + """Remote stub for the SO101 LeRobot env protocol.""" + + def __init__(self, client: RpcClient): + self._client = client + self._spec: dict | None = None + + def reset(self) -> tuple[dict, Any]: + """Drive the arm to its rest pose and return ``(obs, info)``. + + ``obs`` is ``{"state": {"joint_position": (5,) float32, + "gripper_position": (1,) float32}, "frames": {: (H, W, 3) + uint8, ...}}``. + """ + return self._client.call("env.reset", timeout_s=_TIMEOUT_S["env.reset"]) + + def step(self, action) -> tuple[dict, float, bool, bool, dict]: + """Send one absolute joint-target command ``[q1..q5, gripper]``. + + Returns ``(obs, reward, terminated, truncated, info)``. Targets are + clipped to the arm's joint limits server-side. + """ + action = np.asarray(action, dtype=np.float32).reshape(-1) + return self._client.call( + "env.step", args=(action,), timeout_s=_TIMEOUT_S["env.step"] + ) + + def get_spec(self) -> dict: + """Return (and cache) the env's static self-description. + + Keys: ``action_dim``, ``arm_joints``, ``action_low``, + ``action_high``, ``camera_names``, ``max_episode_steps``. + """ + if self._spec is None: + self._spec = self._client.call( + "env.get_spec", timeout_s=_TIMEOUT_S["env.get_spec"] + ) + return self._spec + + def get_ee_pose(self) -> dict: + """Live FK: gripper pose in the base (world) frame. + + Returns ``{xyz, quat_wxyz, joints_deg, T_base_gripper, frame}`` (or + ``{error}`` if FK is unavailable on the driver). + """ + return self._client.call( + "env.get_ee_pose", timeout_s=_TIMEOUT_S["env.get_ee_pose"] + ) + + def get_scene_camera_meta(self) -> dict: + """Scene-camera intrinsics + depth scale + base extrinsic. + + Returns ``{serial, K, width, height, depth_scale, frame, calibrated, + T_base_cam}`` (``T_base_cam`` is ``None`` until calibrated). + """ + return self._client.call( + "env.get_scene_camera_meta", + timeout_s=_TIMEOUT_S["env.get_scene_camera_meta"], + ) + + def move_to( + self, + xyz, + *, + gripper: float | None = None, + approach: str = "free", + yaw_deg: float | None = None, + ) -> dict: + """Move the gripper to a world-frame (base_link) XYZ via IK. + + ``approach="free"`` uses position-only IK (wrist orientation free); + ``approach="down"`` drives the gripper to point straight down for + grasping, with ``yaw_deg`` setting the jaw-line heading (auto-searched + when ``None``). The target is clipped to the workspace box and + approached in capped waypoints server-side; holds the current gripper + unless ``gripper`` is given. Returns the move log (``reached``, + ``pos_error_m``, ``approach_tilt_deg``, ...). + """ + return self._client.call( + "env.move_to", args=(xyz,), + kwargs={"gripper": gripper, "approach": approach, "yaw_deg": yaw_deg}, + timeout_s=_TIMEOUT_S["env.move_to"], + ) + + def move_joints_delta( + self, + delta_deg, + *, + gripper_delta: float | None = None, + ) -> dict: + """Nudge each arm joint by a relative amount (degrees). + + ``delta_deg`` is ``[d_pan, d_lift, d_elbow, d_wrist_flex, d_wrist_roll]`` + added to the current joints (each capped + clamped to limits + server-side). Optionally nudge the gripper by ``gripper_delta``. Returns + the achieved joints + EE pose. Use for fine alignment move_to cannot + express. + """ + return self._client.call( + "env.move_joints_delta", args=(delta_deg,), + kwargs={"gripper_delta": gripper_delta}, + timeout_s=_TIMEOUT_S["env.move_joints_delta"], + ) + + def get_obs(self) -> dict: + """Fetch the current observation without moving the arm.""" + return self._client.call("env.get_obs", timeout_s=_TIMEOUT_S["env.get_obs"]) diff --git a/robots/lerobot/env_server.py b/robots/lerobot/env_server.py new file mode 100644 index 00000000..f1dc303e --- /dev/null +++ b/robots/lerobot/env_server.py @@ -0,0 +1,1325 @@ +"""Standalone LeRobot SO101 env host for the LLM-in-the-loop agent (env-only). + +Drives a physical SO101 follower arm through LeRobot's synchronous Python +API (:class:`lerobot.robots.so_follower.SO101Follower`) and exposes a minimal +``reset`` / ``step`` gym-style surface over an RPC server +(:class:`rpent.utils.rpc.RpcFacade`, http transport by default) — the same +wire protocol the LIBERO driver uses, so the agent side talks to both +identically. + +Unlike the LIBERO driver, this server does **not** wrap an RLinf env class: +importing ``rlinf.envs.realworld`` runs node-level ROS setup side effects at +import time (it kills ``roscore``/``rosmaster`` processes), which is +inappropriate for a standalone driver. Instead this file talks to LeRobot +directly, reusing the device / action / observation recipe from RLinf's +``rlinf.envs.realworld.so101.SO101Env``. + +Run it inside the ``lerobot`` conda env:: + + conda activate lerobot + python robots/lerobot/env_server.py --output-dir /tmp/so101_run + +Hardware defaults match the current bench setup: follower on ``/dev/ttyACM1`` +(calibration id ``my_awesome_follower_arm``), an OpenCV hand/arm camera on +``/dev/video2``, and an Intel RealSense D405 scene camera (serial +``409122274720``). Every default is overridable from the CLI. + +Spawned by ``rpent/cli/main.py`` when ``--env lerobot`` is selected (requires +the ``lerobot`` extra in the agent venv), or run manually in the ``lerobot`` +conda env and attached to via ``--env-endpoint http://host:port``. +""" +from __future__ import annotations + +import argparse +import os +import signal +import sys +import time +from pathlib import Path + +import numpy as np + +# Make ``rpent`` importable when this file is run as a script from +# the ``lerobot`` conda env (which need not have rpent installed). +_PHYSICALAGENT_ROOT = Path(__file__).resolve().parents[2] +if str(_PHYSICALAGENT_ROOT) not in sys.path: + sys.path.insert(0, str(_PHYSICALAGENT_ROOT)) + +from robots.lerobot import calibration as scene_calib # noqa: E402 +from robots.lerobot import geometry as geom # noqa: E402 +from robots.lerobot.kinematics import SO101Kinematics # noqa: E402 +from robots.lerobot.scene_camera import SceneCameraD405 # noqa: E402 +from rpent.utils.logging import get_logger, init_output_dir # noqa: E402 +from rpent.utils.rpc import RpcFacade # noqa: E402 + +logger = get_logger("lerobot_driver") + + +# --------------------------------------------------------------------------- +# SO101 joint / limit constants (mirrors rlinf SO101Env defaults) +# --------------------------------------------------------------------------- + +# Arm joints in bus-ID order; LeRobot keys obs/action by ``.pos``. +_ARM_JOINTS = ( + "shoulder_pan", + "shoulder_lift", + "elbow_flex", + "wrist_flex", + "wrist_roll", +) +_GRIPPER = "gripper" +_NUM_ARM_JOINTS = len(_ARM_JOINTS) + +# Arm joint limits in degrees (arm only, no gripper), aligned to the SO101 URDF +# (deg = rad * 180/pi). Matching the URDF matters for ``move_to``'s top-down +# mode: its IK solutions use wrist_flex up to ~95 deg and wrist_roll beyond +# +/-150 deg, so a tighter clip here would silently alter the solved pose (and +# re-tilt the gripper). placo already enforces these limits; this clip is a +# secondary safety net. +_JOINT_LIMIT_LOW_DEG = np.array([-110.0, -100.0, -96.8, -95.0, -157.2], dtype=np.float32) +_JOINT_LIMIT_HIGH_DEG = np.array([110.0, 100.0, 96.8, 95.0, 162.8], dtype=np.float32) +_GRIPPER_LIMIT_LOW = 0.0 +_GRIPPER_LIMIT_HIGH = 90.0 +# Home pose ``[q1..q5, gripper]``. The five arm joints use the all-zero +# calibrated pose (as requested). The gripper homes to a mid opening, NOT 0: +# driving the gripper to 0 stalls it closed against its mechanical stop and +# trips the motor's overload protection (observed: gripper motor id 6 dropped +# off the bus after a 0 command). Keep this strictly between the limits. +_RESET_QPOS = np.array([0.0, -100.0, 90.0, 65.0, 0.0, 50.0], dtype=np.float32) + +# Conservative tabletop workspace box in the base (world) frame, meters. +# ``move_to`` clips targets to this so an LLM-supplied coordinate can't drive +# the arm into the table / out of reach. Sized to the SO101's ~0.35 m reach and +# the observed pick region (the base origin sits above the table, so the plate/ +# table surface is around z ~ -0.06): x forward, y lateral, z up. move_to now +# closes the loop and actually reaches the commanded z, so the z floor must stop +# the fingertips just ABOVE the plate (grasps straddle the cube's upper body). +# Tune for your bench. +_WORKSPACE_MIN = np.array([0.08, -0.28, -0.055], dtype=np.float64) +_WORKSPACE_MAX = np.array([0.38, 0.28, 0.30], dtype=np.float64) + +# --- grasp control point (TCP) ------------------------------------------ +# IK/FK use the ``gripper_frame_link`` frame, which actually sits on the FIXED +# jaw's fingertip -- ~2.5 cm ABOVE the fingertip plane and offset toward the +# fixed side, NOT between the fingers. This vector (meters, in the +# gripper_frame_link LOCAL frame) shifts the controlled/reported point to the +# grasp point between the fingertips, so ``move_to`` targets and ``get_ee_pose`` +# refer to where an object is actually grasped. TUNE on hardware / via touch calibration +# if grasps land consistently off-centre; set to zeros for the raw frame. +_TCP_OFFSET_GRIPPER = np.array([0.025, -0.01, 0.02], dtype=np.float64) + +# --- motion speed / smoothness (safety) --------------------------------- +# The arm runs in position mode. Two knobs keep motions slow and gentle: +# * Servo acceleration: LeRobot's configure() maxes the Feetech "Acceleration" +# register at 254 (snappy). We override it with a gentler value (0-254; +# lower = softer ramps), applied to every motor after connect. +# * Software velocity cap: point-to-point motions (reset, move_to, +# move_joints_delta) are streamed as interpolated setpoints so no joint +# exceeds ``_MAX_JOINT_VEL_DEG_S`` (deg/s), instead of snapping to the target +# at full servo speed. Both are overridable from the CLI. +_MOTOR_ACCELERATION = 30 +_MAX_JOINT_VEL_DEG_S = 60.0 +_PACE_DT_S = 0.05 # setpoint-streaming period (20 Hz) +# Feetech position gain. LeRobot's configure() lowers P_Coefficient to 16 (from +# the factory default 32) "to avoid shakiness"; that soft gain lets gravity- +# loaded joints (esp. elbow_flex) under-reach their commanded angle, so the +# gripper lands short and low. We raise it so the arm actually holds commanded +# joints. Overridable from the CLI (raise for stiffer tracking; lower if the +# arm oscillates/buzzes). None = leave LeRobot's value untouched. +_POSITION_GAIN = 32 + +# move_to interpolates the full gripper POSE into fine steps (straight-line +# position + slerped orientation) and solves IK warm-started at each, so the tip +# tracks the Cartesian line and the wrist reorients smoothly -- avoiding wrist +# IK-branch flips that otherwise swing the long gripper through the table. +_CART_STEP_M = 0.01 # max Cartesian move per interpolated IK step +_REORIENT_STEP_DEG = 6.0 # max orientation change per interpolated IK step +# If an interpolated step's IK solution jumps more than this (deg) from the +# previous one, the path crossed an IK discontinuity (near-singular / infeasible +# top-down pose). move_to stops at the last safe pose rather than streaming the +# arm through the swing, which could drive the long gripper into the table. +_MAX_STEP_JOINT_JUMP_DEG = 20.0 +# Closed-loop correction: after the open-loop path, the real servos can settle +# short of the commanded pose (gravity sag under load), so move_to re-commands +# with Cartesian feed-forward until the achieved tip is within _CORRECTION_TOL_M +# or the budget runs out. Feed-forward is capped and workspace-clipped for +# safety (it can never drive below the workspace floor / into the plate). +_MAX_POSITION_CORRECTIONS = 3 +_CORRECTION_TOL_M = 0.008 +_MAX_CORRECTION_M = 0.06 + + +def _build_camera_configs(raw: dict[str, dict]) -> dict: + """Build LeRobot ``CameraConfig`` instances from user-facing dicts. + + Each value must contain a ``"type"`` key (``"opencv"`` or + ``"intelrealsense"``); the remaining keys are forwarded to the matching + ``CameraConfig`` subclass. Mirrors ``SO101Env._build_camera_configs``: + the config subclasses self-register on import, so the two camera modules + must be imported before ``get_choice_class`` can resolve the type name. + """ + import lerobot.cameras.opencv.configuration_opencv # noqa: F401 registers "opencv" + import lerobot.cameras.realsense.configuration_realsense # noqa: F401 registers "intelrealsense" + from lerobot.cameras.configs import CameraConfig + + result: dict = {} + for name, cfg in raw.items(): + cfg = dict(cfg) + cam_type = cfg.pop("type") + result[name] = CameraConfig.get_choice_class(cam_type)(**cfg) + return result + + +def _to_lerobot_action(arm_targets: np.ndarray, gripper_target: float) -> dict: + """Build the ``{motor}.pos`` dict LeRobot's ``send_action`` expects. + + LeRobot filters incoming keys via ``key.endswith(".pos")``; a missing + suffix silently drops that motor. + """ + action = {f"{name}.pos": float(arm_targets[i]) for i, name in enumerate(_ARM_JOINTS)} + action[f"{_GRIPPER}.pos"] = float(gripper_target) + return action + + +def _topdown_rotation(yaw: float) -> np.ndarray: + """Target gripper rotation for a top-down approach (in the base frame). + + Maps the gripper's local approach axis (local +z, the wrist->fingertips + direction) to world -z (straight down). ``yaw`` (radians) rotates the + gripper about the vertical: the gripper's local +x maps to the horizontal + direction ``(cos yaw, sin yaw, 0)``, which sets the jaw-line heading. + """ + c, s = np.cos(yaw), np.sin(yaw) + return np.array( + [[c, s, 0.0], [s, -c, 0.0], [0.0, 0.0, -1.0]], dtype=np.float64 + ) + + +def _grasp_point(T_base_gripper: np.ndarray) -> np.ndarray: + """World xyz of the grasp point (between the fingertips) for a + ``gripper_frame_link`` pose, applying ``_TCP_OFFSET_GRIPPER`` in the gripper + frame. ``gripper_frame_link`` itself is on the fixed jaw, so this is what + ``move_to`` should target and ``get_ee_pose`` should report for grasping. + """ + T = np.asarray(T_base_gripper, dtype=np.float64) + return T[:3, 3] + T[:3, :3] @ _TCP_OFFSET_GRIPPER + + +def _approach_tilt_deg(R: np.ndarray) -> float: + """Angle (deg) between the gripper approach axis (local +z) and world -z. + + 0 deg = pointing straight down. Used to verify a top-down ``move_to``. + """ + z_world = np.asarray(R, dtype=np.float64) @ np.array([0.0, 0.0, 1.0]) + return float(np.degrees(np.arccos(np.clip(-z_world[2], -1.0, 1.0)))) + + +def _rotation_angle_deg(R0: np.ndarray, R1: np.ndarray) -> float: + """Geodesic angle (deg) between two rotation matrices.""" + cos = ( + np.trace(np.asarray(R0, dtype=np.float64).T @ np.asarray(R1, dtype=np.float64)) + - 1.0 + ) / 2.0 + return float(np.degrees(np.arccos(np.clip(cos, -1.0, 1.0)))) + + +def _slerp_rotation(R0: np.ndarray, R1: np.ndarray, t: float) -> np.ndarray: + """Interpolate rotation ``R0`` -> ``R1`` by fraction ``t`` in [0, 1]. + + Rotates about the fixed axis of the relative rotation (Rodrigues) -- a + matrix slerp -- so the gripper reorients along one smooth shortest arc. + """ + R0 = np.asarray(R0, dtype=np.float64) + R1 = np.asarray(R1, dtype=np.float64) + R_rel = R0.T @ R1 + ang = np.arccos(np.clip((np.trace(R_rel) - 1.0) / 2.0, -1.0, 1.0)) + if ang < 1e-8: + return R0.copy() + axis = np.array( + [ + R_rel[2, 1] - R_rel[1, 2], + R_rel[0, 2] - R_rel[2, 0], + R_rel[1, 0] - R_rel[0, 1], + ], + dtype=np.float64, + ) / (2.0 * np.sin(ang)) + th = ang * float(t) + K = np.array( + [ + [0.0, -axis[2], axis[1]], + [axis[2], 0.0, -axis[0]], + [-axis[1], axis[0], 0.0], + ], + dtype=np.float64, + ) + return R0 @ (np.eye(3) + np.sin(th) * K + (1.0 - np.cos(th)) * (K @ K)) + + +class SO101LeRobotEnv: + """Minimal ``reset`` / ``step`` driver for a physical SO101 arm. + + Action: ``(6,)`` float array ``[q1..q5, gripper]`` of absolute joint + position targets in degrees. Arm targets are clipped to the configured + joint limits and the gripper to ``[0, 90]`` before being sent to the + motor bus. + + Observation:: + + {"state": {"joint_position": (5,) float32, # arm joints, degrees + "gripper_position": (1,) float32, # gripper opening + "ee_pose_base": (3,) float32, # gripper xyz in base (FK) + "ee_quat_base": (4,) float32}, # gripper quat wxyz (FK) + "frames": {"arm": (H, W, 3) uint8, "scene": (H, W, 3) uint8}, + "depth": {"scene": (H, W) float32}} # metric, aligned to color + + ``ee_pose_base`` / ``ee_quat_base`` are present only when FK is available; + ``scene`` frames/depth only when the scene camera is configured. All values + are plain numpy / floats so they pickle across the RPC wire (the agent + process does not import torch). + """ + + def __init__( + self, + *, + port: str, + calibration_id: str, + arm_camera_cfgs: dict[str, dict], + scene_serial: str | None = None, + scene_size: tuple[int, int] = (720, 1280), + scene_fps: int = 30, + urdf_path: str | None = None, + max_relative_target: float | None = None, + max_episode_steps: int = 200, + step_frequency: float = 30.0, + motor_acceleration: int | None = _MOTOR_ACCELERATION, + max_joint_vel_deg_s: float = _MAX_JOINT_VEL_DEG_S, + position_gain: int | None = _POSITION_GAIN, + auto_calibrate: bool = False, + ) -> None: + self._max_episode_steps = max_episode_steps + self._step_frequency = step_frequency + self._max_joint_vel_deg_s = float(max_joint_vel_deg_s) + self._pace_dt = _PACE_DT_S + self._num_steps = 0 + self._action_low = np.append(_JOINT_LIMIT_LOW_DEG, _GRIPPER_LIMIT_LOW).astype(np.float32) + self._action_high = np.append(_JOINT_LIMIT_HIGH_DEG, _GRIPPER_LIMIT_HIGH).astype(np.float32) + + from lerobot.robots.so_follower import SO101Follower + from lerobot.robots.so_follower.config_so_follower import SO101FollowerConfig + + robot_cfg = SO101FollowerConfig( + port=port, + id=calibration_id, + use_degrees=True, + max_relative_target=max_relative_target, + # Keep torque on at disconnect so the arm holds its parked pose. + disable_torque_on_disconnect=False, + cameras=_build_camera_configs(arm_camera_cfgs), + ) + self._robot = SO101Follower(robot_cfg) + # ``calibrate=False`` loads the on-disk calibration without ever + # prompting on stdin (a server must never block on input). + self._robot.connect(calibrate=auto_calibrate) + self._arm_camera_names = tuple(arm_camera_cfgs.keys()) + + # Soften motion: LeRobot's configure() sets the Feetech "Acceleration" + # register to its max (254). Override with a gentler value so the arm + # ramps smoothly rather than snapping (safety). Done with torque briefly + # off, mirroring LeRobot's own register writes. + self._motor_acceleration = motor_acceleration + if motor_acceleration is not None: + try: + with self._robot.bus.torque_disabled(): + for motor in self._robot.bus.motors: + # Keep the gripper snappy so grasps close promptly; only + # slow the (heavier, safety-relevant) arm joints. + if motor == _GRIPPER: + continue + self._robot.bus.write( + "Acceleration", motor, int(motor_acceleration) + ) + except Exception as e: + logger.warning( + "could not set motor Acceleration=%s (motions stay fast): %s", + motor_acceleration, e, + ) + + # Stiffen position holding: LeRobot lowers P_Coefficient to 16, which + # lets gravity-loaded arm joints under-reach their target. Raise it on + # the arm joints (leave the gripper as LeRobot set it). + self._position_gain = position_gain + if position_gain is not None: + try: + with self._robot.bus.torque_disabled(): + for motor in self._robot.bus.motors: + if motor == _GRIPPER: + continue + self._robot.bus.write( + "P_Coefficient", motor, int(position_gain) + ) + except Exception as e: + logger.warning( + "could not set motor P_Coefficient=%s (tracking may be soft): %s", + position_gain, e, + ) + + # Scene camera (fixed, depth) is managed directly via pyrealsense2 so we + # get depth aligned to color + intrinsics (LeRobot's wrapper gives + # neither). The arm camera stays under LeRobot above (color only). + self._scene_serial = scene_serial or None + self._scene_cam: SceneCameraD405 | None = None + if self._scene_serial: + scene_h, scene_w = scene_size + self._scene_cam = SceneCameraD405( + self._scene_serial, width=scene_w, height=scene_h, fps=scene_fps, + ) + + # Forward kinematics for end-effector pose in the base (world) frame. + self._kin: SO101Kinematics | None = None + try: + self._kin = SO101Kinematics(urdf_path=urdf_path) + except Exception as e: + logger.warning("FK unavailable (%s); ee_pose will be omitted", e) + + # Scene-cam -> base extrinsic, if it has been calibrated (touch/Kabsch). + # A fit whose saved RMSE is too large is REJECTED (treated as + # uncalibrated): a bad extrinsic makes back_project return badly wrong + # world coordinates, so the arm would chase unreachable targets. The + # operator must recalibrate rather than have the agent act on garbage. + self._T_base_cam = None + self._calib_rmse_m: float | None = None + if self._scene_serial: + record = scene_calib.load_extrinsic_record(self._scene_serial) + if record is not None: + self._calib_rmse_m = record.get("rmse_m") + rmse = self._calib_rmse_m + if rmse is not None and rmse > scene_calib.MAX_ACCEPTABLE_RMSE_M: + logger.warning( + "scene-cam extrinsic REJECTED: rmse=%.3fm > %.3fm limit; " + "back_project would be unreliable. Recalibrate with " + "robots/lerobot/auto_calibrate_scene_cam.py " + "(or calibrate_scene_cam.py).", + rmse, scene_calib.MAX_ACCEPTABLE_RMSE_M, + ) + else: + self._T_base_cam = record["T_base_cam"] + + if self._T_base_cam is not None: + extr_status = ( + f"loaded (rmse={self._calib_rmse_m:.3f}m)" + if self._calib_rmse_m is not None else "loaded" + ) + elif self._calib_rmse_m is not None: + extr_status = f"rejected (rmse={self._calib_rmse_m:.3f}m)" + else: + extr_status = "uncalibrated" + + cam_list = list(self._arm_camera_names) + (["scene"] if self._scene_cam else []) + logger.info( + "SO101 connected on %s (calibration_id=%s); cameras=[%s]; " + "FK=%s; scene_extrinsic=%s", + port, calibration_id, ", ".join(cam_list) or "none", + "on" if self._kin else "off", + extr_status, + ) + logger.info( + "motion pacing: max_joint_vel=%.0f deg/s; motor_acceleration=%s", + self._max_joint_vel_deg_s, + self._motor_acceleration if self._motor_acceleration is not None + else "default (254)", + ) + logger.info( + "position gain: P_Coefficient=%s", + self._position_gain if self._position_gain is not None + else "default (LeRobot 16)", + ) + + # ------------------------------------------------------------------ + # gym-like surface + # ------------------------------------------------------------------ + + def _stream_joint_path(self, q_from, q_to, gripper_value) -> None: + """Glide the arm from ``q_from`` to ``q_to`` (arm-joint vectors, degrees). + + Streams interpolated joint setpoints so no joint moves faster than + ``self._max_joint_vel_deg_s`` (deg/s), instead of commanding the target + in one shot and letting the servos snap there at full speed. The gripper + is held at ``gripper_value`` throughout. + """ + q_from = np.asarray(q_from, dtype=np.float64).reshape(-1)[:_NUM_ARM_JOINTS] + q_to = np.asarray(q_to, dtype=np.float64).reshape(-1)[:_NUM_ARM_JOINTS] + max_delta = float(np.max(np.abs(q_to - q_from))) if q_to.size else 0.0 + per_step = max(1e-6, self._max_joint_vel_deg_s * self._pace_dt) + n = max(1, int(np.ceil(max_delta / per_step))) + for i in range(1, n + 1): + q = q_from + (q_to - q_from) * (i / n) + self._robot.send_action(_to_lerobot_action(q, gripper_value)) + time.sleep(self._pace_dt) + + def reset(self) -> tuple[dict, dict]: + """Send the arm to its rest pose (paced), reset the counter, return obs.""" + self._num_steps = 0 + obs = self._robot.get_observation() + cur = np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64 + ) + self._stream_joint_path( + cur, _RESET_QPOS[:_NUM_ARM_JOINTS], float(_RESET_QPOS[_NUM_ARM_JOINTS]) + ) + time.sleep(0.3) + return self._get_observation(), {} + + def step(self, action) -> tuple[dict, float, bool, bool, dict]: + """Command one absolute joint target and return the gym 5-tuple. + + Returns ``(obs, reward, terminated, truncated, info)``. This minimal + driver computes no task reward (``reward == 0.0``) and never + auto-terminates; ``truncated`` flips once ``max_episode_steps`` is + reached. Task success / stopping is decided by the agent. + """ + t0 = time.time() + action = np.asarray(action, dtype=np.float32).reshape(-1) + expected = _NUM_ARM_JOINTS + 1 + if action.shape[0] != expected: + raise ValueError( + f"action must have {expected} entries [q1..q5, gripper]; " + f"got {action.shape[0]}" + ) + action = np.clip(action, self._action_low, self._action_high) + self._robot.send_action( + _to_lerobot_action(action[:_NUM_ARM_JOINTS], float(action[_NUM_ARM_JOINTS])) + ) + self._num_steps += 1 + + obs = self._get_observation() + truncated = self._num_steps >= self._max_episode_steps + + # Pace the control loop to the requested frequency. + dt = time.time() - t0 + time.sleep(max(0.0, (1.0 / self._step_frequency) - dt)) + return obs, 0.0, False, truncated, {} + + def get_spec(self) -> dict: + """Static self-description for the agent side (action bounds, cams).""" + cam_names = list(self._arm_camera_names) + (["scene"] if self._scene_cam else []) + return { + "action_dim": _NUM_ARM_JOINTS + 1, + "arm_joints": list(_ARM_JOINTS), + "action_low": self._action_low.tolist(), + "action_high": self._action_high.tolist(), + "camera_names": cam_names, + "scene_camera": "scene" if self._scene_cam else None, + "has_ee_pose": self._kin is not None, + "world_frame": "base_link", + "max_episode_steps": self._max_episode_steps, + } + + # ------------------------------------------------------------------ + # localization surface (base frame == world) + # ------------------------------------------------------------------ + + def get_ee_pose(self) -> dict: + """Live FK: gripper pose in the base (world) frame.""" + obs = self._robot.get_observation() + joints = np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float32 + ) + ee = self._compute_ee_pose(joints) + if ee is None: + return {"error": "FK unavailable (URDF/placo missing)"} + return { + "xyz": _grasp_point(ee["T"]).tolist(), + "quat_wxyz": ee["quat"].tolist(), + "joints_deg": joints.tolist(), + "T_base_gripper": ee["T"].tolist(), + "frame": "base_link", + } + + def get_scene_camera_meta(self) -> dict: + """Scene-camera intrinsics + depth scale + base extrinsic (if any).""" + if self._scene_cam is None: + return {"error": "scene camera not configured"} + meta = self._scene_cam.meta() + meta["frame"] = "scene_cam" + meta["calibrated"] = self._T_base_cam is not None + meta["rmse_m"] = self._calib_rmse_m + meta["T_base_cam"] = ( + self._T_base_cam.tolist() if self._T_base_cam is not None else None + ) + return meta + + def get_scene_frame(self) -> dict: + """Live scene color + metric depth (for calibration / ad-hoc queries).""" + if self._scene_cam is None: + return {"error": "scene camera not configured"} + rgb, depth = self._scene_cam.read() + return {"color": rgb, "depth": depth, "K": self._scene_cam.K.tolist()} + + def get_obs(self) -> dict: + """Return the current observation without moving the arm. + + Used to refresh the agent-side cache after a primitive (e.g. move_to) + that changes the world but does not itself return an observation. + """ + return self._get_observation() + + def set_torque(self, enabled: bool) -> dict: + """Enable/disable arm motor torque. + + Disabling lets the operator free-drive the arm by hand (used by the + touch/Kabsch scene-camera calibration). Joint encoders remain readable + with torque off, so FK still works. Re-enable to hold position. + """ + bus = self._robot.bus + try: + if enabled: + bus.enable_torque() + else: + bus.disable_torque() + except Exception as e: + return {"ok": False, "error": str(e)} + logger.info("arm torque %s", "enabled" if enabled else "disabled") + return {"ok": True, "torque_enabled": bool(enabled)} + + def _read_arm_joints(self) -> np.ndarray: + """Current arm joint angles (deg) from a fresh observation.""" + obs = self._robot.get_observation() + return np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64 + ) + + def _servo_pose_path( + self, from_joints, R_start, p_start, R_end, p_end, grip_val, orient_w, + settle_s, + ) -> bool: + """Stream the gripper along an interpolated pose (straight-line position + + slerped orientation), warm-starting IK per fine step so solutions stay + on ONE branch and the tip tracks the line. Returns ``halted`` (an IK + discontinuity forced an early stop). Settles ``settle_s`` before + returning so the caller can measure the achieved pose. + """ + from_joints = np.asarray(from_joints, dtype=np.float64) + p_start = np.asarray(p_start, dtype=np.float64) + p_end = np.asarray(p_end, dtype=np.float64) + dist = float(np.linalg.norm(p_end - p_start)) + reorient_deg = _rotation_angle_deg(R_start, R_end) if orient_w > 0 else 0.0 + n_steps = max( + 1, + int(np.ceil(dist / _CART_STEP_M)), + int(np.ceil(reorient_deg / _REORIENT_STEP_DEG)), + ) + seed = from_joints.copy() + prev_q = from_joints.copy() + halted = False + for i in range(1, n_steps + 1): + frac = i / n_steps + T_des = np.eye(4) + T_des[:3, :3] = ( + _slerp_rotation(R_start, R_end, frac) if orient_w > 0 else R_end + ) + T_des[:3, 3] = p_start + (p_end - p_start) * frac + q = self._kin.ik( + seed, T_des, position_weight=1.0, orientation_weight=orient_w + ) + q_arm = np.clip( + q[:_NUM_ARM_JOINTS], _JOINT_LIMIT_LOW_DEG.astype(np.float64), + _JOINT_LIMIT_HIGH_DEG.astype(np.float64), + ) + # A large jump between consecutive fine-step solutions = the path + # crossed an IK discontinuity (near-singular / infeasible top-down + # pose); stop at the last safe pose rather than swinging through it. + if float(np.max(np.abs(q_arm - prev_q))) > _MAX_STEP_JOINT_JUMP_DEG: + halted = True + break + seed = q_arm + self._stream_joint_path(prev_q, q_arm, grip_val) + prev_q = q_arm + time.sleep(settle_s) + return halted + + def move_to( + self, + xyz, + *, + gripper: float | None = None, + approach: str = "free", + yaw_deg: float | None = None, + settle_s: float = 0.4, + pos_tol_m: float = 0.02, + tilt_tol_deg: float = 15.0, + max_corrections: int = _MAX_POSITION_CORRECTIONS, + ) -> dict: + """Move the gripper to a world-frame (base_link) XYZ via IK. + + ``approach`` selects the wrist-orientation policy: + + * ``"free"`` (default): position-only IK; the wrist settles at whatever + orientation placo converges to. Maximal reach, but the fingertips' + location relative to the returned EE point is unpredictable (the TCP + is ~0.1 m out along the gripper axis), so it is unreliable for + grasping. + * ``"down"``: top-down IK -- the gripper approach axis is driven to + vertical (pointing straight down) so the fingertips descend along + world -z, which makes grasping predictable. ``yaw_deg`` sets the + jaw-line heading about the vertical (0 = +x/forward); if ``None`` a + reachable yaw is searched automatically. + + The target is clipped to the workspace box and approached by + interpolating the full gripper pose (straight-line position + smoothly + slerped orientation) into fine, warm-started IK steps. It then CLOSES + THE LOOP: it measures the achieved tip and re-commands with feed-forward + (up to ``max_corrections`` times) to cancel the servos' steady-state sag + under load, so the tip lands on the commanded xyz -- callers should pass + the true target, NOT a hand-tuned over-shoot. Holds the current gripper + opening unless ``gripper`` is given. + + Returns a log dict with ``reached`` (position within ``pos_tol_m`` and, + for ``"down"``, approach tilt within ``tilt_tol_deg``), the + commanded/achieved xyz, the position error, and the achieved approach + tilt in degrees. + """ + if self._kin is None: + return {"error": "IK unavailable (URDF/placo missing)"} + approach = str(approach).lower() + if approach not in ("free", "down"): + return {"error": f"approach must be 'free' or 'down'; got {approach!r}"} + target = np.asarray(xyz, dtype=np.float64).reshape(-1) + if target.shape[0] != 3: + return {"error": f"xyz must be 3 numbers; got {target.shape[0]}"} + + clipped = np.clip(target, _WORKSPACE_MIN, _WORKSPACE_MAX) + was_clipped = not np.allclose(clipped, target, atol=1e-6) + + obs = self._robot.get_observation() + cur_joints = np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64 + ) + cur_gripper = float(obs.get(f"{_GRIPPER}.pos", 0.0)) + grip_val = cur_gripper if gripper is None else float(gripper) + + T_cur = self._kin.fk(cur_joints) + p0 = T_cur[:3, 3] + R0 = T_cur[:3, :3] + + # Target orientation for the move: reorient to vertical for "down"; + # leave it to the IK (weight 0) for "free". + if approach == "down": + orient_w = 1.0 + yaw = ( + self._search_topdown_yaw(cur_joints, clipped) + if yaw_deg is None else float(np.radians(yaw_deg)) + ) + R_target = _topdown_rotation(yaw) + else: + orient_w = 0.0 + yaw = None + R_target = R0 # ignored (orientation_weight=0) + + # Held orientation: top-down for "down", the current (ignored) rotation + # for "free". Move the tip along the interpolated pose, then close the + # loop to cancel the servos' steady-state sag under load. + R_hold = R_target if approach == "down" else R0 + # move_to controls the GRASP POINT (between the fingertips). Convert that + # target into the gripper_frame_link goal the IK/servo actually drives. + # The offset is only well-defined for the fixed "down" orientation; + # "free" leaves the tip frame uncorrected (its final orientation, hence + # the offset direction, is unknown until IK converges). + tcp_off_world = ( + R_hold @ _TCP_OFFSET_GRIPPER if approach == "down" + else np.zeros(3, dtype=np.float64) + ) + halted = self._servo_pose_path( + cur_joints, R0, p0, R_hold, clipped - tcp_off_world, + grip_val, orient_w, settle_s, + ) + final_joints = self._read_arm_joints() + T_final = self._kin.fk(final_joints) + grasp_final = _grasp_point(T_final) if approach == "down" else T_final[:3, 3] + + # Closed-loop correction: the real arm can settle short of the commanded + # pose (gravity sag), even though the IK path is exact. Feed the residual + # forward and re-command (workspace-clipped + capped) until the achieved + # GRASP POINT is within tolerance or the correction budget runs out. + n_corr = 0 + ff = np.zeros(3, dtype=np.float64) + while ( + not halted + and n_corr < max_corrections + and float(np.linalg.norm(clipped - grasp_final)) > _CORRECTION_TOL_M + ): + ff = np.clip( + ff + (clipped - grasp_final), -_MAX_CORRECTION_M, _MAX_CORRECTION_M + ) + corr_goal = ( + np.clip(clipped + ff, _WORKSPACE_MIN, _WORKSPACE_MAX) - tcp_off_world + ) + halted = self._servo_pose_path( + final_joints, R_hold, T_final[:3, 3], R_hold, corr_goal, + grip_val, orient_w, settle_s, + ) + final_joints = self._read_arm_joints() + T_final = self._kin.fk(final_joints) + grasp_final = _grasp_point(T_final) if approach == "down" else T_final[:3, 3] + n_corr += 1 + + err = float(np.linalg.norm(grasp_final - clipped)) + tilt = _approach_tilt_deg(T_final[:3, :3]) + reached = ( + not halted + and err <= pos_tol_m + and (approach != "down" or tilt <= tilt_tol_deg) + ) + result = { + "reached": bool(reached), + "approach": approach, + "target_xyz": [round(float(v), 4) for v in target], + "commanded_xyz": [round(float(v), 4) for v in clipped], + "final_xyz": [round(float(v), 4) for v in grasp_final], + "pos_error_m": round(err, 4), + "approach_tilt_deg": round(tilt, 1), + "yaw_deg": (None if yaw is None else round(float(np.degrees(yaw)), 1)), + "clipped_to_workspace": was_clipped, + "pos_corrections": n_corr, + "halted_early": bool(halted), + "joints_deg": [round(float(v), 2) for v in final_joints], + "gripper": round(grip_val, 2), + } + if halted: + result["note"] = ( + "stopped partway: the straight-line top-down path crossed an " + "unreachable / near-singular pose. Try a nearer target, a " + "different yaw_deg, or move in smaller hops." + ) + elif not reached: + result["note"] = ( + f"settled {round(err * 1000)} mm short after {n_corr} " + "correction(s); the target may be past the arm's reach here. " + "Try a nearer / higher target." + ) + return result + + def _search_topdown_yaw(self, seed_joints, target, *, n_yaw: int = 12) -> float: + """Pick a top-down wrist yaw (rad) that reaches ``target`` best. + + Whether a vertical approach is reachable depends on the jaw-line yaw + (wrist_roll/shoulder_pan coupling), so we sweep candidate yaws, run IK + for each, and keep the one with the smallest FK position error (ties + broken by approach tilt). Pure CPU -- the arm does not move. + """ + best_yaw, best_key = 0.0, None + seed = np.asarray(seed_joints, dtype=np.float64) + for k in range(n_yaw): + yaw = 2.0 * np.pi * k / n_yaw + T_des = np.eye(4) + T_des[:3, :3] = _topdown_rotation(yaw) + T_des[:3, 3] = target + q = self._kin.ik(seed, T_des, position_weight=1.0, orientation_weight=1.0) + T_q = self._kin.fk(q[:_NUM_ARM_JOINTS]) + perr = float(np.linalg.norm(T_q[:3, 3] - target)) + tilt = _approach_tilt_deg(T_q[:3, :3]) + key = (round(perr, 4), round(tilt, 1)) + if best_key is None or key < best_key: + best_key, best_yaw = key, yaw + return best_yaw + + def move_joints_delta( + self, + delta_deg, + *, + gripper_delta: float | None = None, + max_step_deg: float = 15.0, + settle_s: float = 0.4, + ) -> dict: + """Nudge each arm joint by a relative amount (degrees). + + ``delta_deg`` is 5 values ``[d_pan, d_lift, d_elbow, d_wrist_flex, + d_wrist_roll]`` added to the current arm joints. Each entry is capped to + +/- ``max_step_deg`` and the result is clamped to the joint limits, so a + single call makes a small, safe adjustment. Optionally nudge the gripper + by ``gripper_delta`` (clamped to its limits). Use this for fine + alignment ``move_to`` cannot express -- e.g. tweak wrist_roll to line the + jaws up with an object, or descend a few millimetres -- reading the new + EE pose back from the result. + + Returns a log dict with the applied delta, achieved joints, gripper, and + (when FK is available) the new EE xyz and approach tilt. + """ + delta = np.asarray(delta_deg, dtype=np.float64).reshape(-1) + if delta.shape[0] != _NUM_ARM_JOINTS: + return { + "error": f"delta_deg must have {_NUM_ARM_JOINTS} entries " + f"[pan, lift, elbow, wrist_flex, wrist_roll]; got {delta.shape[0]}" + } + cap = abs(float(max_step_deg)) + delta = np.clip(delta, -cap, cap) + + obs = self._robot.get_observation() + cur_joints = np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64 + ) + cur_gripper = float(obs.get(f"{_GRIPPER}.pos", 0.0)) + + target_joints = np.clip( + cur_joints + delta, + _JOINT_LIMIT_LOW_DEG.astype(np.float64), + _JOINT_LIMIT_HIGH_DEG.astype(np.float64), + ) + if gripper_delta is None: + grip_val = cur_gripper + else: + grip_val = float( + np.clip(cur_gripper + float(gripper_delta), + _GRIPPER_LIMIT_LOW, _GRIPPER_LIMIT_HIGH) + ) + + self._stream_joint_path(cur_joints, target_joints, grip_val) + time.sleep(settle_s) + + final_obs = self._robot.get_observation() + final_joints = np.array( + [final_obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64 + ) + result: dict = { + "applied_delta_deg": [round(float(v), 2) for v in delta], + "joints_deg": [round(float(v), 2) for v in final_joints], + "gripper": round(grip_val, 2), + } + ee = self._compute_ee_pose(final_joints) + if ee is not None: + result["ee_xyz"] = [round(float(v), 4) for v in _grasp_point(ee["T"])] + result["approach_tilt_deg"] = round(_approach_tilt_deg(ee["T"][:3, :3]), 1) + return result + + # ------------------------------------------------------------------ + # automatic scene-camera calibration (markerless, gripper-motion) + # ------------------------------------------------------------------ + + def _set_gripper_hold(self, gripper_value: float) -> None: + """Set the gripper opening while freezing the arm at its current joints. + + Used during calibration so that between the two capture frames ONLY the + gripper fingers move (clean motion segmentation). + """ + obs = self._robot.get_observation() + q = np.array([obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64) + self._robot.send_action(_to_lerobot_action(q, float(gripper_value))) + + @staticmethod + def _calibration_targets() -> list[list[float]]: + """A wide, non-coplanar grid of tip targets inside the workspace. + + Free-orientation IK reaches these, giving a large spread in x/y/z AND a + variety of wrist orientations. That orientation variety is what makes + the constant tip-detector offset identifiable (see + ``geometry.solve_extrinsic_with_offset``), so the calibration keeps the + default free approach rather than a fixed top-down one. Ordered z-fastest + so an early stop (``n_points``) still spans all three heights. + """ + xs = [0.15, 0.21, 0.27, 0.33] + ys = [-0.16, -0.05, 0.05, 0.16] + zs = [0.08, 0.15, 0.22] + return [[x, y, z] for x in xs for y in ys for z in zs] + + def _capture_scene_median( + self, n_frames: int = 5 + ) -> tuple[np.ndarray, np.ndarray]: + """Grab several scene frames and return per-pixel temporal medians. + + Median-averaging across frames suppresses the camera's per-frame color + and depth noise, which otherwise becomes lateral error once a pixel is + back-projected through the oblique view. Depth invalids (<=0 / + non-finite) are ignored per pixel; pixels with no valid sample stay 0. + """ + import warnings + + rgbs: list[np.ndarray] = [] + depths: list[np.ndarray] = [] + for _ in range(max(1, int(n_frames))): + rgb, depth = self._scene_cam.read() + rgbs.append(np.asarray(rgb)) + d = np.asarray(depth, dtype=np.float32) + depths.append(np.where(np.isfinite(d) & (d > 0), d, np.nan)) + rgb_med = np.median(np.stack(rgbs, axis=0), axis=0).astype(np.uint8) + with warnings.catch_warnings(): # nanmedian warns on all-invalid pixels + warnings.simplefilter("ignore", category=RuntimeWarning) + depth_med = np.nanmedian(np.stack(depths, axis=0), axis=0) + depth_med = np.nan_to_num(depth_med, nan=0.0).astype(np.float32) + return rgb_med, depth_med + + def auto_calibrate_scene_camera( + self, + *, + n_points: int = 24, + gripper_open: float = 90.0, + gripper_closed: float = 20.0, + settle_s: float = 0.8, + ransac_thresh_m: float = 0.015, + save: bool = True, + ) -> dict: + """Markerless automatic scene-cam -> base calibration. + + Drives the tip to a wide, non-coplanar grid of base-frame positions + (move_to needs no extrinsic) at varied wrist orientations. At each pose + it toggles the gripper with the arm frozen and segments the motion in + the (temporally median-filtered) scene image to locate the moving jaw + (blob centroid + depth AT that centroid -> camera point); FK gives the + tip pose (origin + rotation). A joint fit + (:func:`geometry.solve_extrinsic_with_offset`) then recovers both + ``T_base_cam`` and the constant offset between the detected centroid and + the tip frame -- so that offset no longer inflates the residual -- and + the extrinsic is saved and hot-loaded so back_project returns world + coords immediately. + + Returns a summary dict (``n_used``, ``rmse_m``, ``tip_offset_local_m``, + per-pose diagnostics). + """ + if self._scene_cam is None: + return {"error": "scene camera not configured"} + if self._kin is None: + return {"error": "IK unavailable (URDF/placo missing)"} + + targets = self._calibration_targets() + cam_pts: list[list[float]] = [] + tip_origins: list[list[float]] = [] + tip_rots: list[list[list[float]]] = [] + poses: list[dict] = [] + + for tgt in targets: + if len(cam_pts) >= n_points: + break + mv = self.move_to(tgt, gripper=gripper_open, settle_s=settle_s) + if "error" in mv or not mv.get("reached"): + poses.append({"target": tgt, "skipped": "unreachable"}) + continue + time.sleep(settle_s) + + # Free-orientation FK tip pose (origin + rotation) at this pose. The + # rotation is what lets the fit solve out the constant tip-detector + # offset (geometry.solve_extrinsic_with_offset), so orientation must + # vary across the grid -- hence the free (not top-down) approach. + T_tip = self._kin.fk(self._read_arm_joints()) + o_i = T_tip[:3, 3] + R_i = T_tip[:3, :3] + + self._scene_cam.read() # drop the in-flight frame from the move + rgb_open, _ = self._capture_scene_median() + self._set_gripper_hold(gripper_closed) + time.sleep(settle_s) + rgb_closed, depth = self._capture_scene_median() + self._set_gripper_hold(gripper_open) # reopen for the next pose + + det = geom.detect_tip_pixel_by_motion( + rgb_open, rgb_closed, depth, self._scene_cam.K, + ) + if det is None: + poses.append({"target": tgt, "skipped": "no_tip_detected"}) + continue + cam_pts.append(det["xyz_cam"]) + tip_origins.append(o_i.tolist()) + tip_rots.append(R_i.tolist()) + poses.append({"target": tgt, "base_xyz": o_i.round(4).tolist(), + "pixel": [round(v, 1) for v in det["pixel"]], + "depth_m": round(det["depth_m"], 4), "area": det["area"]}) + + if len(cam_pts) < 4: + return {"error": f"only {len(cam_pts)} usable points (need >= 4)", + "poses": poses} + + # Joint fit: recover T_base_cam AND the constant offset between the + # detected motion-blob centroid and the FK tip frame, so that offset no + # longer pollutes the residual (the old ~1 cm RMSE floor). + T, rmse, inliers, tip_offset = geom.solve_extrinsic_with_offset( + cam_pts, tip_origins, tip_rots, thresh_m=ransac_thresh_m, + ) + accepted = bool(rmse <= scene_calib.MAX_ACCEPTABLE_RMSE_M) + result = { + "n_targets": len(targets), + "n_used": len(cam_pts), + "n_inliers": int(np.asarray(inliers).sum()), + "rmse_m": round(float(rmse), 4), + "tip_offset_local_m": [round(float(v), 4) for v in tip_offset], + "T_base_cam": T.tolist(), + "accepted": accepted, + "saved": False, + "poses": poses, + } + if not accepted: + # A high RMSE means the cam/base correspondences are inconsistent + # (poor tip detection, lighting, or occlusion). Saving it would + # silently corrupt every back_project, so refuse and ask for a rerun. + result["error"] = ( + f"calibration RMSE {rmse * 1000:.1f} mm exceeds the " + f"{scene_calib.MAX_ACCEPTABLE_RMSE_M * 1000:.0f} mm limit; not " + "saved. Clear the workspace, improve gripper visibility/lighting, " + "and rerun." + ) + logger.warning( + "scene-cam calibration REJECTED: rmse=%.4fm (> %.3fm); not saved", + rmse, scene_calib.MAX_ACCEPTABLE_RMSE_M, + ) + elif save: + path = scene_calib.save_extrinsic( + self._scene_serial, T, K=self._scene_cam.K, + rmse_m=rmse, num_points=int(np.asarray(inliers).sum()), + ) + self._T_base_cam = T # hot-load so back_project works immediately + self._calib_rmse_m = round(float(rmse), 4) + result["saved"] = True + result["path"] = str(path) + logger.info("scene-cam calibrated: rmse=%.4fm, saved %s", rmse, path) + + # Park the arm at rest after the sweep (paced, gentle). + try: + obs = self._robot.get_observation() + cur = np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64 + ) + self._stream_joint_path( + cur, _RESET_QPOS[:_NUM_ARM_JOINTS], float(_RESET_QPOS[_NUM_ARM_JOINTS]) + ) + except Exception: + pass + return result + + def close(self) -> None: + """Park the arm at rest (paced + torque held) and disconnect cleanly.""" + try: + obs = self._robot.get_observation() + cur = np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64 + ) + self._stream_joint_path( + cur, _RESET_QPOS[:_NUM_ARM_JOINTS], float(_RESET_QPOS[_NUM_ARM_JOINTS]) + ) + time.sleep(0.3) + except Exception as e: + logger.warning("failed to park arm on close: %s", e) + if self._scene_cam is not None: + try: + self._scene_cam.close() + except Exception as e: + logger.warning("error closing scene camera: %s", e) + try: + self._robot.disconnect() + logger.info("SO101 disconnected") + except Exception as e: + logger.warning("error disconnecting robot: %s", e) + + # ------------------------------------------------------------------ + # internals + # ------------------------------------------------------------------ + + def _get_observation(self) -> dict: + obs = self._robot.get_observation() + joint_position = np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float32 + ) + gripper_position = np.array([obs.get(f"{_GRIPPER}.pos", 0.0)], dtype=np.float32) + + frames: dict = {} + for cam in self._arm_camera_names: + frame = obs.get(cam) + if frame is not None: + frames[cam] = np.ascontiguousarray(np.asarray(frame, dtype=np.uint8)) + + depth: dict = {} + if self._scene_cam is not None: + scene_rgb, scene_depth = self._scene_cam.read() + frames["scene"] = scene_rgb + depth["scene"] = scene_depth + + state: dict = { + "joint_position": joint_position, + "gripper_position": gripper_position, + } + ee = self._compute_ee_pose(joint_position) + if ee is not None: + state["ee_pose_base"] = _grasp_point(ee["T"]).astype(np.float32) + state["ee_quat_base"] = ee["quat"] + + out: dict = {"state": state, "frames": frames} + if depth: + out["depth"] = depth + return out + + def _compute_ee_pose(self, joints_deg) -> dict | None: + """FK -> gripper pose in the base (world) frame, or None if FK is off.""" + if self._kin is None: + return None + try: + T = self._kin.fk(joints_deg) + except Exception as e: + logger.warning("FK failed: %s", e) + return None + xyz = T[:3, 3].astype(np.float32) + quat = geom.rotation_to_quat(T[:3, :3]).astype(np.float32) + return {"xyz": xyz, "quat": quat, "T": T} + + +# --------------------------------------------------------------------------- +# RPC plumbing +# --------------------------------------------------------------------------- + +class LerobotEnvFacade(RpcFacade): + """Serve :class:`SO101LeRobotEnv` over the RPC boundary. + + ``shutdown`` and ``healthz`` are handled by :class:`RpcFacade` and the + transport; this only routes ``env.*`` methods to the env instance. + """ + + def __init__(self, env: SO101LeRobotEnv): + super().__init__() + self._env = env + + def _dispatch(self, method: str, args: tuple, kwargs: dict): + if method.startswith("env."): + return getattr(self._env, method[len("env."):])(*args, **kwargs) + raise ValueError(f"unknown RPC method: {method!r}") + + def request_shutdown(self) -> None: + """Signal the serve loop to exit (used by the signal handlers).""" + self._shutdown_event.set() + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def _build_argparser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description="Standalone LeRobot SO101 env server") + p.add_argument("--serial-port", default="/dev/ttyACM0", + help="Serial port of the SO101 follower arm.") + p.add_argument("--calibration-id", default="my_awesome_follower_arm", + help="LeRobot calibration id (loads .json).") + p.add_argument("--arm-camera-path", default="/dev/v4l/by-id/usb-icSpring_icspring_camera-video-index0", + help="OpenCV device path for the arm/hand camera " + "(empty string to disable).") + p.add_argument("--scene-camera-serial", default="409122274720", + help="Intel RealSense serial/name for the scene camera " + "(empty string to disable).") + p.add_argument("--camera-width", type=int, default=640) + p.add_argument("--camera-height", type=int, default=480) + p.add_argument("--camera-fps", type=int, default=30) + p.add_argument("--scene-camera-width", type=int, default=1280, + help="Scene (RealSense) color/depth width; default 1280 " + "(720p) for finer localization. <=0 falls back to " + "--camera-width.") + p.add_argument("--scene-camera-height", type=int, default=720, + help="Scene (RealSense) color/depth height; default 720. " + "<=0 falls back to --camera-height.") + p.add_argument("--no-cameras", action="store_true", + help="Disable all cameras (state-only observations).") + p.add_argument("--max-relative-target", type=float, default=None, + help="Per-step joint movement cap in degrees (safety). " + "Default: no cap (matches RLinf SO101Env).") + p.add_argument("--max-episode-steps", type=int, default=200) + p.add_argument("--step-frequency", type=float, default=30.0) + p.add_argument("--max-joint-vel", type=float, default=_MAX_JOINT_VEL_DEG_S, + help="Max joint speed (deg/s) for paced point-to-point moves " + "(reset/move_to/move_joints_delta). Lower = slower/safer.") + p.add_argument("--motor-acceleration", type=int, default=_MOTOR_ACCELERATION, + help="Feetech servo Acceleration register (0-254; LeRobot " + "default 254). Lower = gentler ramps. Set 254 for the " + "original snappy motion.") + p.add_argument("--position-gain", type=int, default=_POSITION_GAIN, + help="Feetech servo P_Coefficient / position gain (LeRobot " + "uses 16, factory default 32). Higher = stiffer, holds " + "commanded joints under load (fixes gripper landing " + "short/low); too high can buzz/oscillate. Applied to arm " + "joints only.") + p.add_argument("--auto-calibrate", action="store_true", + help="Allow LeRobot to run interactive calibration if the " + "arm is uncalibrated (may block on stdin). Off by default.") + p.add_argument("--urdf-path", default=None, + help="SO101 URDF for FK / EE pose. Default: " + "~/.cache/huggingface/lerobot/urdf/so101.urdf") + p.add_argument("--output-dir", required=True) + p.add_argument("--transport", choices=["socket", "http"], default="http", + help="RPC transport (default http). Socket (pickle) is also " + "available for the numpy obs/action payloads.") + p.add_argument("--host", type=str, default="127.0.0.1") + p.add_argument("--port", type=int, default=0, + help="RPC port. 0 asks the OS for a free port.") + return p + + +def _build_arm_camera_cfgs(args: argparse.Namespace) -> dict[str, dict]: + """Assemble the LeRobot (arm/hand) camera spec from CLI args. + + The scene camera is NOT included here — it is managed directly via + pyrealsense2 (see :class:`SceneCameraD405`) so we get depth + intrinsics. + """ + if args.no_cameras or not args.arm_camera_path: + return {} + return { + "arm": { + "type": "opencv", + "index_or_path": args.arm_camera_path, + "width": args.camera_width, + "height": args.camera_height, + "fps": args.camera_fps, + } + } + + +def main() -> int: + args = _build_argparser().parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + init_output_dir(args.output_dir) + + arm_camera_cfgs = _build_arm_camera_cfgs(args) + scene_serial = None if args.no_cameras else (args.scene_camera_serial or None) + logger.info( + "starting SO101 env server: serial_port=%s rpc_port=%s output_dir=%s arm_cams=%s scene=%s", + args.serial_port, args.port, args.output_dir, list(arm_camera_cfgs), scene_serial, + ) + + env = SO101LeRobotEnv( + port=args.serial_port, + calibration_id=args.calibration_id, + arm_camera_cfgs=arm_camera_cfgs, + scene_serial=scene_serial, + scene_size=( + args.scene_camera_height if args.scene_camera_height > 0 else args.camera_height, + args.scene_camera_width if args.scene_camera_width > 0 else args.camera_width, + ), + scene_fps=args.camera_fps, + urdf_path=args.urdf_path, + max_relative_target=args.max_relative_target, + max_episode_steps=args.max_episode_steps, + step_frequency=args.step_frequency, + motor_acceleration=args.motor_acceleration, + max_joint_vel_deg_s=args.max_joint_vel, + position_gain=args.position_gain, + auto_calibrate=args.auto_calibrate, + ) + + facade = LerobotEnvFacade(env) + + # Park the arm on SIGTERM / SIGINT (launcher kill, Ctrl-C). The handler + # only flags shutdown; the actual parking runs in ``env.close()`` in the + # ``finally`` below, never inside the signal context. (SIGKILL / kill -9 + # cannot be caught.) RpcFacade also exits on parent death (stdin EOF) and + # on the ``shutdown`` RPC. + def _handle_signal(signum, _frame): + logger.warning( + "received %s; parking arm and shutting down", + signal.Signals(signum).name, + ) + facade.request_shutdown() + + for _sig in (signal.SIGTERM, signal.SIGINT): + signal.signal(_sig, _handle_signal) + + try: + facade.serve(transport=args.transport, host=args.host, port=args.port) + finally: + env.close() + logger.info("driver exited cleanly") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/robots/lerobot/geometry.py b/robots/lerobot/geometry.py new file mode 100644 index 00000000..198c8df3 --- /dev/null +++ b/robots/lerobot/geometry.py @@ -0,0 +1,350 @@ +"""Pure-numpy geometry helpers for SO101 scene-camera localization. + +No hardware / lerobot imports — safe to unit-test offline. All transforms use +the convention ``T_a_b`` = pose of frame ``b`` in frame ``a`` so that +``p_a = T_a_b @ [p_b; 1]``. The world frame is the arm ``base_link``. +""" +from __future__ import annotations + +import numpy as np + + +def backproject_pixel(K, col: float, row: float, depth_m: float) -> np.ndarray: + """Backproject a pixel + metric depth into the camera frame (meters). + + Args: + K: 3x3 pinhole intrinsics (of the stream the pixel was taken from; + for our driver, depth is aligned to color so the color ``K`` + applies to both). + col: pixel x (column, u). + row: pixel y (row, v). + depth_m: metric depth at ``(row, col)`` in meters. + + Returns: + ``(3,)`` point ``[x, y, z]`` in the camera frame. + """ + K = np.asarray(K, dtype=np.float64) + fx, fy = K[0, 0], K[1, 1] + cx, cy = K[0, 2], K[1, 2] + z = float(depth_m) + x = (float(col) - cx) * z / fx + y = (float(row) - cy) * z / fy + return np.array([x, y, z], dtype=np.float64) + + +def transform_points(T, pts) -> np.ndarray: + """Apply a 4x4 homogeneous transform to a point or array of points. + + Args: + T: 4x4 transform. + pts: ``(3,)`` or ``(N, 3)`` points. + + Returns: + Transformed points, same leading shape as ``pts``. + """ + T = np.asarray(T, dtype=np.float64) + pts = np.asarray(pts, dtype=np.float64) + single = pts.ndim == 1 + p = np.atleast_2d(pts) + ph = np.concatenate([p, np.ones((p.shape[0], 1))], axis=1) # (N, 4) + out = (ph @ T.T)[:, :3] + return out[0] if single else out + + +def invert_transform(T) -> np.ndarray: + """Invert a 4x4 rigid transform (R, t) -> (R^T, -R^T t).""" + T = np.asarray(T, dtype=np.float64) + R = T[:3, :3] + t = T[:3, 3] + out = np.eye(4) + out[:3, :3] = R.T + out[:3, 3] = -R.T @ t + return out + + +def kabsch_umeyama(src, dst) -> tuple[np.ndarray, float]: + """Best-fit rigid transform mapping ``src`` -> ``dst`` (no scaling). + + Solves for ``T`` minimizing ``sum_i || T @ src_i - dst_i ||^2`` using the + SVD (Kabsch/Umeyama) with a reflection guard. + + Args: + src: ``(N, 3)`` source points (e.g. camera-frame). + dst: ``(N, 3)`` destination points (e.g. base-frame). + + Returns: + ``(T_4x4, rmse_meters)``. + """ + src = np.asarray(src, dtype=np.float64) + dst = np.asarray(dst, dtype=np.float64) + if src.shape != dst.shape or src.ndim != 2 or src.shape[1] != 3: + raise ValueError("src and dst must both be (N, 3) with matching N") + if src.shape[0] < 3: + raise ValueError("need at least 3 correspondences (4+ recommended)") + + c_src = src.mean(axis=0) + c_dst = dst.mean(axis=0) + s = src - c_src + d = dst - c_dst + + H = s.T @ d + U, _, Vt = np.linalg.svd(H) + # Reflection guard: ensure a proper rotation (det = +1). + D = np.eye(3) + D[2, 2] = np.sign(np.linalg.det(Vt.T @ U.T)) + R = Vt.T @ D @ U.T + t = c_dst - R @ c_src + + T = np.eye(4) + T[:3, :3] = R + T[:3, 3] = t + + resid = transform_points(T, src) - dst + rmse = float(np.sqrt((resid ** 2).sum(axis=1).mean())) + return T, rmse + + +def ransac_kabsch( + src, + dst, + *, + thresh_m: float = 0.015, + iters: int = 300, + min_inliers: int = 4, + seed: int = 0, +) -> tuple[np.ndarray, float, np.ndarray]: + """Robust rigid fit ``src`` -> ``dst`` with RANSAC over :func:`kabsch_umeyama`. + + Drops outlier correspondences (e.g. a mis-detected tip). Samples 4 points, + fits, counts inliers within ``thresh_m``, keeps the best consensus, then + refits on all inliers. + + Returns ``(T_4x4, inlier_rmse_m, inlier_mask)``. Falls back to a plain fit + on all points if no good consensus is found. + """ + src = np.asarray(src, dtype=np.float64) + dst = np.asarray(dst, dtype=np.float64) + n = src.shape[0] + if n < 4: + T, rmse = kabsch_umeyama(src, dst) + return T, rmse, np.ones(n, dtype=bool) + + rng = np.random.default_rng(seed) + idx = np.arange(n) + best_mask = None + best_count = 0 + for _ in range(iters): + sample = rng.choice(idx, size=4, replace=False) + try: + T, _ = kabsch_umeyama(src[sample], dst[sample]) + except Exception: + continue + resid = np.linalg.norm(transform_points(T, src) - dst, axis=1) + mask = resid < thresh_m + count = int(mask.sum()) + if count > best_count: + best_count = count + best_mask = mask + + if best_mask is None or best_count < min_inliers: + T, rmse = kabsch_umeyama(src, dst) + return T, rmse, np.ones(n, dtype=bool) + + T, rmse = kabsch_umeyama(src[best_mask], dst[best_mask]) + return T, rmse, best_mask + + +def solve_extrinsic_with_offset( + cam_pts, + tip_origins, + tip_rotations, + *, + thresh_m: float = 0.015, + offset_iters: int = 15, + outer_iters: int = 3, + seed: int = 0, +) -> tuple[np.ndarray, float, np.ndarray, np.ndarray]: + """Joint fit of ``T_base_cam`` AND a constant gripper-local tip offset. + + The markerless routine detects the moving jaw's motion-blob centroid in the + camera, but matches it against the FK tip-frame ORIGIN in the base frame. + Those are not the same physical point: the centroid is displaced from the + tip-frame origin by an (approximately) constant vector ``d`` expressed in + the gripper LOCAL frame. In a free-orientation grid that displacement points + a different way in the base frame at every pose, so a single rigid + ``T_base_cam`` cannot absorb it -- it lands in the residual and is the main + driver of the ~1 cm fit RMSE. + + This models the displacement explicitly. For pose ``i`` with tip-frame + origin ``o_i`` and rotation ``R_i`` (base frame, from FK) and detected + camera point ``c_i``:: + + T_base_cam @ c_i == o_i + R_i @ d + + It alternates two convex steps: (1) with ``d`` fixed, Kabsch-fit ``T`` to the + corrected targets ``b_i = o_i + R_i @ d``; (2) with ``T`` fixed, solve the + linear least squares ``R_i @ d = T @ c_i - o_i`` (closed form + ``d = mean_i R_i^T (T c_i - o_i)`` since each ``R_i`` is orthonormal). + Gross outliers (bad detections) are rejected with RANSAC. ``d`` is only + identifiable when the tip ROTATIONS vary across poses; with (near-)constant + orientation the displacement is indistinguishable from ``T``'s translation + and the solver returns ``d ~= 0`` (reducing to the plain rigid fit). + + Args: + cam_pts: ``(N, 3)`` detected points in the camera frame. + tip_origins: ``(N, 3)`` FK tip-frame origins in the base frame. + tip_rotations: ``(N, 3, 3)`` FK tip-frame rotations in the base frame. + thresh_m: RANSAC inlier threshold (m) on the offset-corrected fit. + offset_iters: max inner alternations per outer round. + outer_iters: rounds of (refine offset -> re-select inliers). + seed: RANSAC RNG seed. + + Returns: + ``(T_4x4, inlier_rmse_m, inlier_mask, d_local)``. + """ + cam = np.asarray(cam_pts, dtype=np.float64) + o = np.asarray(tip_origins, dtype=np.float64) + Rs = np.asarray(tip_rotations, dtype=np.float64) + n = cam.shape[0] + if n < 4 or Rs.shape != (n, 3, 3): + T, rmse = kabsch_umeyama(cam, o) + return T, rmse, np.ones(n, dtype=bool), np.zeros(3) + + # Stage 1: a loose RANSAC (tolerating the still-unknown offset) drops gross + # outliers -- mis-detected tips -- before we estimate the offset. + loose = max(thresh_m, 0.07) + T, rmse, mask = ransac_kabsch(cam, o, thresh_m=loose, seed=seed) + if int(mask.sum()) < 4: + mask = np.ones(n, dtype=bool) + + d = np.zeros(3) + for _ in range(outer_iters): + # Alternate: solve the offset from residuals, refit T to the corrected + # targets, until the offset stops moving. + for _ in range(offset_iters): + r = transform_points(T, cam[mask]) - o[mask] # (M, 3) + # mean_i R_i^T r_i (R_i^T r_i via 'mba,mb->ma' since R_i is (i,j)=(row,col)) + d_new = np.einsum("mba,mb->ma", Rs[mask], r).mean(axis=0) + b = o + np.einsum("nij,j->ni", Rs, d_new) # o_i + R_i d + T, _ = kabsch_umeyama(cam[mask], b[mask]) + if np.linalg.norm(d_new - d) < 1e-6: + d = d_new + break + d = d_new + # Re-select inliers with the tight threshold on the corrected targets. + b = o + np.einsum("nij,j->ni", Rs, d) + T, rmse, mask = ransac_kabsch(cam, b, thresh_m=thresh_m, seed=seed) + if int(mask.sum()) < 4: + mask = np.ones(n, dtype=bool) + T, rmse = kabsch_umeyama(cam, b) + return T, rmse, mask, d + + +def rotation_to_quat(R) -> np.ndarray: + """Convert a 3x3 rotation matrix to a quaternion ``[w, x, y, z]``.""" + R = np.asarray(R, dtype=np.float64) + tr = np.trace(R) + if tr > 0: + s = np.sqrt(tr + 1.0) * 2 + w = 0.25 * s + x = (R[2, 1] - R[1, 2]) / s + y = (R[0, 2] - R[2, 0]) / s + z = (R[1, 0] - R[0, 1]) / s + elif R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]: + s = np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) * 2 + w = (R[2, 1] - R[1, 2]) / s + x = 0.25 * s + y = (R[0, 1] + R[1, 0]) / s + z = (R[0, 2] + R[2, 0]) / s + elif R[1, 1] > R[2, 2]: + s = np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) * 2 + w = (R[0, 2] - R[2, 0]) / s + x = (R[0, 1] + R[1, 0]) / s + y = 0.25 * s + z = (R[1, 2] + R[2, 1]) / s + else: + s = np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) * 2 + w = (R[1, 0] - R[0, 1]) / s + x = (R[0, 2] + R[2, 0]) / s + y = (R[1, 2] + R[2, 1]) / s + z = 0.25 * s + return np.array([w, x, y, z], dtype=np.float64) + + +def sample_depth_patch(depth_m, col: int, row: int, radius: int = 2) -> float: + """Median of the valid (>0, finite) depths in a small patch around a pixel. + + Robustifies a single-pixel depth read (sensor noise / dropouts). Returns + ``nan`` if no valid depth is found in the patch. + """ + depth_m = np.asarray(depth_m, dtype=np.float64) + h, w = depth_m.shape[:2] + r0, r1 = max(0, row - radius), min(h, row + radius + 1) + c0, c1 = max(0, col - radius), min(w, col + radius + 1) + patch = depth_m[r0:r1, c0:c1].reshape(-1) + valid = patch[np.isfinite(patch) & (patch > 0)] + if valid.size == 0: + return float("nan") + return float(np.median(valid)) + + +def detect_tip_pixel_by_motion( + rgb_open, + rgb_closed, + depth_m, + K, + *, + diff_thresh: int = 18, + min_area: int = 40, + max_area: int = 40000, +) -> dict | None: + """Locate the gripper in the scene image via gripper-toggle motion. + + Given two scene frames that differ only by the gripper opening (arm held + still), the changed pixels are the gripper fingers. Returns the centroid of + the largest valid motion blob, the median depth over it, and the + backprojected camera-frame point — or ``None`` if no usable blob is found. + + ``cv2`` is imported lazily so the rest of this module stays import-light. + """ + import cv2 + + a = cv2.cvtColor(np.asarray(rgb_open, dtype=np.uint8), cv2.COLOR_RGB2GRAY) + b = cv2.cvtColor(np.asarray(rgb_closed, dtype=np.uint8), cv2.COLOR_RGB2GRAY) + diff = cv2.GaussianBlur(cv2.absdiff(a, b), (5, 5), 0) + _, mask = cv2.threshold(diff, int(diff_thresh), 255, cv2.THRESH_BINARY) + kernel = np.ones((3, 3), np.uint8) + mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) + mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8)) + + num, labels, stats, centroids = cv2.connectedComponentsWithStats(mask, connectivity=8) + if num <= 1: + return None + + depth_m = np.asarray(depth_m, dtype=np.float64) + # Largest component first (skip background label 0). + for comp in np.argsort(stats[1:, cv2.CC_STAT_AREA])[::-1] + 1: + area = int(stats[comp, cv2.CC_STAT_AREA]) + if area < min_area or area > max_area: + continue + col, row = float(centroids[comp][0]), float(centroids[comp][1]) + # Depth AT the centroid (small patch), so the returned (col, row, z) all + # describe the SAME point. The blob spans a depth gradient under the + # oblique scene view, so a whole-blob median would pair the centroid + # pixel with some other pixel's depth and bias the back-projection. + z = sample_depth_patch(depth_m, int(round(col)), int(round(row)), radius=3) + if not np.isfinite(z): + # Centroid fell on a depth dropout -- fall back to the blob median. + dvals = depth_m[labels == comp] + dvals = dvals[np.isfinite(dvals) & (dvals > 0)] + if dvals.size == 0: + continue + z = float(np.median(dvals)) + p_cam = backproject_pixel(K, col, row, z) + return { + "pixel": [row, col], + "depth_m": z, + "area": area, + "xyz_cam": p_cam.tolist(), + } + return None diff --git a/robots/lerobot/kinematics.py b/robots/lerobot/kinematics.py new file mode 100644 index 00000000..588a46ca --- /dev/null +++ b/robots/lerobot/kinematics.py @@ -0,0 +1,100 @@ +"""SO101 forward / inverse kinematics in the arm base frame. + +Thin wrapper over LeRobot's placo-based +:class:`lerobot.model.kinematics.RobotKinematics`, pinned to the SO101 arm +joints and the ``gripper_frame_link`` tip frame. Forward kinematics returns +``T_base_gripper`` (the gripper pose in the ``base_link`` world frame); inverse +kinematics maps a desired base-frame gripper pose back to joint targets. +""" +from __future__ import annotations + +import os + +import numpy as np + +# Arm joints in bus-ID order (no gripper); must match the URDF joint names. +_ARM_JOINTS = ( + "shoulder_pan", + "shoulder_lift", + "elbow_flex", + "wrist_flex", + "wrist_roll", +) +_DEFAULT_URDF = "~/.cache/huggingface/lerobot/urdf/so101.urdf" +_TARGET_FRAME = "gripper_frame_link" + + +class SO101Kinematics: + """FK/IK for the SO101 arm, expressed in the ``base_link`` world frame.""" + + def __init__(self, urdf_path: str | None = None, target_frame: str = _TARGET_FRAME): + from lerobot.model.kinematics import RobotKinematics + + urdf = os.path.expanduser(urdf_path or _DEFAULT_URDF) + if not os.path.isfile(urdf): + raise FileNotFoundError( + f"SO101 URDF not found at {urdf}. Set urdf_path or download the " + "SO101 URDF (see toolkits/lerobot/compute_ee_pose.py --urdf help)." + ) + self._kin = RobotKinematics( + urdf_path=urdf, + target_frame_name=target_frame, + joint_names=list(_ARM_JOINTS), + ) + + def fk(self, joints_deg) -> np.ndarray: + """Forward kinematics: 5 arm joint angles (deg) -> ``T_base_gripper`` (4x4).""" + q = np.asarray(joints_deg, dtype=np.float64).reshape(-1)[: len(_ARM_JOINTS)] + return np.asarray(self._kin.forward_kinematics(q), dtype=np.float64) + + def ik( + self, + current_joints_deg, + T_base_gripper_des, + *, + position_weight: float = 1.0, + orientation_weight: float = 0.01, + max_iters: int = 60, + pos_tol_m: float = 0.002, + orient_tol_deg: float = 2.0, + ) -> np.ndarray: + """Inverse kinematics: desired ``T_base_gripper`` (4x4) -> arm joints (deg). + + placo's solver advances one step per call, so a single solve undershoots + on large targets. We iterate (feeding each solution back as the seed) + until the FK error is within tolerance or ``max_iters`` is hit; the + result is the best effort (check the FK error if it matters). + + With ``orientation_weight <= 0`` only the position error (``pos_tol_m``) + gates convergence and the target rotation is ignored. With a positive + ``orientation_weight`` the solve also matches the rotation of + ``T_base_gripper_des``, and convergence additionally requires the + orientation error below ``orient_tol_deg`` (used by ``move_to``'s + top-down mode). + + ``current_joints_deg`` seeds the first iteration (its gripper entry, if + any, is preserved by LeRobot's solver). + """ + q = np.asarray(current_joints_deg, dtype=np.float64).reshape(-1) + T = np.asarray(T_base_gripper_des, dtype=np.float64) + target = T[:3, 3] + R_des = T[:3, :3] + n = len(_ARM_JOINTS) + for _ in range(max_iters): + q = np.asarray( + self._kin.inverse_kinematics( + q, T, + position_weight=position_weight, + orientation_weight=orientation_weight, + ), + dtype=np.float64, + ) + T_cur = self.fk(q[:n]) + if np.linalg.norm(T_cur[:3, 3] - target) >= pos_tol_m: + continue + if orientation_weight <= 0.0: + break + cos_ang = (np.trace(R_des.T @ T_cur[:3, :3]) - 1.0) / 2.0 + if np.degrees(np.arccos(np.clip(cos_ang, -1.0, 1.0))) < orient_tol_deg: + break + return q diff --git a/robots/lerobot/prompt.py b/robots/lerobot/prompt.py new file mode 100644 index 00000000..320ae853 --- /dev/null +++ b/robots/lerobot/prompt.py @@ -0,0 +1,159 @@ +"""LeRobot SO101 prompt fragments and assembly.""" +from __future__ import annotations + +from rpent.context.prompt_utils import BulletList, Numbered, PromptNode +from rpent.context.prompts import prompt as base_prompt + +# --- system-prompt sections ------------------------------------------------ + +PREAMBLE = """ +You are a physical agent that drives a robot to accomplish a manipulation task. You act by calling tools: observe the scene through cameras and robot state, reason about where things are in the robot's coordinate frame, and command the arm. (Under the Claude Code / Codex CLIs the tools may appear namespaced as ``mcp__rpent__`` — call them by whatever name your tool list shows.) +""" + +GOAL = """ +Accomplish the task specified by the user. +""" + +ENVIRONMENT = BulletList([ + """ + Robot: SO101 — a 5-DOF arm plus a 1-DOF gripper. Command it two ways: + move_to (drive to a world-frame xyz) and move_joints_delta (relative + per-joint nudge in degrees; negative gripper_delta closes the gripper). + move_to positions the point BETWEEN THE FINGERTIPS (not the wrist) at the + target, with the ~7 cm fingers hanging below it. move_joints_delta is for fine alignment. + """, + """ + World frame = the arm base (``base_link``), in meters: x forward, z up, y + lateral. back_project, get_ee_pose, and move_to all use this one frame — + call get_ee_pose to ground yourself. + """, + """ + Reachable box (move_to clips to it and flags clipped_to_workspace): x + [0.08, 0.38], y [-0.28, 0.28], z [-0.055, 0.30] m. The table/plate surface + is near z = -0.06; the z floor stops the fingertips just above it, so you + can descend to the floor without hitting the table. + """, + """ + Gripper opening is in degrees: ~90 open, ~10-20 grasping. NEVER command 0 — + it stalls the motor against its stop. + """, + """ + Scene camera: fixed, HAS depth — back_project a pixel to a world xyz. Use it + to locate the target and to check from the side that the fingers are placed + correctly around it. Arm camera: on the gripper, looking straight down, NO + depth — so never back_project it. + """, + """ + Tools: view_env_state (state + scene/arm images), get_scene_camera_meta + (intrinsics + calibration flag), back_project (scene pixel -> world xyz), + get_ee_pose (the fingertip point's xyz in world), move_to, move_joints_delta, + finish. Plus read_text_file / write_text_file / list_dir for the scratch + dir ({{output_dir}}), and read_memory / write_memory for lessons carried + across runs. + """, +]) + +RULES = BulletList([ + """ + Observe before acting — never guess coordinates. Locate objects with + back_project on the scene image. + """, + """ + back_project only yields world coordinates when the scene camera is + calibrated (check get_scene_camera_meta); if ``calibrated`` is false its + output is camera-frame and unusable — stop and report. + """, + """ + After each action, check ``reached``, ``pos_error_m``, and (for + approach="down") ``approach_tilt_deg`` (~0 = vertical). If ``reached`` is + false with a reach / near-singular note, don't repeat the same command — figure out what's wrong from the images, try a different xyz or yaw_deg, or move_joints_delta. + """, + """ + The scene image is your ground truth: study it to confirm each result (the + view updates after every move_to; call view_env_state only when you need + another look). Verify each sub-goal before building on it — that you are + positioned correctly before committing an action, and that the action + succeeded before the next one; if a precondition isn't met, re-localize or + re-position instead of forcing it. + """, + """ + When grasping, close the gripper with move_joints_delta (negative + gripper_delta), which holds the arm still — never close with move_to (it + re-solves IK and can shift off the target). + """, + """ + Trust images more than coordinates. For instance, move_to might not move the center of the gripper to the coordinates you specify --- the jaws are in-symmetrical, so the center of the jaws might be offset from the center of the gripper. You may want to either find out the offset and pad your xyz accordingly, or use move_joints_delta to nudge the gripper into alignment after move_to. + """, + """ + If the env server returns an error, stop and report it — don't continue + blindly. + """, +]) + +WORKFLOW = Numbered([ + """ + Consult memory first: call read_memory (no arguments) to read the MEMORY.md + index of lessons learned on past runs, then read_memory(name) any entry + relevant to this task or robot and apply it. This is fast and often saves + you from repeating a known failure. + """, + """ + Understand the task: from the user's instruction, identify the target + object or site and what will count as success. + """, + """ + Localize before acting: on the (calibrated) scene image, back_project a few + pixels on the target to get a stable world point P — never guess coordinates. + """, + """ + Plan the approach this task needs (grasp, press, push, place, ...). As a + rule, go to a safe standoff above or beside the target first, then move in + along the task axis — approaching at object level tends to sweep things + aside. Use approach="down" for vertical actions. + """, + """ + Act, then verify: after each action check its result and the scene image, + confirm the sub-goal before the next, and re-localize if anything moved. + """, + """ + If the task needs a grasp: the move_to point sits between the fingertips, so + put them around the object (target z at or below its top), confirm from the + scene image / get_ee_pose (not the arm cam) that the object is between the + jaws, then close with move_joints_delta and lift. + """, + """ + Record a lesson: if this run taught you something non-obvious and verified — + a fix that took more than one try, a useful magic number/offset, or a gotcha + about this robot or scene — call write_memory(name, hook, content) so future + runs reuse it. Skip it for routine runs and never record guesses. If a past + memory proved wrong, call write_memory with the same name to correct it. + """, + """ + Finish: verify the success condition, report success or failure with a short + summary, and call finish(). + """, +]) + +# --- user-prompt sections -------------------------------------------------- + +USER_CONTEXT = """Pick up the green cube on the table.""" + + +# --- prompt tree factories ------------------------------------------------- + + +def system_prompt() -> dict[str, PromptNode]: + """Return the system prompt tree.""" + return { + "Intro": PREAMBLE, + "Goal": GOAL, + "Rules": RULES, + "Workflow": WORKFLOW, + "Environment": ENVIRONMENT, + "Output": base_prompt.OUTPUT, + } + + +def user_prompt() -> dict[str, PromptNode]: + """Return the first user message tree.""" + return USER_CONTEXT diff --git a/robots/lerobot/scene_camera.py b/robots/lerobot/scene_camera.py new file mode 100644 index 00000000..3f93ac74 --- /dev/null +++ b/robots/lerobot/scene_camera.py @@ -0,0 +1,108 @@ +"""Direct pyrealsense2 scene camera for the SO101 env (depth aligned to color). + +LeRobot's RealSense wrapper does not align depth to color and does not expose +intrinsics, both of which we need for pixel -> 3D backprojection. So the scene +camera is managed here directly via ``pyrealsense2``: a single pipeline streams +color + depth, ``rs.align`` registers depth into the color frame, and the color +intrinsics + depth scale are read from the active profile. + +The arm (hand) camera stays under LeRobot (color only); only the fixed scene +camera needs depth. +""" +from __future__ import annotations + +import numpy as np + + +class SceneCameraD405: + """Color + depth (aligned to color) from an Intel RealSense (e.g. D405).""" + + def __init__( + self, + serial: str, + *, + width: int = 640, + height: int = 480, + fps: int = 30, + warmup_frames: int = 15, + ) -> None: + import pyrealsense2 as rs + + self._rs = rs + self._serial = str(serial) + self._width = int(width) + self._height = int(height) + + self._pipeline = rs.pipeline() + cfg = rs.config() + cfg.enable_device(self._serial) + cfg.enable_stream(rs.stream.color, self._width, self._height, rs.format.rgb8, int(fps)) + cfg.enable_stream(rs.stream.depth, self._width, self._height, rs.format.z16, int(fps)) + self._profile = self._pipeline.start(cfg) + + # Align depth into the color frame so depth[row, col] matches the + # color pixel (row, col). + self._align = rs.align(rs.stream.color) + + depth_sensor = self._profile.get_device().first_depth_sensor() + self._depth_scale = float(depth_sensor.get_depth_scale()) # meters / unit + + color_stream = self._profile.get_stream(rs.stream.color).as_video_stream_profile() + intr = color_stream.get_intrinsics() + self._K = np.array( + [[intr.fx, 0.0, intr.ppx], [0.0, intr.fy, intr.ppy], [0.0, 0.0, 1.0]], + dtype=np.float64, + ) + + for _ in range(max(0, warmup_frames)): + self._pipeline.wait_for_frames() + + # -- capture ----------------------------------------------------------- + + def read(self) -> tuple[np.ndarray, np.ndarray]: + """Return ``(color_rgb_uint8 [H,W,3], depth_m_float32 [H,W])``. + + Depth is metric (meters), aligned to the color frame; invalid/no-return + pixels are ``0.0``. + """ + frames = self._pipeline.wait_for_frames() + frames = self._align.process(frames) + color = frames.get_color_frame() + depth = frames.get_depth_frame() + if not color or not depth: + raise RuntimeError("scene camera: incomplete frameset (no color/depth)") + color_img = np.ascontiguousarray(np.asanyarray(color.get_data()), dtype=np.uint8) + depth_raw = np.asanyarray(depth.get_data()) # uint16, depth units + depth_m = (depth_raw.astype(np.float32) * self._depth_scale) + return color_img, np.ascontiguousarray(depth_m) + + # -- metadata ---------------------------------------------------------- + + @property + def K(self) -> np.ndarray: + return self._K + + @property + def depth_scale(self) -> float: + return self._depth_scale + + @property + def size(self) -> tuple[int, int]: + """``(height, width)``.""" + return (self._height, self._width) + + def meta(self) -> dict: + """JSON-able camera metadata (intrinsics, size, depth scale, serial).""" + return { + "serial": self._serial, + "K": self._K.tolist(), + "width": self._width, + "height": self._height, + "depth_scale": self._depth_scale, + } + + def close(self) -> None: + try: + self._pipeline.stop() + except Exception: + pass diff --git a/robots/lerobot/toolkit.py b/robots/lerobot/toolkit.py new file mode 100644 index 00000000..30024f66 --- /dev/null +++ b/robots/lerobot/toolkit.py @@ -0,0 +1,95 @@ +"""LeRobot SO101 toolkit: common tools + SO101 primitives. + +Inherits the common file/IO tools (including ``finish``) from :class:`Toolkit` +and registers the SO101-specific tools (``view_env_state``, ``back_project``, +driver readers, and the move primitives) on top. +""" +from __future__ import annotations + +from functools import partial +from typing import Any + +from robots.lerobot import tools as lerobot_tools +from rpent.dashboard.events import DashboardEventSink +from rpent.tools.state import EnvState +from rpent.tools.toolkit import Toolkit +from rpent.utils.logging import get_output_dir + + +class LerobotToolkit(Toolkit): + """Toolkit for the LeRobot SO101 environment.""" + + _VIEW_IMAGE_SLOTS = { + "_image_bytes": "scene.png", + "_image_cam_bytes": "arm.png", + } + # Per-env artifact names for the dashboard's live frame images. + _FRAME_ARTIFACTS = { + "camera": "scene.png", + "wrist": "arm.png", + } + + def __init__( + self, + *, + env: Any, + model: Any | None = None, + dashboard_events: DashboardEventSink, + ) -> None: + state = EnvState(get_output_dir()) + super().__init__(dashboard_events=dashboard_events, state=state) + self.init_driver_clean(env=env, model=model) + self._register_tools() + + # ------------------------------------------------------------------ + # Registration + # ------------------------------------------------------------------ + def _register_tools(self) -> None: + # Read-only tools whose handlers aren't driver methods (they need the + # run's EnvState bound in). Every other spec binds to its primitive- + # driver method; @updatestate on the method decides state capture. + state_handlers = { + "view_env_state": partial( + self._state.view, image_slots=self._VIEW_IMAGE_SLOTS + ), + "back_project": partial(lerobot_tools.back_project, state=self._state), + } + for spec in lerobot_tools.TOOLS_SPEC: + name = spec["name"] + if name in state_handlers: + handler = state_handlers[name] + else: + handler = getattr(self._driver, name, None) + if handler is None: + continue # spec without a backing driver method + self.add_tool(name, spec, handler) + + def get_env_state( + self, + *, + command: dict[str, Any], + result: dict[str, Any], + elapsed_s: float, + ) -> dict[str, Any]: + self._driver._refresh() + record = lerobot_tools.dump_state( + self._driver, + self._state, + log={"command": command, "result": result, "elapsed_s": elapsed_s}, + ) + out = self._state.view(record.step_idx, image_slots=self._VIEW_IMAGE_SLOTS) + out["agent_elapsed_s"] = elapsed_s + return out + + def init_driver_clean(self, *, env: Any, model: Any | None = None) -> None: + """Wipe stale run artifacts, build the primitive driver, dump step 0.""" + self._state.reset() + driver = lerobot_tools.LerobotPrimitives(env=env, model=model) + driver.reset() + record = lerobot_tools.dump_state(driver, self._state, log=None) + self._driver = driver + self._publish_step(record) + + def close(self) -> None: + """End-of-run cleanup hook. TODO: flush an episode video if desired.""" + return None diff --git a/robots/lerobot/tools.py b/robots/lerobot/tools.py new file mode 100644 index 00000000..e1b8294e --- /dev/null +++ b/robots/lerobot/tools.py @@ -0,0 +1,425 @@ +"""LeRobot SO101 tool implementation. + +Structure mirrors :mod:`robots.libero.tools`: + +* :class:`LerobotPrimitives` — the primitive driver the toolkit owns. Holds + the env client (and an optional policy/VLA model) plus per-run state, and + exposes one method per primitive tool. +* per-step state dump (:func:`dump_state`) plus reader tools bound to the + run's :class:`rpent.tools.state.EnvState`. +* :data:`TOOLS_SPEC` — Anthropic-shaped tool schemas. + +NOTE: this is a scaffold. The concrete robot primitives (move / grasp / +release / ...) and their schemas are intentionally left as TODOs — only the +loop infrastructure (state dump + view) is implemented so the env +loads and the pattern is in place. +""" +from __future__ import annotations + +from typing import Any + +import numpy as np + +from robots.lerobot.env_client import LerobotEnvClient +from rpent.tools.common import robust_surface_centroid +from rpent.tools.state import EnvState, StepRecord +from rpent.tools.toolkit import updatestate +from rpent.utils.logging import get_logger + +logger = get_logger("lerobot") + +_BACKPROJECT_RADIUS = 6 + + +def _to_list(x) -> list: + """Coerce a numpy array / sequence / scalar into a plain list[float].""" + if x is None: + return [] + arr = np.asarray(x, dtype=np.float32).reshape(-1) + return [round(float(v), 4) for v in arr] + + +# --------------------------------------------------------------------------- +# Primitive driver +# --------------------------------------------------------------------------- + + +class LerobotPrimitives: + """Wraps a single SO101 env client (+ optional policy) with primitive- + level methods. + + The toolkit constructs this from the env RPC client and calls + :meth:`reset` once at start-up; :func:`dump_state` reads back the latest + observation via :meth:`get_state` / :meth:`latest_frames` after each tool. + + TODO: add the concrete robot primitives (e.g. ``move_to``, ``grasp``, + ``release``, ``home``) on top of the low-level :meth:`step` passthrough. + Each should return a ``dict`` log and leave ``self._last_obs`` current. + """ + + def __init__(self, env: LerobotEnvClient, model: Any | None = None): + self.env = env + self.model = model # optional policy/VLA client; None for scripted prims + self._last_obs: dict | None = None + self._spec: dict | None = None + self._scene_meta: dict | None = None + self._num_steps = 0 + + # -- lifecycle --------------------------------------------------------- + + def reset(self) -> tuple[dict, Any]: + """Reset the env (arm → rest) and cache the first observation.""" + self._spec = self.env.get_spec() + obs, info = self.env.reset() + self._last_obs = obs + self._num_steps = 0 + return obs, info + + def step(self, action) -> dict: + """Low-level passthrough: one ``env.step``. Higher-level primitives + are layered on top of this (TODO). + """ + obs, reward, terminated, truncated, info = self.env.step(action) + self._last_obs = obs + self._num_steps += 1 + return { + "reward": float(reward), + "terminated": bool(terminated), + "truncated": bool(truncated), + "num_steps": self._num_steps, + } + + # -- state accessors used by dump_state -------------------------------- + + def get_state(self) -> dict: + """Return the robot proprioceptive state from the last observation.""" + if self._last_obs is None: + return {} + st = self._last_obs.get("state", {}) + out = { + "joint_position": _to_list(st.get("joint_position")), + "gripper_position": _to_list(st.get("gripper_position")), + "num_steps": self._num_steps, + } + if st.get("ee_pose_base") is not None: + out["ee_pose_base"] = _to_list(st.get("ee_pose_base")) + if st.get("ee_quat_base") is not None: + out["ee_quat_base"] = _to_list(st.get("ee_quat_base")) + return out + + def latest_frames(self) -> dict: + """Return the camera frames dict from the last observation.""" + if self._last_obs is None: + return {} + return dict(self._last_obs.get("frames", {})) + + def latest_depth(self) -> np.ndarray | None: + """Return the scene depth map (meters) from the last observation.""" + if self._last_obs is None: + return None + depth = self._last_obs.get("depth", {}) + scene = depth.get("scene") if isinstance(depth, dict) else None + return None if scene is None else np.asarray(scene, dtype=np.float32) + + # -- localization (base/world frame) ----------------------------------- + + def get_ee_pose(self) -> dict: + """Live FK: gripper pose in the base (world) frame.""" + return self.env.get_ee_pose() + + def get_scene_camera_meta(self) -> dict: + """Scene-camera intrinsics + depth scale + base extrinsic (cached).""" + if self._scene_meta is None: + self._scene_meta = self.env.get_scene_camera_meta() + return self._scene_meta + + # -- primitives (move the robot; toolkit capture refreshes observation) -- + + @updatestate + def move_to( + self, + xyz, + gripper: float | None = None, + approach: str = "free", + yaw_deg: float | None = None, + ) -> dict: + """Move the gripper to a world-frame (base_link) XYZ via driver IK. + + ``approach="down"`` keeps the gripper pointing straight down (for + grasping); ``yaw_deg`` sets the jaw heading. See the driver for details. + """ + return self.env.move_to( + xyz, gripper=gripper, approach=approach, yaw_deg=yaw_deg + ) + + @updatestate + def move_joints_delta( + self, + delta_deg, + gripper_delta: float | None = None, + ) -> dict: + """Nudge each arm joint relatively (degrees) for fine alignment.""" + return self.env.move_joints_delta( + delta_deg, + gripper_delta=gripper_delta, + ) + + def _refresh(self) -> None: + """Refresh the cached observation for toolkit state capture.""" + try: + self._last_obs = self.env.get_obs() + except Exception as e: + logger.warning("obs refresh failed: %s", e) + + +def dump_state( + driver: LerobotPrimitives, + state: EnvState, + log: dict | None = None, +) -> StepRecord: + """Dump camera artifacts and proprioceptive state through ``EnvState``.""" + log = log or {} + with state.record_step( + state=driver.get_state(), + command=log.get("command"), + result=log.get("result"), + elapsed_s=log.get("elapsed_s"), + ) as step_idx: + for camera, frame in driver.latest_frames().items(): + state.save( + f"{camera}.png", + frame, + step=step_idx, + ) + + depth = driver.latest_depth() + if depth is not None: + state.save( + "scene_depth.npy", + depth, + step=step_idx, + ) + + scene_meta = driver.get_scene_camera_meta() + if isinstance(scene_meta, dict) and "error" not in scene_meta: + state.save( + "scene_metadata.json", + scene_meta, + step=step_idx, + ) + + return state.get(step_idx) + + +def back_project( + row: int, + col: int, + step: int = -1, + radius: int | None = _BACKPROJECT_RADIUS, + *, + state: EnvState, +) -> dict: + """Backproject a scene-camera pixel neighborhood to a robust world point. + + The scene camera views the table at a steep oblique angle, so a single + pixel's depth error becomes a large lateral error. Instead of trusting one + pixel, this back-projects EVERY valid pixel in a ``(2*radius+1)`` square + window around ``(row, col)``, keeps those on the dominant surface (depth + within a narrow band of the window median, rejecting background, table, + and dropouts), and returns the MEDIAN world ``xyz`` of that surface: a + stable object centroid rather than one face pixel. Use ``radius=0`` for the + old single-pixel behavior. + + Pick ``(row, col)`` on the saved ``scene.png`` observation; depth is + aligned in ``scene_depth.npy``. Returns base/world ``xyz`` when calibrated, else + camera-frame ``xyz_cam`` with a note. + """ + try: + record = state.get(step) + except Exception as exc: + return {"error": f"state step not available: {exc}"} + nn = record.step_idx + metadata_name = "scene_metadata.json" + depth_name = "scene_depth.npy" + if metadata_name not in record.artifacts: + return {"error": f"scene camera metadata not recorded for step {nn}"} + try: + meta = state.load(metadata_name, step=nn) + if depth_name not in record.artifacts: + raise FileNotFoundError(depth_name) + depth = state.load(depth_name, step=nn) + except Exception as exc: + return {"error": f"depth for step {nn} not found: {exc}"} + + out = robust_surface_centroid( + depth, + meta["K"], + meta.get("T_base_cam"), + row, + col, + radius=_BACKPROJECT_RADIUS if radius is None else radius, + ) + if "error" in out: + return out + + out["step"] = nn + out["camera"] = "scene" + out["camera_frame"] = meta.get("frame", "scene_cam") + if meta.get("T_base_cam") is not None: + out["frame"] = "base_link" + else: + out["frame"] = meta.get("frame", "scene_cam") + out["note"] = ( + "scene camera not calibrated (no T_base_cam); returning camera-frame " + "xyz only. Run robots/lerobot/calibrate_scene_cam.py." + ) + return out + + +# --------------------------------------------------------------------------- +# Tool schema declarations (Anthropic-shaped) +# --------------------------------------------------------------------------- + +TOOLS_SPEC: list[dict[str, Any]] = [ + { + "name": "view_env_state", + "description": ( + "Read one recorded state and its observation artifacts. Step -1 " + "selects the latest entry. Embeds scene and arm camera frames." + ), + "input_schema": { + "type": "object", + "properties": { + "step": { + "type": "integer", + "default": -1, + "description": "Step number; 0 = initial, -1 = latest.", + }, + }, + }, + }, + { + "name": "get_ee_pose", + "description": ( + "Live forward kinematics: the gripper tip pose in the WORLD frame " + "(arm base_link). Returns xyz (meters), quat_wxyz, and joints_deg. " + "Use this to know where the gripper currently is in world coords." + ), + "input_schema": {"type": "object", "properties": {}}, + }, + { + "name": "get_scene_camera_meta", + "description": ( + "Scene-camera calibration: intrinsics K, depth scale, and whether " + "the camera->base extrinsic (T_base_cam) is calibrated. If " + "calibrated is false, back_project returns camera-frame coords only." + ), + "input_schema": {"type": "object", "properties": {}}, + }, + { + "name": "back_project", + "description": ( + "Backproject a SCENE-camera pixel to a 3D point in the WORLD frame " + "(arm base_link), using the saved aligned depth. Pick (row, col) on " + "the scene color image from view_env_state, near the CENTER of " + "the target. It samples a small window around the pixel and returns " + "the robust MEDIAN world `xyz` of the object surface (not one noisy " + "pixel), plus `n_points` and `xy_spread_m` (a small spread means a " + "confident estimate). Returns world `xyz` when calibrated (else " + "camera-frame `xyz_cam`). This is the primary tool for locating " + "objects in the robot's coordinate system." + ), + "input_schema": { + "type": "object", + "properties": { + "row": {"type": "integer", "description": "Pixel row (y) in the scene image, near the target center."}, + "col": {"type": "integer", "description": "Pixel column (x) in the scene image, near the target center."}, + "step": { + "type": "integer", + "default": -1, + "description": "Step whose depth to use; -1 = latest.", + }, + "radius": { + "type": ["integer", "null"], + "description": "Half-size (px) of the sampling window; null = default (6). Use a smaller value for tiny/cluttered targets, 0 for a single pixel.", + }, + }, + "required": ["row", "col"], + }, + }, + { + "name": "move_to", + "description": ( + "Move the gripper to a target [x, y, z] in the WORLD frame (arm " + "base_link), meters. The target is clipped to a safe workspace box " + "and approached in small capped steps. `approach` controls the " + "wrist orientation: 'free' (default) lets IK pick any orientation " + "(maximal reach, but the fingertips' exact location is " + "unpredictable -- not for grasping); 'down' keeps the gripper " + "pointing STRAIGHT DOWN so the fingers descend vertically (use this " + "to grasp). With approach='down', `yaw_deg` sets the jaw-line " + "heading about vertical (0=+x/forward, 90=+y/left); leave null to " + "auto-pick a reachable heading. Optionally set the gripper opening. " + "Returns `reached`, `pos_error_m`, and `approach_tilt_deg` (0 = " + "perfectly vertical). Use get_ee_pose / back_project to choose " + "targets in the same frame." + ), + "input_schema": { + "type": "object", + "properties": { + "xyz": { + "type": "array", + "description": "World-frame target [x, y, z] in meters (base_link).", + "items": {"type": "number"}, + "minItems": 3, + "maxItems": 3, + }, + "gripper": { + "type": ["number", "null"], + "description": "Gripper opening degrees (~90 open .. ~15 grasp); null keeps current. Never 0.", + }, + "approach": { + "type": "string", + "enum": ["free", "down"], + "description": "'down' = gripper points straight down (for grasping); 'free' = any orientation. Default 'free'.", + }, + "yaw_deg": { + "type": ["number", "null"], + "description": "With approach='down', jaw-line heading about vertical in degrees (0=forward, 90=left). Null = auto-pick a reachable heading.", + }, + }, + "required": ["xyz"], + }, + }, + { + "name": "move_joints_delta", + "description": ( + "Fine-adjust the arm by nudging each joint RELATIVELY (degrees). " + "`delta_deg` is 5 values added to the current joints: [shoulder_pan, " + "shoulder_lift, elbow_flex, wrist_flex, wrist_roll]. Each is capped " + "to +/-15 deg/call and clamped to joint limits. Positive wrist_roll " + "rotates the jaw line; wrist_flex tilts the gripper up/down. Use " + "this when move_to gets you close but the grasp needs a small tweak " + "(align the jaws across the object, or descend a few mm). Optionally " + "nudge the gripper with `gripper_delta`. Returns the new joints and " + "EE xyz. Prefer move_to for big moves; this is for fine alignment." + ), + "input_schema": { + "type": "object", + "properties": { + "delta_deg": { + "type": "array", + "description": "Relative joint deltas in degrees [pan, lift, elbow, wrist_flex, wrist_roll].", + "items": {"type": "number"}, + "minItems": 5, + "maxItems": 5, + }, + "gripper_delta": { + "type": ["number", "null"], + "description": "Relative gripper opening change in degrees; null keeps current.", + }, + }, + "required": ["delta_deg"], + }, + }, +] diff --git a/robots/libero/__init__.py b/robots/libero/__init__.py index 7bf9323e..a5fd060e 100644 --- a/robots/libero/__init__.py +++ b/robots/libero/__init__.py @@ -46,7 +46,6 @@ def get_toolkit( *, primitives_kwargs: dict[str, Any], dashboard_events: DashboardEventSink, - video_path: str | None = None, ): """Return the LIBERO toolkit (common tools + LIBERO primitives).""" from robots.libero.toolkit import LiberoToolkit @@ -54,7 +53,6 @@ def get_toolkit( return LiberoToolkit( primitives_kwargs=primitives_kwargs, dashboard_events=dashboard_events, - video_path=video_path, ) diff --git a/robots/libero/guides/env_calibration.md b/robots/libero/guides/env_calibration.md index 1dbab00d..c3eb4203 100644 --- a/robots/libero/guides/env_calibration.md +++ b/robots/libero/guides/env_calibration.md @@ -2,12 +2,15 @@ ## Current LIBERO MCP Runtime Contract -Use this file as a calibration reference only. For current MCP-based runs, use -structured MCP tools, do not issue file-based protocol commands, and do not manually manage -`env_server.py`. Do not read BDDL files or hidden task definition files to infer -coordinates. Do not expect object world coordinates in `states.json`; localize -objects through images_cam + depth/back_project, segment, and wrist/high-res -artifacts when available. +Use this file as a calibration reference for structured-tool runs. The runner +owns the environment server and the `EnvState` lifecycle. Do not issue +file-based driver commands, inspect observation storage directly, or read BDDL +files for coordinates. + +Start with `view_env_state({"step": 0})`. It returns the initial robot state, +top-level task language, logical observation references, and embedded camera +images. Use `back_project` or `segment` for geometry and `view_camera_meta` for +calibration. Step `-1` selects the latest record. Measured 2026-05-20 on `libero_10_with_mug` t0 (LIVING_ROOM frame) and t8 (KITCHEN frame). All probes use `move_to` with `gripper=-1` and tight @@ -17,9 +20,9 @@ Measured 2026-05-20 on `libero_10_with_mug` t0 (LIVING_ROOM frame) and t8 Each task scene uses one of the table fixtures below, which sets the entire world-frame z origin. The OSC workspace and all pick/place altitudes shift -accordingly. **Check `states.json[0].state.robot0_eef_pos[2]` in the initial -state and branch on it.** Do not read BDDL files for this; use runtime state and -visual evidence. +accordingly. **Check `state.robot0_eef_pos[2]` in the result of +`view_env_state({"step": 0})` and branch on it.** Do not read BDDL files for +this; use runtime state and visual evidence. | Fixture | eef home z | Table top z | Used by tasks | |---|---|---|---| @@ -119,7 +122,7 @@ limit 1.15). My libero_10 t0 used z=0.95 for travel — safe and consistent. ## Practical rules going forward -1. **Always read `states.json[0].state.robot0_eef_pos[2]` before computing any z target.** +1. **Always inspect the initial returned `state.robot0_eef_pos[2]` before computing any z target.** ≈ 0.68 → LIVING_ROOM; ≈ 1.17 → KITCHEN; ≈ 0.26 → OBJECT. Use the matching frame table above. 2. **Never command an eef z below the per-frame floor.** Going to z=0.42 @@ -143,24 +146,25 @@ limit 1.15). My libero_10 t0 used z=0.95 for travel — safe and consistent. `move_to` + `set_gripper` (last resort; unreliable for objects <6 cm — see `resources/libero/memory/feedback_scripted_pick_limits.md`). -## Calibration log files +## Calibration Records + +Each calibration motion returns its state, command result, elapsed time, and +embedded observation. Use that tool result immediately. To revisit a recorded +step, call `view_env_state({"step": N})`; use `-1` for the latest step. -Raw probe logs are preserved in: -- `{output_dir}/states.json` (one step entry per command — the per-command - audit; each entry has `command`, `result`, `state`, `elapsed_s`). -- Only kept for the most recent agent session; reproduce by re-running - the calibration with the snippet in the next section. +The internal `states.json` file is a versioned manifest owned by `EnvState`, not +a list for manual indexing. Calibration analysis should use structured tool +results rather than parsing storage files. ## Reproducer -Legacy calibration notes below describe the old file-protocol flow. In the -current MCP runtime, use the runner-managed environment and call structured MCP -tools instead. +Use the runner-managed environment and call structured tools: ```bash # For each z in 0.65 .. 0.42, call the MCP tool: move_to({"xyz": [-0.20, 0.10, 0.65], "gripper": -1, "tol": 0.008, "step_clip": 0.010, "max_steps": 80}) -# Then read states.json entry NN for final_eef_pos & final_dist_m. +# Inspect the returned result for final_eef_pos and final_dist_m. +# Call view_env_state({"step": -1}) only when the latest state is needed again. ``` diff --git a/robots/libero/guides/pro_hybrid_guide.md b/robots/libero/guides/pro_hybrid_guide.md index d7999405..7b3054cb 100644 --- a/robots/libero/guides/pro_hybrid_guide.md +++ b/robots/libero/guides/pro_hybrid_guide.md @@ -1,502 +1,197 @@ -# LIBERO-Pro Hybrid (Pi0.5 + LLM-in-the-loop) — Perception-Isolated Guide - -You are picking up the **LIBERO-Pro** evaluation track in **perception-isolated** -mode — the only mode in this repository (the legacy oracle-state mode is not -included here). - -> **Pi0.5 only does the grasp (`pi0_pick`). The LLM (you) handles every motion -> (`move_to`), every release, sequencing, retries — and you do not get GT object -> coordinates. You localize objects yourself from the depth + camera calibration -> the runtime dumps each step.** - -This document layers on the base playbook. Read it first: - -- [`strict_hybrid_guide.md`](./strict_hybrid_guide.md) — the perception protocol: - back-projection localization, the perception artifacts - (`images_cam/image_cam_NN.png` / `depths/depth_NN.npy` / `camera_meta.json`), - the primitive vocabulary, the Rules (0/1/2/4/5), and the `strict_perception` - audit format. **This is the source of truth for *how you localize*.** -- **This file** — LIBERO-Pro–specific setup, the four perturbation axes, the Pi0 - fullshot baseline, the frame split, and how perception isolation changes the - P2-swap story. - -> Your task comes from `states.json[0].task_language`; the BDDL is FORBIDDEN. -> Read the authoritative instruction (the BDDL `:language` tag, coord-free) that -> the runtime injects and obey it verbatim. Do **not** scrape the BDDL or import -> the benchmark: that is error-prone (wrong task-map index → wrong task) and a -> perception-isolation breach (the `:init` block holds the GT coordinates this -> mode withholds). For swap fixtures, localize them **visually** (§3.5), never -> from `:init`. You never hand-roll projection math — `back_project` applies -> `K⁻¹` + the cam→world extrinsic and returns the surface `world_xyz` under a -> pixel; just use it. - -Whenever this guide says "see the perception protocol" or "the localization -snippet", it means `strict_hybrid_guide.md`. **Every rule there applies here**; -this guide only *adds* PRO constraints and tooling. Runner contract: call the -structured MCP tools (the runner owns the env server — do not start/stop it, and -issue no file-based commands); it is a single-episode run — no `reset`/`exit`, -recover in place or write an honest failure audit and `finish`. - -## 0. What's different from oracle PRO mode (read this first) - -| | oracle mode (not in this repo) | **perception (this guide)** | -|---|---|---| -| how launched | oracle-state run | `rpent/cli/main.py --libero-type pro` (or `LIBERO_TYPE=pro`); perception artifacts always dumped, coords withheld | -| `states.json` objects | full `objects:{name:[x,y,z]}` | **`object_names:[…]` only — NO coords** | -| how you learn the task | env prompt / scrape BDDL | **`states.json[0].task_language`** (authoritative `:language`, coord-free) — never read the BDDL | -| extra obs artifacts | agentview RGB only | **+ `images_cam/`, `depths/`, `world/` (agentview); `images_wrist/`, `depths_wrist/`, `world_wrist/` (wrist); hi-res pairs; `camera_meta.json`** — all via `back_project` | -| cameras | agentview only | **agentview (fixed, ~1m → ±8-13cm) + eye-in-hand wrist (moves with gripper, ±1-2cm when <20cm to target)** — see §3.6 | -| how P2 swap is solved | read swapped coords from `states.json[0]` | **localize the swapped objects by `back_project`** | -| how a swapped *fixture* site is found | read the swap BDDL `:init` block | **localize the fixture visually** (see §3.5) | -| run budget | short | **larger** — perception localization + manipulation is slower (raise `--max-turns` / `--planner-timeout-s`) | -| audit `regime` | `strict` | `strict_perception` | -| which image you pick pixels in | agentview RGB | `images_cam/image_cam_NN.png` (or hi-res) for **pixel-picking**; `images/image_NN.png` only for a sanity glance | - -**The single most important conceptual shift.** In oracle PRO mode the headline -is "the hybrid beats Pi0 on P2 because it reads the *swapped* coordinates straight -out of `states.json[0]`, while Pi0 is prompt-/memory-blind." In **perception** -mode there are no coordinates to read — so the hybrid's P2 win now comes from -**seeing where the object is and back-projecting it**. This is a *stronger* claim -(no oracle state at all), and it is the whole point of running PRO in perception -mode: P1 (Task) is still won by reading the *language*; P2 (Position) is now won by -*perception*, not by an oracle. - -## 1. Why LIBERO-Pro - -LIBERO-Pro -([paper](https://arxiv.org/pdf/2510.03827), [repo](https://github.com/Zxy-MLlab/LIBERO-PRO)) -perturbs each base task along five axes; all end-to-end VLAs (OpenVLA / Pi0 / -Pi0.5 / UniVLA) collapse on the two strongest: - -| Axis | Suffix | Paper column | What changes | Headline result | -|---|---|---|---|---| -| **Task** | `_task` | **P1** | Instruction + goal predicate inverted | All VLAs ≈ 0.0 | -| **Position** | `_swap` | **P2** | Object **and fixture** initial positions swapped | All VLAs 0.0–0.4 | -| Semantic | `_lan` | — | Instruction paraphrased; goal unchanged | VLAs handle (memorize visual) | -| Object | `_object` | — | Object appearance / colour / scale | VLA visual policy stressed | -| Environment | `_environment` | — | Table / scene swapped | Visual policy stressed | - -The agentic hybrid wins on **P1 and P2** by routing the language channel and the -*perceived* spatial-state channel through the LLM. Object and Environment -perturbations enter through the Pi0 vision channel and the hybrid inherits the -VLA's weakness there — declare that upfront, don't oversell. **In perception -mode, the Object/Environment axes also stress your own localization** (a -recolored or rescaled object is harder to pixel-pick), so lean on the -multi-pixel-median tip from the perception protocol. - -## 2. Setup (do these once on a fresh checkout) - -Run the idempotent installer from the repo root: - -```bash -bash scripts/install_libero_pro_plus.sh -``` +# LIBERO-Pro Hybrid Perception Guide + +This guide extends [strict_hybrid_guide.md](./strict_hybrid_guide.md) for the +LIBERO-Pro evaluation tracks. Read the strict guide first; its runtime contract, +localization discipline, and single-attempt rules apply unchanged. -It does all four steps below: the liberopro editable install, applies the -benchmark-registration patch, syncs the authoritative HF dataset snapshot, and -verifies with the `get_benchmark(...).get_task(0).language` check. The perception -observables (depth + both cameras + hi-res) are unconditional once the runner -launches with `--libero-type pro`; there is nothing perception-specific to -install beyond the PRO setup itself. +LIBERO-Pro perturbs object placement, fixture placement, spatial relations, and +task language. The agent must solve the observed scene rather than replaying a +seed-specific command sequence. -### 2.1. LIBERO-PRO repo +## Start Of Run -Cloned at `${LIBERO_PRO_PATH:-/path/to/LIBERO-PRO}/` from -`https://github.com/RLinf/LIBERO-PRO.git` and installed editable into the openpi -venv: +Call: -```bash -python -m pip show liberopro -# Name: liberopro Version: 0.1.0 Location: ${LIBERO_PRO_PATH:-/path/to/LIBERO-PRO} +```json +{"step": 0} ``` -### 2.2. Apply the benchmark-registration patch +through `view_env_state`. From the returned tool result: -The upstream `__init__.py` does **not** expose the 16 perturbation suites through -`get_benchmark()`. Our patch -[`scripts/liberopro_register_perturbations.patch`](../../../../scripts/liberopro_register_perturbations.patch) -adds them and overrides `Task.language` to read each BDDL's actual `:language` -tag (so the perturbed instruction reaches Pi0 / hybrid). +- read top-level `task_language` verbatim; +- inspect `state.robot0_eef_pos` to identify the scene frame; +- inspect object names only as a scene inventory, never as coordinates; +- inspect the embedded agentview image for semantic identity and relations; +- inspect the wrist image only when it contains useful close geometry. -```bash -cd ${LIBERO_PRO_PATH:-/path/to/LIBERO-PRO} -git apply /scripts/liberopro_register_perturbations.patch -``` +Step `-1` means the latest record. Do not inspect or index `states.json` +directly; it is an internal manifest. -If already applied (likely), `git status -s` shows clean. If you reinstall -liberopro, re-apply. +## Perturbation Axes -### 2.3. Huggingface dataset (authoritative) +### Object perturbation -The LIBERO-PRO git repo ships **incomplete / broken** init files for several -perturbation suites (e.g. `libero_spatial_swap` has 0 BDDLs; some -`libero_spatial_task` `.pruned_init` files are 0 bytes). Treat the git repo as -unreliable for perturbation data. The full, correct set lives on Huggingface -([`zhouxueyang/LIBERO-Pro`](https://huggingface.co/datasets/zhouxueyang/LIBERO-Pro)), -persisted locally at: +Objects move relative to their seed positions. Re-localize every target and +destination from the current images. Seed recipes remain useful for primitive +ordering, prompt wording, safe heights, and known failure modes, but never for +coordinates. -``` -${LIBEROPRO_DATASET_PATH:-/path/to/liberopro_hf}/ -├── bddl_files/ 16 perturbation suites, 10 BDDLs each -└── init_files/ 16 perturbation suites, 10 init files each -``` +### Spatial perturbation -Covers `{libero_spatial, libero_object, libero_goal, libero_10} × {swap, task, -lan, object}`. The installer syncs this into the liberopro install (overwriting -the broken upstream files); if the persistent copy is gone, re-download with: +Relations such as left/right, on-top-of, or next-to may select a different +instance than in the reference scene. Identify the candidate satisfying the +current visual relation, then localize it with `back_project` or `segment`. -```bash -python -c " -from huggingface_hub import snapshot_download -snapshot_download(repo_id='zhouxueyang/LIBERO-Pro', repo_type='dataset', - local_dir='${LIBEROPRO_DATASET_PATH:-/path/to/liberopro_hf}', - allow_patterns=['bddl_files/**','init_files/**'])" -``` +### Goal perturbation -### 2.4. Verify +The destination or requested interaction changes. Re-read `task_language` and +classify the destination semantically. Do not infer the goal from the suite, +task index, object list, or a sibling result file. -```bash -LIBERO_TYPE=pro python -c " -import liberopro.liberopro.benchmark as bench -for n in ['libero_spatial_task','libero_spatial_swap','libero_spatial_lan']: - b = bench.get_benchmark(n)(); t = b.get_task(0) - print(f'{n} t0: {t.language!r} trials={len(b.get_task_init_states(0))}')" -``` +### Task-language perturbation -Expected: -``` -libero_spatial_task t0: 'Pick the akita black bowl not between the plate and the ramekin and place it on the plate' trials=50 -libero_spatial_swap t0: 'Pick the akita black bowl between the plate and the ramekin and place it on the plate' trials=50 -libero_spatial_lan t0: 'lift the black bowl between the plate and ramekin and set it on the plate' trials=50 -``` +The authoritative instruction is the current top-level `task_language`. Do not +read BDDL files. They combine language with hidden initialization data and +therefore violate perception isolation. -## 3. PRO-specific environment gotchas +## Scene Frame Selection -Everything in `strict_hybrid_guide.md` applies. The following are **additional** -PRO constraints (most live in detail at [`env_calibration.md`](./env_calibration.md)). +The initial end-effector height distinguishes the principal scene frames: -### 3.1. Three scene frames, picked per-task +| Initial EEF z | Frame | Typical fixtures | +|---|---|---| +| approximately 0.26 m | object/low-table | grocery and basket tasks | +| approximately 0.68 m | living-room table | plates, baskets, pudding | +| approximately 1.17 m | kitchen table | stove, cabinet, drawer, microwave | -PRO scenes use one of three table fixtures; the eef home z differs by up to -~0.9 m. +Use `view_env_state({"step": 0})["state"]["robot0_eef_pos"][2]` as the +measurement. Then use the matching safe-height guidance from +[env_calibration.md](./env_calibration.md). -| Fixture | eef home z | Table top z | xy reachable | Where | -|---|---|---|---|---| -| `living_room_table` | ≈ 0.68 | ≈ 0.43 | `(x∈±0.30, y∈±0.30)` | basket / plate / pudding | -| `kitchen_table` | ≈ 1.17 | ≈ 0.90 | `(x∈±0.30, y∈±0.30)` | stove / cabinet / drawer / microwave | -| `object` (low table) | ≈ 0.26 | ≈ 0.0 | `(x∈±0.30, y∈±0.30)` | `libero_object` grocery-into-basket | +## Mandatory Perception Table -**Mandatory check at session start: read `states.json[0].state.robot0_eef_pos[2]`** -(via `view_driver_state({"step": 0})`). ≈ 0.68 → LIVING_ROOM; ≈ 1.17 → KITCHEN; -≈ 0.26 → OBJECT. Pick `pre_pos_z` / `carry_z` / `release_z` accordingly (per-item -OBJECT-frame altitudes are in -`resources/libero/memory/project_libero_object_pro_done.md`). Sending a -wrong-frame z (e.g. KITCHEN coordinates while the env is in LIVING_ROOM frame) -crashes the env worker (EOFError, silent state loss). +Before manipulating, create one row for every task-relevant entity: -> **Perception note.** This proprioceptive z is *not* an object coordinate — it's -> the robot's own pose, which `states.json` still gives you in perception mode. -> Reading it to pick the frame is fine. It also doubles as a free depth sanity -> check: your back-projected table z should sit ≈0.25 m below eef home z (≈0.43 in -> LIVING_ROOM, ≈0.90 in KITCHEN). If a back-projection lands far from that, you -> picked a wrong pixel. +| field | meaning | +|---|---| +| role | target, destination, support, fixture, or relation landmark | +| semantic evidence | visual properties and relation establishing identity | +| agentview pixels | several interior pixels or a verified segment mask | +| agentview xyz | robust projection from the global camera | +| wrist refinement | accepted, rejected, basket-confirmed, or unnecessary | +| final xyz | coordinate used for planning | +| uncertainty | duplicates, occlusion, rim bias, label ambiguity, etc. | -For `libero_spatial_*` you are always in **KITCHEN frame**. Standard heights: +Do not begin manipulation while a required row lacks a defensible identity or +coordinate. -``` -pre_pos_z = 1.05 # ~7 cm above objects at z≈0.97 -carry_z = 1.10 # safe traversal, well under upper limit 1.15 -release_z = 1.01 # ~7 cm above plate top at z≈0.90 -``` +## Swapped Objects And Fixtures -These are *altitude* knobs (robot-frame), not object positions — you still derive -the object's **xy** by `back_project` every time. - -### 3.2. xy single-step ±0.30 cap - -OSC flips IK branches if you command `|x|>0.30` or `|y|>0.30` in a single -`move_to` (the eef lands in the wrong half-space and corrupts the run). **Never -command beyond ±0.30 in a single move** — split into carry-z waypoints. (Detail -in `env_calibration.md`.) - -### 3.3. Slow long-distance carry for swap variants - -P2 swap can move a large object 15 cm across the table. `step_clip=0.025` lets it -slip in the gripper and the object ends up centimetres off target. Mitigation -(proven on `_swap` t0, carries over directly): - -- `carry_z = 1.15` (higher than usual 1.10) -- `step_clip = 0.020` (slower) -- **Re-localize mid-travel** instead of trusting a cached `object_xyz - eef_xyz` - offset: `Read images_cam/image_cam_NN.png` (or the hi-res `images_cam_hi/`) - mid-carry and call `back_project` on the carried object's pixel; if the - perceived offset drifted >5 mm from post-pick, re-pre-position and re-`pi0_pick`. - -### 3.4. Task language is from BDDL, not filename - -After the patch, `get_task(i).language` returns the perturbed `:language` tag, and -the env passes it to Pi0 as the prompt (surfaced to you as -`states.json[0].task_language`). For `_task` and `_lan` this is the perturbed -instruction. **Don't override it** — falsifying the VLA's prompt-blindness is the -point. You read the same language to decide *which* object to localize and place. - -### 3.5. ⚠ Swap moves FIXTURES too — and you must localize them visually - -This is the biggest perception-mode trap, and it is **specific to `_swap`**. The -P2 perturbation does not only swap loose objects; for `libero_goal_swap` it swaps -entire **fixtures** (stove ↔ cabinet ↔ wine_rack), so a goal predicate like -`On(bowl, flat_stove_1_cook_region)` now points at wherever the *stove* was -relocated to, and there are no coordinates in `states.json` to read. See -`resources/libero/memory/feedback_swap_perturbs_fixtures.md`. - -In **oracle** mode the documented fix is to read the swap BDDL `:init` block and -recompute the fixture site's world coordinates. **That is forbidden here** — the -`:init` block is ground-truth geometry. In **perception** mode you instead: - -1. Identify the target fixture by name from the (perturbed) task language and - `state.object_names`. -2. **Localize the fixture's predicate site visually** in - `images_cam_hi/image_cam_hi_NN.png`: pick pixels on the stove's cook region / - cabinet top surface / rack top shelf, then `back_project` (sample 3–5 pixels, - median the xy — fixtures are large and the surface you want is the *placement* - surface, not the nearest edge). -3. Carry target xy = perceived site xy; descend to the perceived site z + a small - clearance, `release`, then **retreat with gripper open** — the predicate often - fires *during* the settle, not at the release step itself. - -So the swap-fixture problem becomes "find the fixture in the image" rather than -"read where the BDDL put it." Loose-object swaps (e.g. `libero_spatial_swap`, -where only bowls/plates move) are simpler: just localize each object by -`back_project` as usual; their new positions fall straight out of the depth. - -### 3.6. Two cameras — agentview = IDENTITY, wrist = GEOMETRY - -Do **not** restate the two-camera protocol here — the strict guide's **First-step -perception protocol** is the source of truth: agentview / agentview-hi is the -semantic identity authority (decides *which* object/surface satisfies the task -language + relation); wrist / wrist-hi refines geometry for the *same* candidate -(accept only within ~3–5 cm of the agentview anchor, never average, basket/cavity -excepted). Both maps share one world frame, read via -`back_project({"camera":"wrist"})`; the optional `segment` honours both cameras -(`{"camera":"agentview"|"wrist"}`, `min_score` default 0.2, or `point:[row,col]`). - -**PRO implication:** `_swap` (and `_task`) can *invert* target or destination -semantics, so **never reuse base-task or recipe coordinates** — the object -satisfying the relation may now sit where the base task never put it. Let -agentview decide *what*; the near-vertical wrist is BAD at identity (locks onto -look-alikes) and only refines *where*. - -### 3.6c. Mandatory pre-task perception pass - -Also owned by the strict guide's First-step perception protocol: before ANY -pick/place, build the localization table (one row per task-relevant entity) and -pass the FINAL READY CHECK. **PRO implication:** in single-attempt mode a -wrong-target first grab is unrecoverable, and under `_swap` the "right" target is -exactly the one you'd get wrong by habit — so localizing all entities up front is -cheap insurance, not optional. - -### 3.6b. Hi-res perception channel (1024×1024) - -Also owned by the strict guide's Hi-res perception channel section (1024×1024 -pairs each step; prefer the hi image for identification; `back_project` defaults -`resolution="high"`; never mix pixel grids). **PRO implication:** hi-res fixes -*which* object you point at — decisive for the recolored/rescaled Object axis and -for telling swapped same-shape groceries apart — but does NOT change metric -accuracy or replace the wrist coarse→fine refinement. **Identity caveat:** SAM3 -gives two same-shaped, different-brand objects the SAME category (it labelled the -tomato-sauce and alphabet-soup cans identically) — use its mask only as a category -candidate, then READ the label yourself in the hi-res crop to assign identity; -never let SAM3's category settle a brand choice. - -## 4. The four-cell experiment per (base task, seed) - -For each base task you claim coverage on, generate four runs — every run is a -perception cell: - -| Suite | Variant | Perception-mode expectation | -|---|---|---| -| `libero_spatial` | base sanity | Pi0 and hybrid both pass | -| `libero_spatial_task` | **P1 Task** | Pi0 ✗ (picks base target); hybrid ✓ (LLM flips target from instruction, then localizes it) | -| `libero_spatial_swap` | **P2 Position** | Pi0 mixed; hybrid ✓ — **by localizing the swapped object/fixture from depth, not from state** | -| `libero_spatial_lan` | Semantic | Both pass (paraphrase invariant) | +For swap-style perturbations: -Replace `spatial` with `object`, `goal`, or `10` for the other base suites. +1. Identify both swapped entities in the embedded agentview image. +2. Classify each by appearance and current relation, not expected seed layout. +3. Localize each independently. +4. Verify the chosen target still satisfies `task_language`. +5. Re-localize after any contact that could move either entity. -### 4.1. Hybrid run — the MCP runner +The runtime withholds privileged coordinates, so there is no coordinate field +to fall back on. The current images and geometry tools are the source of truth. -Launch a cell with the CLI; the runner owns `env_server.py`, exposes the -structured tools, and runs single-attempt: +## Destination Classification -```bash -python rpent/cli/main.py --env libero --suite --task --seed \ - --libero-type pro --planner claude_code --model claude-opus-4-8 +Pro scenes frequently contain look-alike surfaces. Before placement, explicitly +classify all plausible destinations: -# e.g. --suite libero_spatial_task 0 (P1) · libero_spatial_swap 0 (P2) · -# libero_goal_swap 2 (P2 fixture swap) · libero_10_task 5 (long horizon) -``` +- plate versus stove burner; +- cabinet top versus drawer opening; +- basket interior versus rim; +- microwave cavity versus door or surrounding counter; +- movable lid versus fixed fixture surface. -`--libero-type pro` may be given as `LIBERO_TYPE=pro` instead. +Only after semantic classification should you call `back_project` or `segment` +for coordinates. -Audit + recipe land in the run's `output_dir` (`output_dir` and `recipe_tag` -arrive in your first message): +## Mid-Carry Re-Localization -``` -{output_dir}/{recipe_tag}.json <- you write this (write_text_file) -{output_dir}/recipe_{recipe_tag}.jsonl <- exported automatically by the runner -``` +Long-horizon Pro tasks often move or occlude objects during earlier steps. +Before each new pick or placement: -Do NOT write into `resources/libero/results_*_pert/` — that tree is a **read-only -seed-0 reference corpus**, not a write target. - -### 4.2. Environment server is runner-owned - -There is no manual server to launch and no REPL to drive: the MCP runner starts, -manages, and tears down `env_server.py` for you and blocks each tool call until -the next `states.json` entry is dumped. Do not start, stop, or background it, and -do not poll for readiness. - -### 4.3. Pi0 fullshot baseline - -The baseline is Pi0.5 driving the task end-to-end with the runtime's own -(perturbed) `task_language`. There is **no standalone baseline CLI in this repo**; -the numbers to compare against are the recorded full-shot results (the team's -`SUCCESS_RATES` table). **Do not invent a `pi0_baseline.py` path.** - -Pi0 never sees object coords in either mode, so there is no perception variant of -the baseline. Expected behavior: - -- **P1 (task):** Pi0 "succeeds at the wrong task" — it picks the *base*-task - target object, places it on the plate, and `libero_terminated=False` because the - goal predicate names a different object. This is exactly the gap the hybrid - closes. -- **P2 (position):** Pi0 picks / places at the *base* (un-swapped) location. - -### 4.4. Audit JSON — PRO + perception fields - -Start from the `strict_perception` audit schema in the perception protocol, and -add the PRO fields: - -```jsonc -{ - "suite": "libero_spatial_swap", - "task_id": 0, - "seed": 0, - "regime": "strict_perception", - "perturbation_type": "swap (P2 Position perturbation)", - "perturbed_task_language": "", - "perturbation_semantics": "", - "expected_baseline_behavior": "", - "strategy_notes": "HOW you localized — which pixel(s) in images_cam, back-projected world xyz; for swap, how you found the relocated object/fixture", - "pick_result": { /* the pi0_pick step's result */ }, - "final_state": { /* latest states.json entry's `state` field */ }, - "libero_terminated": true -} -``` +1. call `view_env_state({"step": -1})` if the previous primitive result is + no longer in context; +2. inspect the newest embedded images; +3. re-localize any entity that may have moved; +4. update the working perception table; +5. verify the remaining primitive order still matches `task_language`. -`strategy_notes` **must** describe the localization (pixel → `back_project` → -world xyz). For swap cells, explicitly note that the relocated object/fixture was -found by perception, not by reading coords. If unrecoverable after honest -exploration, write `libero_terminated: false` with what you tried and which step -failed — never warp (teleport primitives are deleted; see Rule 4 in the -perception protocol). Write the audit with `write_text_file` to -`{output_dir}/{recipe_tag}.json`, then call `finish`. - -## 5. Perception-protocol rules — PRO clarifications - -- **Rule 0 (use images for reasoning).** Even more critical under PRO: P2 swap can - move a large object/fixture clear across the table, and you have *no* - coordinates to fall back on. `images_cam/image_cam_NN.png` + depth are your only - spatial truth — open `images_cam/` and describe the scene before deciding - targets. -- **Rule 1 (no `pi0_end_to_end`).** Pi0 does the grasp via `pi0_pick`; the LLM - scripts every motion + release. Under PRO this is doubly important — handing - back to Pi0 means handing back to the prompt-blind / memorized-place habit you - are trying to falsify. -- **Rule 2 (single-episode current run).** This is a one-shot eval: do not call - `reset` or `exit`. Recover *within* the episode when safe (re-localize, - re-pre-position, re-`pi0_pick`, walk the prompt ladder, - `rotate_pitch`/`move_pose`); otherwise write an honest stuck/failure audit and - `finish`. `_swap` typically needs an in-episode retry — document it. -- **Rule 4 (no teleport).** `set_object_pose`, `articulate_to`, `js_move_to`, - `carry_object` are deleted. A goal past OSC reach with no physical approach → - honest `libero_terminated:false`. -- **Rule 5 (assume solvable).** A localization that moves the gripper into thin - air means you picked a wrong pixel (a reflection, a rim, a decoy object under - the Object perturbation), not that the cell is unreachable. Re-look, re-pick, - re-`back_project` before concluding failure. - -## 6. Existing corpus +Never assume the initial coordinate remains valid after contact, release, or a +fixture interaction. -``` -robots/libero/guides/ -├── pro_hybrid_guide.md <- this file -├── strict_hybrid_guide.md <- perception protocol + Rules (source of truth) -└── env_calibration.md <- OSC frame bounds + safe altitudes -scripts/ -└── liberopro_register_perturbations.patch -resources/libero/memory/ <- MEMORY.md index + feedback_*/project_* notes -resources/libero/results_spatial_pert/ <- read-only seed-0 reference corpus -resources/libero/results_{object,goal,10}_pert/ <- same, other suites (seed-0) -``` +## Contact Tasks -Before you start, **read the auto-memory**: `resources/libero/memory/MEMORY.md` -(one-line hooks, auto-injected via CLAUDE.md). For perception PRO cells always -open `feedback_no_teleport_rule.md` and — for any `_swap` cell — -`feedback_swap_perturbs_fixtures.md` (what swaps, and why you re-find the -relocated fixture visually). For bowl→plate spatial tasks also read -`feedback_bowl_eef_y_offset.md`; for cluttered picks, `feedback_pi0_pick_full_prompt.md`; -after two failed retries, `feedback_failure_forensics.md`. The -`resources/libero/results_*_pert/` recipes are **inputs** (technique priors) — -consult them for prompt ladders, staging, and target zones, but never reuse their -coordinates (re-derive every xyz from THIS scene) and never write there. - -## 7. What to do next (priority order) - -1. **Extend spatial to all 10 tasks at seed 0**, four perception cells each - (base / `_task` / `_swap` / `_lan`). For hybrid runs, use the seed-0 reference - recipes in `resources/libero/results_spatial_pert/` as *technique* starting - points — the pick step is usually identical; the place target changes for - `_swap`, the target object changes for `_task`. Never reuse their coordinates. -2. **Scale to seeds beyond 0** (50 trials per task). Recipes must re-localize per - scene; a perception recipe is *data-flow* (perceive → plan → act), never - hard-coded xyz. -3. **Replicate on `libero_object`, `libero_goal`, `libero_10`.** Frame split - applies (read `states.json[0].state.robot0_eef_pos[2]`). For - `libero_goal_swap`, apply §3.5 — localize the **swapped fixture** visually. -4. **Aggregate into a main table** `(suite × perturbation × {Pi0, hybrid})`. The - headline number is the conditional: of the seeds Pi0 fails on, what fraction - does the *perception* hybrid solve? Because there is no oracle state anywhere in - the hybrid's reasoning here, this is the strongest single statistic the agentic - decomposition can claim. - -## 8. Quick reference for a brand-new session - -```bash -# 1. Sanity-check the liberopro patch (perturbed language must show) -LIBERO_TYPE=pro python -c \ - "import liberopro.liberopro.benchmark as b; print(b.get_benchmark('libero_spatial_task')().get_task(0).language)" -# -> must read 'Pick the akita black bowl not between ...' (the perturbed text) - -# 2. Read the auto-memory: resources/libero/memory/MEMORY.md - -# 3. Launch a perception cell (runner owns env_server; single-attempt) -python rpent/cli/main.py --env libero --suite libero_spatial_swap --task --seed 0 \ - --libero-type pro --planner claude_code --model claude-opus-4-8 -``` +Use `pi0_doubled` for learned contact behavior such as turning a knob or +opening/closing a drawer. Its success flag mirrors the benchmark termination +predicate, so an intermediate contact can be useful even when success is false. +Inspect state and image evidence after every contact attempt. + +Use short scripted alignment motions around contact skills. Avoid long blind +pushes, which can destabilize MuJoCo or move the end effector into an invalid IK +branch. + +## Planning Across Multiple Objects + +For multi-object tasks: + +- process objects in an order that minimizes collision and occlusion; +- preserve already-correct placements; +- route later carries around placed objects rather than over them; +- re-confirm each source and destination immediately before use; +- verify intermediate relations visually rather than assuming they held. + +When order is specified by the task, obey it even if another order appears +easier. + +## Reference Results And Memory + +Reference results and memory entries provide reusable strategy: + +- prompt ladders; +- manipulation ordering; +- safe approach and carry heights; +- fixture-specific contact patterns; +- known failure and recovery modes. + +They do not provide valid coordinates for the current run. Re-derive all +positions through perception. + +## Outcome And Audit + +The current benchmark outcome is top-level `libero_terminated` in the latest +environment tool result. It is not stored inside `state`. + +The audit should record: + +- exact perturbed `task_language`; +- perturbation type when known; +- semantic identification evidence; +- agentview and accepted wrist localization results; +- primitive sequence and recovery decisions; +- memory files consulted; +- final state and `libero_terminated`. + +Recipe export is runtime-managed from recorded primitives and successful +segmentation events. Artifact identifiers returned by tools are for audit and +traceability, not manual file access. + +## Quick Checklist -Then, inside the run: - -4. `view_driver_state({"step": 0})` → read `state.robot0_eef_pos[2]` to pick the - frame (§3.1). `Read images_cam/image_cam_00.png` (or hi-res) and - `view_camera_meta`. Localize the target (and, for `_swap`, the relocated - object/fixture) with `back_project` — run the mandatory pre-task perception pass - (§3.6c). Plan, then execute one structured tool at a time. -5. `write_text_file` the audit to `{output_dir}/{recipe_tag}.json` - (`regime: strict_perception`) before `finish`; the recipe - `{output_dir}/recipe_{recipe_tag}.jsonl` is exported automatically by the runner. - -When in doubt about *how to localize* or a primitive, the source of truth is -[`strict_hybrid_guide.md`](./strict_hybrid_guide.md); about *PRO setup / -perturbation semantics*, see -[`scripts/install_libero_pro_plus.sh`](../../../../scripts/install_libero_pro_plus.sh) -and §2. +- Read strict guide and relevant memory. +- Call `view_env_state({"step": 0})`. +- Select the scene frame from initial EEF z. +- Read top-level `task_language`. +- Build the complete perception table. +- Re-derive all coordinates for this scene. +- Use agentview for identity and wrist for consistent refinement. +- Re-localize after contacts and placements. +- Check top-level `libero_terminated` after each relevant action. +- Write an honest audit and call `finish` without resetting. diff --git a/robots/libero/guides/strict_hybrid_guide.md b/robots/libero/guides/strict_hybrid_guide.md index 683feabe..417b380e 100644 --- a/robots/libero/guides/strict_hybrid_guide.md +++ b/robots/libero/guides/strict_hybrid_guide.md @@ -1,611 +1,206 @@ -# Strict Hybrid LLM + Pi0.5 — Perception-Isolated Guide - -You are taking over a hybrid LIBERO experiment in **perception-isolated** mode -— the only mode in this repository (the legacy oracle-state mode, where the -state JSON carried GT object coordinates, is not included here). - -> **Pi0.5 only does the grasp (`pi0_pick`). The LLM (you) handles every motion -> (`move_to`), every release, sequencing, retries — and you do not get GT -> object coordinates. You localize objects yourself from the depth + camera -> calibration the toolkit dumps each step.** - -## What's different from legacy oracle mode (read this first) - -| | legacy oracle mode (not in this repo) | **perception (this guide)** | -|---|---|---| -| how object coords are withheld | (none — full GT coords in `state`) | **the runner withholds them: `states.json` carries `object_names` only, no coordinates** | -| `states.json` objects | full `objects:{name:[x,y,z]}` | **`object_names:[…]` only — NO coords** | -| how you learn the task | env prompt / BDDL | **`states.json[NN].task_language`** (authoritative `:language`, coord-free) — never scrape the BDDL | -| extra obs artifacts | `images/image_NN.png` only | **+ `images_cam/`, `depths/`, `world/` (agentview); `images_wrist/`, `depths_wrist/`, `world_wrist/`, `wrist_meta/` (wrist); top-level `camera_meta.json`; hi-res `images_cam_hi/`, `world_hi/`, `images_wrist_hi/`, `world_wrist_hi/`** | -| cameras | agentview only | **agentview (fixed, ~1m → ±8–13 cm) + eye-in-hand wrist (moves with gripper, ±1–2 cm when <20 cm to target)** — coarse→fine, see below | -| how you get an object's xyz | read `state["objects"][name]` | **pick the object pixel, then `back_project({"row":ROW,"col":COL,"step":NN})` (K⁻¹+extrinsic already done); refine with the wrist map up close** | -| how you confirm the grasp | GT object-lift oracle | **grasp-only, no oracle — you judge the grasp from gripper width + wrist cam (Rule 1 / 1b)** | -| how you pick which of two identical objects | read their distinct coords | **by SPATIAL RELATION from `task_language` (elevation / left-right), never by `_1`/`_2` name** | -| cell budget | 600 (short suites) | **1200** — perceptual localization + manipulation is slower | -| audit `regime` | `strict` | `strict_perception` | - -How localization works (the core of this mode): - -The toolkit already back-projects EVERY pixel for you into world coordinates and -saves them as `world/world_NN.npy` (agentview) and `world_wrist/world_wrist_NN.npy` -(wrist) — both in the SAME world frame. **You do NOT write back-projection math; -you pick a pixel and call `back_project`, which indexes the map for you.** - -> **Path convention:** below, a bare file name refers to the file inside its own -> `output_dir` subdirectory — e.g. `image_cam_hi_NN.png` lives at -> `images_cam_hi/image_cam_hi_NN.png`, `world_NN.npy` at `world/world_NN.npy`, -> `image_wrist_hi_NN.png` at `images_wrist_hi/image_wrist_hi_NN.png` (see the obs -> artifacts row above for the full directory list). `NN` is the zero-padded step -> index. In practice you rarely hand-build these paths — each primitive tool -> already returns the resolved `image_cam_hi_path` etc. for the new step. - -1. Read `image_cam_hi_NN.png` — agentview RGB **in the calibration frame** - (vertical-flip of the raw buffer). This is the image you pick pixels in. - `image_NN.png` is the Pi0-rotation frame; **do not pick pixels there**. -2. Find the target object visually → pixel `(row, col)`. -3. Read its world xyz directly: `back_project({"row":ROW,"col":COL,"step":NN})`. - Sample 3–5 pixels on the object's top surface and median the returned xy - (robust to one mis-picked rim/edge pixel). That's the object's surface point - in world frame. - -### Hi-res perception channel (ON BY DEFAULT, 1024×1024) - -The toolkit dumps, every step, a 1024×1024 pair per camera IN ADDITION to the -256 files: `images_cam_hi/image_cam_hi_NN.png` + `world_hi/world_hi_NN.npy` -(agentview) and `images_wrist_hi/image_wrist_hi_NN.png` + -`world_wrist_hi/world_wrist_hi_NN.npy` (wrist). -**PREFER the hi pair for looking and identification** — a far object spans 4× -more pixels, and package text/label art is actually legible (e.g. "Cream -Cheese" vs "BUTTER" boxes are readable at 1024 and indistinguishable smears -at 256). If the hi files are absent, everything below works unchanged at 256. - -- A hi-res pixel `(row,col)` indexes ONLY the hi-res world map (`back_project` - default `resolution:"high"`; float16, same world frame). Never index a 1024 - pixel into the 256 map or vice versa — pass `resolution:"low"` only for a - pixel taken from a 256 image (convert by dividing/multiplying by 4 if needed). -- The `segment` command automatically uses the hi-res frame when present - (its `centroid_pixel`/`box` are then in 1024 coords). -- Metric accuracy of mask-median localization is the SAME at both resolutions - (the residual ~2 cm error is the surface-vs-center offset, not pixel size) — - the hi channel is for **identification**, not for replacing the wrist-cam - fine-localization protocol. -- Disk note: hi files keep only the LAST ~5 steps (rolling window); the 256 - history is complete. If the hi files are absent, everything below works - unchanged at 256. - -### Two cameras — agentview = IDENTITY, wrist = GEOMETRY - -**Roles are NOT symmetric (this is the core discipline, from the 80-task -localization sweep):** -- **agentview** (`image_cam_hi_NN.png` + `world_hi/world_hi_NN.npy`, ~1 m away): - the **semantic IDENTITY authority** — decides WHICH object/surface satisfies - the task language + spatial relation. At 1024 a far object spans enough pixels - to read labels/shape. Its metric precision is only ±8–13 cm, but identity is - its job, not millimetres. -- **wrist** (`image_wrist_hi_NN.png` + `world_wrist_hi/world_wrist_hi_NN.npy`, - <20 cm → ±1–2 cm): a **GEOMETRY refinement** camera for the SAME candidate - agentview already chose. It is near-vertical and **bad at identity** — it - cannot read side labels or tell duplicate/similar items apart, and in failed - probes it locked onto a look-alike hundreds of pixels away. **NEVER let the - wrist freely re-identify a non-basket target.** - -Protocol (non-basket objects/surfaces): -1. **Identity (agentview)** — in `image_cam_hi_NN.png` choose the target by RGB / - label / shape / global spatial relation; sample 3–8 pixels on it and median - `back_project` → the **identity anchor** (rough xy ok). -2. **Approach** — `move_to` ~15–20 cm directly above that anchor xy (toolkit - re-renders both cams after every primitive). -3. **Geometry refine (wrist)** — pick the SAME candidate's pixel in - `image_wrist_hi_NN.png`, `back_project` it with `camera:"wrist"`. **Accept - this corrected xy ONLY if it is within ~3–5 cm of the agentview anchor**; if - it jumps >5 cm, REJECT it (it hit a look-alike/background) and keep the - agentview xyz. The wrist may sharpen coordinates but may NOT override the - agentview semantic choice. Never average the two. -- **Basket / cavity special case:** for `basket`/cavity the wrist MAY also - confirm/refine the true interior centre (basket failures are rim/edge bias, - not semantic confusion). - -### Mandatory pre-task perception pass — localize EVERYTHING, THEN act - -Before ANY pick/place, build a localization table in your reasoning — one row -per task-relevant entity (every movable target, every destination/support/ -fixture, every relation landmark named by `task_language`), each with: -`name_or_role · agentview_evidence (why this candidate) · agentview_pixels · -agentview_xyz (median back_project on hi-res) · wrist_refine -accepted|rejected|basket_confirmed · final_xyz · uncertainty`. Do the -**FINAL READY CHECK** (every entity has a final xyz; non-basket wrist -refinements are spatially consistent with agentview; basket points are -interior-centred) — only then start manipulating. Rationale: in single-attempt -mode a **wrong-target first grab is unrecoverable** (it tips/displaces both the -grabbed object and the target zone), so the cheap insurance is to identify all -entities up front instead of recovering later. - -> The underlying math (already done for you): `cam = [(col−cx)·z/fx, -> (row−cy)·z/fy, z]` with `z=depth[row,col]`, then `P_world = -> extrinsic_cam2world @ [cam,1]`. You must invert `K` BEFORE the extrinsic — the -> old `E @ [col·z, row·z, z, 1]` recipe (no `K⁻¹`) is **wrong** (metres off). -> The `world/` and `world_wrist/` maps bake in the correct `K⁻¹`+extrinsic, and -> `back_project` reads them for you — that is precisely why `back_project` -> exists; you never write this math yourself. Forward world→pixel was verified -> 5/5 vs GT at design time (plate Δ=6 mm). The wrist map may be all-table/null -> early (the wrist only sees gripper + table until you move it over a target). - -## Before you start: READ THE AUTO-MEMORY - -Operating wisdom lives in the in-repo memory: +# Strict Hybrid LLM + Pi0.5 Perception Guide +This guide describes the perception-isolated LIBERO runtime. The agent receives +robot proprioception, object names, task language, and camera observations. It +does not receive privileged object coordinates. + +Pi0.5 performs grasping through `pi0_pick`. The LLM owns semantic perception, +localization, scripted motion, release, verification, and recovery within the +current episode. + +## Runtime Contract + +Every motion primitive appends a `StepRecord`. The record contains: + +- `step_idx`: sequential motion-step number; +- `state`: robot proprioception and coordinate-free scene information; +- `artifacts`: sorted logical base names for all files captured at the step; +- `command`, `result`, and `elapsed_s`; +- `extras`: LIBERO outcome fields such as `task_language`, + `libero_terminated`, and `episode_truncated`. + +Artifact storage is private to `EnvState`. Do not construct filenames or read +observation files manually. Use the structured tools: + +- `view_env_state`: state, artifact names, log, and embedded images; +- `view_camera_meta`: calibration data for one camera and step; +- `back_project`: matching image pixel or region to world coordinates; +- `segment`: text- or point-prompted mask plus projected `world_xyz`. + +Step `0` is the initial observation. Step `-1` means the latest observation. +When a tool omits `step`, its schema default is `-1`. + +`states.json` is an internal versioned manifest. It is not an agent-facing state +API and should not be indexed directly. + +## Camera Roles + +`view_env_state` embeds the best available images for the selected step: + +- **policy image**: the Pi0-oriented agentview input; +- **agentview image**: fixed global view, high resolution when available; +- **wrist image**: eye-in-hand view, high resolution when available. + +Use agentview for semantic identity and global relationships. Use wrist for +close-range geometry after the target identity has already been anchored in +agentview. The wrist view must not freely replace the semantic decision when +similar objects are nearby. + +Pixels passed to `back_project` must come from the same camera and resolution +specified in the call. The tool selects the matching world map internally. + +## Initial Perception Pass + +Before the first manipulation primitive: + +1. Call `view_env_state({"step": 0})`. +2. Read the returned top-level `task_language` verbatim. +3. Identify every movable target, destination, support, and relation landmark + named by the task. +4. Inspect the embedded agentview image and classify candidates by color, + shape, label, container type, and spatial relation. +5. Localize each chosen candidate with several interior pixels and + `back_project`, or use `segment` when a stable text/point prompt exists. +6. Record a working table with semantic evidence, sampled pixels, projected + coordinates, uncertainty, and the selected final coordinate. + +Do not start manipulation until every required target and destination has a +defensible identity and coordinate estimate. + +## Semantic Identity Before Geometry + +Depth and world coordinates cannot distinguish visually similar surfaces. A +plate, stove burner, pot lid, and cabinet top can all be flat circular or planar +regions at similar heights. + +Classify the destination in RGB before localizing it: + +- **plate**: ceramic disc with a clean rim, often white or ringed; +- **stove region**: darker fixture surface, coil, grate, or burner pattern; +- **basket**: open container whose interior center differs from the rim; +- **cabinet or drawer**: fixture geometry associated with the task noun. + +When duplicate objects exist, choose by the relation in `task_language`, not by +internal object suffixes. + +## Coarse-to-Fine Localization + +For each non-basket object: + +1. Select the semantic candidate in agentview. +2. Project three to eight interior pixels and take a robust median. +3. Move approximately 15-20 cm above the agentview anchor. +4. Inspect the wrist image and refine the same physical candidate. +5. Accept the wrist estimate only when it remains within roughly 3-5 cm of the + agentview anchor. Otherwise reject it and retain the global estimate. + +For baskets and cavities, use agentview for global identity and wrist for the +interior center. Avoid rim pixels. `back_project` region mode can summarize a +bounded pixel window with optional `z_min` and `z_max` filtering. + +## Segmentation + +`segment` does not move the robot. Supply exactly one of: + +```json +{"prompt": "the black bowl on the stove", "camera": "agentview", "step": -1} ``` -resources/libero/memory/MEMORY.md + +```json +{"point": [420, 615], "camera": "agentview", "step": -1} ``` -Scan the ~30 one-line hooks. For perception cells **always** open: -- `feedback_no_teleport_rule.md` — the deleted primitives. -- `feedback_redo_cell_timeout_1200.md` — why the cell budget is long here. -(The grasp is oracle-free too: there is no GT-lift oracle — you judge the grasp -from gripper width + the wrist cam, see Rule 1 / 1b.) +The result includes the mask score, projected `world_xyz`, logical artifact +identifiers for audit, and an embedded overlay image when available. Inspect +the overlay before trusting the projection. -For bowl→plate spatial tasks always also read `feedback_bowl_eef_y_offset.md` -(bowl-eef y-offset 4.5 cm: place at `eef_y = plate_y + 0.045`, not `plate_y`). +Use plain visual phrases. Remove benchmark or brand names that the segmenter +cannot ground. If segmentation fails, choose pixels manually and call +`back_project`. -## Rule 1 — `pi0_pick` is grasp-only (no oracle) +## Grasping -`pi0_end_to_end` is FORBIDDEN. `pi0_pick` is for the grasp only; you script every -`move_to` and every `release`. Use: +Pre-position above the localized object before calling Pi0.5: -``` -pi0_pick({ - "prompt": "", +```json +{ + "prompt": "pick up the short red-label can", "max_chunks": 20, "lift_thresh": 0.05, "gripper_closed_thresh": 0.06 -}) +} ``` -`pi0_pick` takes **no object-tracking / oracle argument** — it is grasp-only -and reads NO GT object pose. Passing a name would do nothing even if you tried. -You judge "did I grab the target?" yourself (Rule 1b). NEVER let Pi0 finish the -place — YOU do every `move_to` and the `release`. - -## Rule 1b — JUDGE THE GRASP from perception, NOT from a name - -After a pick, decide "did I grab the target?" from two coord-free signals: - -- **Gripper** (`state.robot0_gripper_qpos` from the latest `states.json` - entry): fingers closed but NOT fully shut (~0.01–0.05 gap) ⇒ holding an - object; fully closed (~0.0) ⇒ grasped air. -- **Wrist cam**: after the lift, read `image_wrist_hi_NN.png` — the target - should be raised into the gripper and the spot it came from now empty (its - surface z jumps up by your lift distance). If needed, `back_project` wrist - pixels to confirm the surface z jumped. - -`pi0_pick`'s returned `success` (an eef-lift + gripper-closure heuristic, pure -proprioception) is a HINT, not proof — confirm with the wrist cam before -carrying. The only authoritative TASK-success signal is -`state.libero_terminated` (the benchmark predicate), which is name-independent. - -## Rule 4 — NO TELEPORT primitives (physics-only) - -`set_object_pose`, `articulate_to`, `js_move_to`, `carry_object` are **deleted -from the codebase**. They are not callable and are not in the tool list. If a -goal is past OSC reach and no physical approach works, write an honest -`libero_terminated:false` audit — never warp. - -## Rule 0 — Use images for reasoning, not just JSON state - -After every primitive tool call, inspect the returned state and image paths -(the tool returns them — no separate read needed). Read the new -`image_cam_hi_NN.png` path (calibration frame — the one you pick pixels in) -and, when close to a target, the `image_wrist_hi_NN.png` path. Even more -important here than in oracle mode: this is *the* signal you use to find -objects. - -`image_cam_hi_NN.png` (and its 256 counterpart `image_cam_NN.png`) is the -calibration-frame RGB — same scene as `image_NN.png` but vertically flipped so -that pixel coordinates align with the camera matrices in `camera_meta.json`. -Pick object pixels from these; `states.json` alone gives only proprioception + -object names. - -## Rule 2 — SINGLE EPISODE, NO RESET - -This is a **one-shot** evaluation: you get **exactly ONE episode**. Do NOT reset -and do NOT restart the episode. You MAY recover *within* this one episode -(re-localize — objects may have moved; re-pre-position; re-`pi0_pick` a missed -grasp; walk the next rung of the Pi0 prompt ladder in Rule 3; re-firm the grip; -`rotate_pitch`/`move_pose`) — that is all one continuous attempt. But the -instant you would want to start over, **STOP instead and write the audit** -(success or an honest `libero_terminated:false`), then call `finish`. Never -warp; never reset. - -## Rule 5 — Assume every task is physically solvable - -Same as oracle mode. A localization that "looks right" but moves the gripper -into thin air usually means your pixel was on a wrong surface (e.g. picked the -bowl's reflection on the table). Re-look at `image_cam_hi_NN.png`, pick a -different pixel firmly on the target's top, re-`back_project`. Don't conclude -"unreachable" until you've validated localization. - -## Rule 6 — Your task is `task_language`; the BDDL is FORBIDDEN - -Each `states.json[NN]` entry carries a **`task_language`** field — the -authoritative task instruction (the BDDL's `:language` tag, which contains -**no** coordinates). **Read it and obey it verbatim.** Do not infer the task -from object names, from sibling recipes, or by guessing a `task_map` index — -that produced *wrong-task* runs (an agent solving a task it was never assigned). - -> ⚠️ **Never read the BDDL files, import the benchmark, or query env object -> poses.** The BDDL is the one place the task language *and* the `:init` -> ground-truth coordinates live together — reading it to get the language also -> leaks the coordinates this mode exists to withhold (a perception-isolation -> breach; observed on several swap cells in earlier runs). You already have the -> task from `task_language`; you get object **positions** ONLY by depth -> back-projection. The runtime strips coords from `state` but does **not** -> sandbox the BDDL — that discipline is on you. - -## Rule 7 — Ground the target by its spatial RELATION, not its name - -When the task names a relation ("the bowl **ON THE COOKIES BOX**", "the mug -**LEFT OF** the plate"), the target is whichever object SATISFIES that relation -in the scene — find it by perception. Identical objects (`akita_black_bowl_1` vs -`_2`) carry NO perceptual difference in their names, so the name can't choose -for you; the RELATION disambiguates: - -- "on the cookies box" ⇒ the bowl that is **elevated** (~0.03–0.06 m above the - table, on the box) — distinguish it by its higher world-z from `back_project` - vs the table-level bowl. -- "left/right/front/back of X" ⇒ compare back-projected world xy to X's xy. - -Pick the target purely from where things ARE. (You never need to know which -`_N` name the target is — no primitive in this mode asks for one.) - -## Mental model - -1. **The runner (`rpent/cli/main.py`) owns a long-lived env server** — Pi0.5 + a - single-env LIBERO sim. It launches and manages the server; you do NOT start, - stop, or restart it. -2. **You call one structured MCP tool per step.** The tool BLOCKS until the - toolkit runs that one primitive and dumps the new step, then RETURNS the new - state + log + image paths. There is no file bus and no polling — the tool's - return value IS your signal. Each dumped step appends an entry to - `states.json` (entry `[NN]`) and writes `images/image_NN.png` + - `images_cam/image_cam_NN.png` + `depths/depth_NN.npy` + `world/world_NN.npy` - (+ the wrist and hi-res dirs) and, once, top-level `camera_meta.json`. -3. **You read the returned state, localize via `back_project`, decide the next - move, and call the next tool.** - -## Launch a session - -The runner (`rpent/cli/main.py`) launches and owns the env server (Pi0.5 + single-env -sim) — do not start/stop it. You call MCP tools; begin by reading step 0 via -`view_driver_state({"step":0})`. - -## The perception artifacts you read each step - -| artifact | what's in it | -|---|---| -| `states.json` (entry `[NN]`) | `step_idx`, `libero_terminated`, **`task_language` (your authoritative task instruction — the BDDL `:language` tag, coord-free; obey it verbatim)**, `state.{robot0_eef_pos, robot0_eef_quat, robot0_gripper_qpos, object_names}`, and the merged `command`/`result`/`elapsed_s` for that step. **No object coordinates.** Read via `view_driver_state({"step": NN})` (omit `step` = latest). | -| `images/image_NN.png` | RGB in Pi0 frame (180° rotated). *Do not pick pixels here for back-projection.* | -| `images_cam/image_cam_NN.png` | agentview RGB in **calibration frame** (vertical flip). Pick object pixels HERE (256 grid → `back_project` with `resolution:"low"`). | -| `images_cam_hi/image_cam_hi_NN.png` | **HI-RES 1024×1024** agentview RGB, calibration frame. **PREFER this for looking / identification**; `back_project` its pixels at the default `resolution:"high"`. | -| `depths/depth_NN.npy` | `(256, 256) float32` agentview metric depth (m), calibration frame. Same row/col as `image_cam_NN.png`. | -| `world/world_NN.npy` | `(256, 256, 3) float32` — **precomputed agentview world xyz per pixel** (K⁻¹+extrinsic done). Prefer `back_project`; read this manually only for debugging. Fixed cam (~1 m). | -| `world_hi/world_hi_NN.npy` | `(1024, 1024, 3) float16` — precomputed agentview world xyz per hi-res pixel. | -| `images_wrist/image_wrist_NN.png` (+ `images_wrist_hi/`) | **wrist (eye-in-hand) RGB**, calibration frame. Moves with the gripper. | -| `depths_wrist/depth_wrist_NN.npy` | wrist metric depth (m). | -| `world_wrist/world_wrist_NN.npy` (+ `world_wrist_hi/`) | `(256, 256, 3)` / `(1024, 1024, 3)` **precomputed wrist world xyz per pixel**, SAME world frame as agentview. ±1–2 cm when <20 cm to target. May be all-table/null until you move over the target. | -| `wrist_meta/wrist_meta_NN.json` | wrist intrinsics + extrinsic **for THAT step only** (the wrist cam moves). Read via `view_camera_meta({"camera":"wrist","step":NN})`. | -| `camera_meta.json` (top-level) | agentview `intrinsic_K` (3×3), `extrinsic_cam2world` (4×4), `depth_near/far`, projection recipe. Read via `view_camera_meta({"camera":"agentview"})`. | -| `segments/segment_NN_XX.json` | (after a completed `segment` response) `{found, mode, prompt|point, camera, source_step, score, box, mask_shape, centroid_pixel, world_xyz, n_pixels}` — SAM3's top mask back-projected via the matching world map. `world_xyz` is a robust MEDIAN over the whole mask. A valid no-detection response carries `{found:false, error}`. `XX` is a per-step index. | -| `segments/segment_overlay_NN_XX.png` | (only after a successful `segment` call) the segmented mask tinted red on the source image — read it to confirm SAM3 grabbed the right object. | - -The command + its result + `elapsed_s` are merged INTO the `states.json[NN]` -entry (and echoed in each primitive tool's return value) — there is no separate -per-step log file to read. - -## Localization with `back_project` (coarse agentview → fine wrist) - -You index the precomputed map through `back_project` — no K⁻¹ math. Coarse -(agentview) then, once the gripper is parked over the target, fine (wrist): +Treat `pi0_pick.success` as a hint. Verify the grasp from: -``` -# COARSE: agentview hi-res (default resolution:"high") — choose which object / rough xy -back_project({"row": ROW, "col": COL, "step": NN}) +- end-effector lift; +- nonzero gripper opening consistent with holding an object; +- wrist or agentview evidence that the correct object moved with the gripper; +- an empty or changed source location. -# FINE: wrist — after move_to ~15-20cm above the target, pick its pixel in -# image_wrist_hi_NN.png and back_project it (±1-2cm; refines the SAME candidate). -back_project({"row": ROW, "col": COL, "step": NN, "camera": "wrist"}) -``` +If the grasp failed, recover in the same episode by re-localizing, +re-positioning, and improving the visual prompt. Do not reset in single-attempt +evaluation mode. -(Pass `"resolution":"low"` only when `ROW,COL` came from a 256 image.) - -Tips: - -- Sample 3–5 pixels on the object (centre + a couple of edge pixels) and - median the back-projected xy — robust to a single mis-picked pixel. Avoid - pixels on the thin rim/edge or the gap to the table (those index a - background/edge depth and give a world point metres away). -- The returned point is the **visible surface** under your chosen pixel. - For a flat object (plate, basket) that surface ≈ the place target. For a - bowl/bottle, the surface is the top of the object; the rim's xy is what - you want for `release`, not the grasp's eef_y (apply - `feedback_bowl_eef_y_offset` — bowl `eef_y_target = perceived_plate_y + 0.045`). -- For the **table z** (when you need a known floor to compare against), - `back_project` a pixel on bare table near the object. - -## The command vocabulary - -Call one structured tool per step. The full primitive set (only the *control -signals* are listed here) — every one blocks until the step is dumped and -returns the new state + log + image paths: - -```jsonc -// === physics-only primitives (the entire allowed set) ===================== - -// Scripted EEF servo. action_scale=0.05 is the env's units; step_clip caps -// per-step Δxyz (m) BEFORE division by action_scale — smaller = slower. -move_to({"xyz": [x, y, z], "gripper": -1, - "tol": 0.012, "step_clip": 0.025, "max_steps": 80, - "action_scale": 0.05, "target_yaw": null}) // gripper: -1 open, +1 close - -// Pi0.5 closed-loop pick. Grasp-only — no oracle/tracking arg (Rule 1). -// You judge the grasp from gripper width + wrist cam (Rule 1b). -pi0_pick({"prompt": "pick up the X", "max_chunks": 20, - "lift_thresh": 0.05, "gripper_closed_thresh": 0.06}) - -// Pi0.5 for a contact skill (knob turn, drawer/door open-close — rare here). -// success mirrors libero_terminated only; inspect image/state for intermediates. -pi0_doubled({"prompt": "turn off the stove", "max_chunks": 20}) - -// Open gripper to place. Triggers libero termination if the On/In predicate met. -release({"max_steps": 20}) - -// Hold pose + drive gripper. Use to firm a grip mid-carry. -set_gripper({"gripper": 1, "steps": 5}) // -1 open, +1 close - -// Wrist yaw (world-z). Provide target_yaw (absolute) OR delta_yaw (relative). -rotate_wrist({"target_yaw": 0.0, "gripper": 1, - "max_steps": 40, "tol": 0.02, "step_clip": 0.10}) - -// Tilt eef pitch (axis-angle X). Cavity entry / micro-aiming. target_pitch OR -// delta_pitch. Use before threading a narrow opening whose face normal is ±y. -rotate_pitch({"target_pitch": 0.9, "gripper": 1, - "max_steps": 40, "tol": 0.02, "step_clip": 0.10}) - -// Co-vary xyz + pitch + yaw — threads cabinet-front IK singularity that move_to -// walls at. gripper defaults to -1 (OPEN) — pass gripper:1 while holding. -move_pose({"xyz": [x, y, z], "target_pitch": 0.0, "target_yaw": 0.0, - "gripper": 1, "step_clip": 0.02, "pitch_step": 0.08, "yaw_step": 0.08, - "tol": 0.012, "ori_tol": 0.05, "max_steps": 150}) - -// SAM3-grounded localization — does NOT move the robot and does NOT -// replace manual back-projection. Segments the most-recent dumped image for the -// prompt, back-projects the mask through the matching world map, and writes -// segments/segment_NN_XX.json {score, box, centroid_pixel, world_xyz (robust MEDIAN over -// the whole mask), n_pixels} + segments/segment_overlay_NN_XX.png. Read world_xyz -// directly instead of eyeballing a pixel. camera "wrist" uses the wrist world -// map for fine refinement (move the eef over the target first). Pass -// "point":[row,col] for a point prompt instead of text; prompt and point are -// mutually exclusive. min_score default 0.2. -// RPent manages the SAM3 service. If one call fails or finds nothing, the result is -// an {"error":...,"fallback":...} dict — fall back to picking a pixel in -// image_cam_hi_NN.png and calling back_project. -segment({"prompt": "the black bowl on the stove", "camera": "agentview", - "point": null, "min_score": 0.2}) - -// === FORBIDDEN — DELETED FROM THE CODE — DO NOT EMIT ============ -// set_object_pose, articulate_to, js_move_to, carry_object -``` +## Scripted Motion + +Use short, staged `move_to` commands. Do not issue a single horizontal move +larger than approximately 0.30 m. Long commands can switch IK branches and move +the end effector into the wrong half-space. + +After every motion: + +1. Inspect the returned `result` and final distance. +2. Inspect the new state and embedded images. +3. Re-localize anything that may have moved. +4. Continue only when the observed state matches the plan. + +Use `rotate_wrist`, `rotate_pitch`, or `move_pose` when orientation matters. +Use `pi0_doubled` for short learned contact interactions such as knobs, +buttons, doors, or drawers. Alternate it with short, capped scripted alignment +motions; never use one long blind push. + +## Placement + +Carry objects at a collision-safe height. Reconfirm the destination before +release, especially when plate/burner or rim/interior confusion is possible. + +For open containers: + +- localize the interior rather than the nearest rim; +- place above the interior center; +- lower enough to avoid a high-energy drop; +- inspect the post-release observation and outcome flag. + +If the task predicate does not fire, reclassify the destination before assuming +the grasp failed. Wrong-surface placement is a common cause. + +## Completion And Audit + +The current outcome is the top-level `libero_terminated` value returned by the +latest environment tool. Do not look for it inside `state`. + +On completion, write the requested audit with: + +- suite, task, seed, and evaluation regime; +- exact memory files consulted; +- semantic identification and localization strategy; +- grasp and placement evidence; +- final state and `libero_terminated`; +- honest failure details when the task remains incomplete. + +Then call `finish`. Recipe generation is handled by the runtime from recorded +primitive and successful segmentation events. + +## Final Checklist -### How to use `segment` (SAM3) — practical tips - -`segment` is the fastest way to localize: one call gives you a `world_xyz` -that is a robust MEDIAN over the whole object mask (hundreds of pixels), which -beats eyeballing 3–5 pixels. Workflow: - -1. Read `image_cam_hi_NN.png`, decide the target by the task's spatial RELATION. -2. Call `segment` with a **plain visual phrase + the relation**, then read the - returned `world_xyz` and the `segments/segment_overlay_NN_XX.png` to confirm the mask - is on the right object before you move. -3. **Two camera views — YOUR choice via the `"camera"` field** (`segment` works - on either; default `"agentview"`): - - `"camera":"agentview"` — the fixed ~1 m cam (±8–13 cm). Use it to choose - WHICH object (global layout / spatial relation) and get a rough xy. - - `"camera":"wrist"` — the eye-in-hand cam (±1–2 cm), reads the wrist world - map in the SAME world frame. Use it to PRECISELY localize once the gripper - is parked over the target. ⚠ It is null / all-table until you `move_to` - ~15–20 cm over the target (the wrist only sees the gripper+table before - that), so segment agentview first, approach, THEN segment wrist. - Typical coarse→fine: agentview select → `move_to` above → wrist refine → grasp. - You decide when the wrist refinement is worth it (small / closely-spaced / - identical objects benefit most; a big isolated plate may not need it). - -**Prompt phrasing (important — SAM3 is sensitive):** -- ✅ Use plain colour + shape + spatial relation: `"the black bowl on the stove"`, - `"the white plate"`, `"the bowl on the cookies box"`. -- ❌ Do NOT use the object's internal/brand NAME from `object_names`/BDDL: - `"the akita black bowl"` scores ~0.03 (no detection) because SAM3 can't ground - "akita"; the same object as `"the black bowl on the stove"` scores ~0.76. Strip - proper nouns (`akita`, `glazed_rim_porcelain_…`) — say what it *looks like*. -- For two identical objects, the relation in the prompt (`…on the stove`, - `…left of the plate`) usually steers SAM3 to the right instance; verify via the - overlay and the world-z (an object *on* a fixture has a higher z than a - table-level twin). The two-camera relation protocol still applies when SAM3 - can't disambiguate from text alone. -- If `segment` returns `{"error":...}` (low score / service down), walk the - prompt (drop the brand word, simplify, add the relation) or fall back to - manual pixel → `back_project`. - -## The strict-hybrid recipe (perception variant) - -A typical bowl→plate cell looks like: - -1. `view_driver_state({"step":0})`; inspect `task_language`, - `image_cam_hi_00.png`, and `camera_meta.json`. -2. Identify the target by the task's spatial RELATION, not its name (Rule 7). -3. Localize it — pick its pixel in `image_cam_hi_00.png`, - `back_project({"row":ROW,"col":COL,"step":0})` (coarse). Refine with the - wrist cam after pre-positioning (fine). -4. **Pre-position** ~15–20 cm above it: `move_to([obj_x, obj_y, carry_z])`, - gripper open. Re-`back_project` a wrist pixel here and correct the xy before - grasping. -5. `pi0_pick` with the right prompt (start "pick up the {object}", escalate per - the Pi0 ladder — see below). Confirm the grasp via gripper + wrist (Rule 1b). -6. `set_gripper({"gripper":1,"steps":5})` to firm the grip. -7. Localize the placement region (basket / plate / drawer slot) the same way. -8. `move_to([place_x, place_y, carry_z])` to traverse at constant height. -9. Optionally descend `move_to([place_x, place_y, place_z])`. -10. `release` — predicate (`On`/`In`) checks → `libero_terminated=True` if hit. -11. Light retreat (`move_to` upward) so the next step's image is clean. - -> **Predicate fire timing.** Most LIBERO `On(X, Y)` predicates fire on -> `release` if `X` is above `Y`'s region. `In(X, container)` needs `X` to -> have actually entered the container's volume before release. If a release -> fires `term=False` but the object is on top of the right region, descend -> 1–2 cm more and re-`release`. - -## Pi0 prompt ladder (Rule 3 — Pi0 IS the delivery service) - -Try in order; each rung uses a slightly more specific prompt or a re-pre-pos. - -1. `"pick up the {object}"` — generic sub-instruction. -2. The full `task_language` verbatim (e.g. `"Pick the akita black bowl on the - cookies box and place it on the plate"`). -3. Add a spatial qualifier (`"…on the wooden cabinet"`, `"…next to the - basket"`). -4. Re-pre-position 5 cm lower or shifted, then retry rung 2 or 3. - -Empirically Pi0 sometimes needs the **full prompt with the spatial -qualifier** for elevated picks (stove, cabinet-top, drawer). See -`feedback_pi0_pick_full_prompt.md` and the prompt-ladder note in MEMORY. - -## Key hyperparameters - -- Single-step `xyz` within ±0.30 m of current eef or OSC flips IK. Split - long traversals into 2–3 carry-z waypoints. -- `lift_thresh`: 0.05 (flat/stable) / 0.08 (slippery tall bottles). -- `step_clip`: 0.025 (empty / box) / 0.015 (cans) / 0.012 (tall bottles). -- Frame z (from `state.robot0_eef_pos[2]` at step 00): - ≈ 0.68 → LIVING_ROOM, ≈ 1.17 → KITCHEN, ≈ 0.26 → OBJECT. -- BOWL: `eef_y_target = perceived_plate_y + 0.045` (bowl-eef y-offset). -- TALL BOTTLES: carry at `z=0.30`, release without descending. -- Approach high-then-vertical; recover by re-`pi0_pick`, not by hovering. - -## Reading state - -After every primitive tool call, the return value already carries the new -`state`, the merged `command`/`result`/`elapsed_s` log, and the image paths — -you do **not** need a separate read. When you need an older step, call -`view_driver_state({"step":NN})` (omit `step` for the latest): - -1. Check `result.success`, `final_dist_m`, `peak_lift_m`, etc. in the returned - log. -2. Read the returned `image_cam_hi_NN.png` path → visual confirmation; pick - pixels for any new localization. -3. Check `state.robot0_eef_pos` / `robot0_gripper_qpos` / `libero_terminated` - in the returned `state`. - -You **do not** open the depth maps yourself — feed a pixel to `back_project`. -Don't call `view_driver_state` immediately after a primitive that already -returned the new state. - -## Common failure modes - -- **Pick missed (gripper closed empty).** `result.peak_lift_m` < `lift_thresh`, - `min_gripper_opening` ≈ 0. Re-pre-position 1–2 cm lower or shifted; retry - Pi0 with next prompt-ladder rung. -- **Object slipped mid-carry.** `release` returns `term=False` and the object - is no longer where you `release`d it. `release`, re-pre-pos above it, - `pi0_pick` again, traverse again. -- **OSC stuck.** `move_to` returns `final_dist_m > 0.05` at `max_steps` and - same xy twice → try `rotate_pitch`, split into more waypoints, or - `move_pose` (co-varying) for the cabinet-front singularity. -- **Placement off because localization was wrong.** The release puts the - object on bare table instead of on the plate. Re-read `image_cam_hi_NN.png`, - pick a different pixel on the target region (sample 3 pixels on the plate's - flat top, median the back-projected xy), redo. Tip: if depth at your chosen - pixel is much closer than the table z, you picked the camera's near edge / - an object rim — pick again. - -## Verifying strict compliance - -Before saving the audit, confirm your command history is physics-only. The -teleport primitives are not even in the tool list, but audit it anyway: read -the `command` field of every `states.json` entry (via -`view_driver_state({"step":NN})` per step, or `read_text_file` on -`{output_dir}/states.json`) and confirm every `command.action` is one of the -allowed physics primitives (`move_to`, `pi0_pick`, `pi0_doubled`, `release`, -`set_gripper`, `rotate_wrist`, `rotate_pitch`, `move_pose`) — no -`set_object_pose` / `articulate_to` / `js_move_to` / `carry_object` appears. - -## Persisting successful runs as audit JSONs - -When `state.libero_terminated == true`: - -a. The working command recipe (`{output_dir}/recipe_{recipe_tag}.jsonl`) is - **auto-exported by the runner** from the non-error primitive commands in - `states.json` plus successful segment calls recorded in - `segments/segment_*.json`, merged in execution order — you do NOT - hand-write it. -b. Write a minimal audit JSON with `write_text_file` to - `{output_dir}/{recipe_tag}.json` with at least: `suite`, `task_id`, `seed`, - `regime: "strict_perception"`, `strategy_notes` (mention HOW you localized — - which pixel, depth, back-projected world xyz), `pick_result` (the `result` - from your `pi0_pick`), `final_state` (the latest `states.json` entry's - `state` field), `libero_terminated: true`. -c. Call `finish({"status":"success","summary":"…"})`. - -If unrecoverable after honest exploration in this one episode, write -`{output_dir}/{recipe_tag}.json` with `libero_terminated: false` + -`strategy_notes` describing what you tried, the back-projected xyz you used, and -which step failed. Then call `finish` (NO reset, NO second attempt). - -> The `regime: "strict_perception"` tag distinguishes these audits from the -> oracle-state `strict` regime in mixed-mode datasets. - -## Iteration heuristics - -- After 2 failed retries on the same step, **stop tuning numerics** and - inspect images: read `image_cam_hi_NN.png` at pick / pre-release / - post-release. The visual disagreement is usually the bug (you picked the - wrong pixel, or Pi0 grabbed the decoy). This is the lesson of - `feedback_failure_forensics.md` — applies even more strongly here. -- If a release reports `term=False` but the object is visibly on the target - region, descend 1–2 cm more and `release` again; predicate often needs - contact, not just hover. -- Once you use `back_project`, trust it: if a localization "feels off", it's - because you picked the wrong pixel — not because the depth / calibration is - wrong. (`back_project` inverts `K` before the extrinsic; the raw - `E @ [col·z, row·z, z, 1]` form skips it and is wrong — which is why you use - `back_project` and never hand-roll the math.) - -## What "strict_perception" means concretely - -- **No GT object coordinates anywhere in your reasoning.** The `state` you - read has none; the only legitimate sources of object xyz are the camera - images + depth + the precomputed `world/` and `world_wrist/` maps (via - `back_project`). -- **No teleport primitives.** The four are deleted. -- **Pi0 only does the grasp.** You script every motion + release. -- **Fully oracle-free, including the grasp.** There is no GT-lift oracle — - `pi0_pick` reads NO GT object pose and takes no tracking argument. You judge - the grasp from gripper width + the wrist cam, and TASK success from - `state.libero_terminated` (the benchmark predicate). -- **Single attempt.** One episode, no reset (Rule 2). -- The expected audit `regime` is `strict_perception`. - -## Reference cases - -The seed-0 sweep results live under `resources/libero/results_*_pert/` -(`results_10_pert`, `results_object_pert`, `results_spatial_pert`, -`results_goal_pert` — PRO swap+task, t0–t9). Each solved cell has an audit JSON -+ a `recipe_{tag}.jsonl` command sequence. The consistent winning pattern: -localize → pre-pos → `pi0_pick` → `set_gripper` → move → `release` in 6–12 -commands. - -When you write a new audit, browse a sibling cell's `recipe_{tag}.jsonl` as a -*technique* template — but never paste its xyz; re-derive every position via -`back_project` from THIS scene's depth. - -Begin by reading `resources/libero/memory/MEMORY.md`, then call -`view_driver_state({"step":0})` and inspect the returned `image_cam_hi_00.png` -(+ `camera_meta.json` via `view_camera_meta`); localize the target object via -`back_project`, then plan and execute. +- Initial state read with `step: 0`. +- Task language copied from the returned tool result. +- All targets and destinations semantically identified in agentview. +- Coordinates obtained through `back_project` or `segment`. +- Wrist refinements checked against the agentview anchor. +- Grasp confirmed visually and proprioceptively. +- Long motion split into safe waypoints. +- Destination reclassified immediately before release. +- Latest result checked for top-level `libero_terminated`. +- Audit written before `finish`. diff --git a/robots/libero/prompts/system.py b/robots/libero/prompts/system.py index 16bd1e93..d93b8df8 100644 --- a/robots/libero/prompts/system.py +++ b/robots/libero/prompts/system.py @@ -10,7 +10,7 @@ > persistence / up to N attempts" instruction anywhere below).** This is a > ONE-SHOT evaluation: you get **exactly ONE episode**. You MUST NOT call > `reset`, and you must not restart the episode. Plan carefully, then execute -> your single best manipulation sequence toward `state.libero_terminated == true`. +> your single best manipulation sequence toward `libero_terminated == true`. > You MAY recover *within* this one episode (re-pre-position, re-`pi0_pick` a > missed grasp, walk the Pi0 prompt ladder, `rotate_pitch`/`move_pose`) — that is > all one continuous attempt — but the instant you would want to reset/start over, @@ -70,7 +70,7 @@ weak at reading side labels or distinguishing similar grocery items (ketchup/BBQ/tomato sauce, soup cans, cream cheese/butter). Do NOT let the wrist freely re-identify a non-basket target; it often locks onto a look-alike. - Instead: choose the target from `image_cam_hi_NN.png`, compute its agentview + Instead: choose the target from the high-resolution agentview image, compute its xyz, move over that candidate, project/track that SAME candidate in wrist, and refine only its surface/center coordinates. SAM3 scores ~0.02-0.06 on brand nouns ("alphabet soup", "tomato sauce") — prompt by colour+shape ("the short @@ -141,63 +141,33 @@ - Call the real structured tools exposed by the runtime. - Use bare tool names in this prompt: `move_to`, `pi0_pick`, `release`, `set_gripper`, `rotate_wrist`, `rotate_pitch`, `move_pose`, `pi0_doubled`, - `view_driver_state`, `view_camera_meta`, `back_project`, `segment`, + `view_env_state`, `view_camera_meta`, `back_project`, `segment`, `read_text_file`, `write_text_file`, `list_dir`, `finish`. - Under some runtimes these same tools may appear namespaced; call the actual tool name shown in your tool list, preserving the same arguments and semantics. -The toolkit writes artifacts in `{{output_dir}}/`: - -- `{{output_dir}}/states.json` — top-level JSON array; each entry has - `step_idx`, `task_language`, `libero_terminated`, `state` (robot - proprioception + object_names; NO object coordinates), `command`, `result`, - `elapsed_s`, and world-map path fields when available. -- `{{output_dir}}/images/image_NN.png` — agentview RGB, 180°-rotated (Pi0 frame; - do NOT use for back-projection). -- `{{output_dir}}/images_cam/image_cam_NN.png` — agentview RGB in the CALIBRATION - frame; use for low-resolution pixel checks. -- `{{output_dir}}/depths/depth_NN.npy` — agentview metric depth (meters), - calibration frame. -- `{{output_dir}}/world/world_NN.npy` — HxWx3 precomputed world xyz per 256px - agentview pixel. Prefer `back_project`; read this manually only for debugging - or if the tool is unavailable. -- `{{output_dir}}/images_wrist/image_wrist_NN.png` — wrist RGB, calibration frame. -- `{{output_dir}}/depths_wrist/depth_wrist_NN.npy` — wrist metric depth (meters). -- `{{output_dir}}/world_wrist/world_wrist_NN.npy` — wrist world xyz map in the - SAME world frame as agentview. -- `{{output_dir}}/wrist_meta/wrist_meta_NN.json` — wrist intrinsics + extrinsic - FOR THAT STEP ONLY (the wrist cam moves, so it changes every step). -- `{{output_dir}}/images_cam_hi/image_cam_hi_NN.png` — HI-RES (1024x1024) - agentview RGB in calibration frame. USE THIS to inspect the scene and identify - objects — a far object spans 4x more pixels than at 256. -- `{{output_dir}}/world_hi/world_hi_NN.npy` — 1024x1024x3 float16 precomputed - world xyz per hi-res agentview pixel. Prefer `back_project`; if you manually - inspect it, never index a low-res pixel into this grid or vice versa. -- `{{output_dir}}/images_wrist_hi/image_wrist_hi_NN.png` / - `{{output_dir}}/world_wrist_hi/world_wrist_hi_NN.npy` — same hi-res pair for - the WRIST cam. - ⚠ Hi-res pixel (row,col) indexes ONLY the hi-res world map (and 256 pixel -> - 256 map). Don't mix grids; if you must convert, divide hi coords by 4. - ⚠ Hi-res files keep only the LAST 5 STEPS (disk); for older before/after - comparisons use the 256 files or `states.json` history. -- `{{output_dir}}/camera_meta.json` — agentview intrinsics K, cam->world - extrinsic, projection recipe. -- `{{output_dir}}/action_videos/step_NN_.mp4` — per-action clips generated - when the Dashboard is enabled. - -NN is zero-padded sequential (`00`, `01`, `02`, ...). Initial state step `00` is -dumped before you begin. Use `view_driver_state({"step": 0})` to read it.""" - -GOAL = """YOUR GOAL: produce `state.libero_terminated == true` in ONE episode. ⛔ NO +The driver records a state and an `observation` dictionary for every motion +step. Observation entries name the available policy, agentview, wrist, depth, +world-map, and metadata artifacts, but storage paths are internal to the +runtime. Do not construct or read artifact paths manually. + +Use `view_env_state` to retrieve a state. It embeds the policy image and the +best available agentview and wrist images, preferring high resolution. Use +`view_camera_meta`, `back_project`, and `segment` to consume metadata, depth, +and world maps. These tools guarantee that the selected camera, resolution, +and step use matching artifacts. + +Step `0` is the initial state. Step `-1` selects the latest state.""" + +GOAL = """YOUR GOAL: produce top-level `libero_terminated == true` in ONE episode. ⛔ NO `reset`, NO retry (SINGLE-ATTEMPT MODE — see the override at the very top; it supersedes any reset/retry wording in the Rules below).""" RULES = """Rule 0 — USE IMAGES. After every primitive tool call, inspect the returned state - and image paths. If you need a state again, call `view_driver_state`. Read the - new `image_cam_hi_NN.png` path (calibration frame — the one you pick pixels in) - and, when close to a target, the `image_wrist_hi_NN.png` path. The image is - your spatial-reasoning input; `states.json` only gives proprioception + object - names. + and embedded images. If you need a state again, call `view_env_state`. + Use agentview for global layout and wrist for close-range geometry. The image + is your spatial-reasoning input; the JSON state only gives proprioception + + object names. Rule 1 — Pi0 is ONLY for the grasp. Use: pi0_pick({ @@ -215,19 +185,19 @@ Rule 1b — JUDGE THE GRASP from perception, NOT from a name. After a pick, decide "did I grab the target?" from two coord-free signals: • GRIPPER (proprioception): `state.robot0_gripper_qpos` from the latest - `states.json` entry — fingers closed but NOT fully shut (~0.01–0.05 gap) + state record — fingers closed but NOT fully shut (~0.01–0.05 gap) ⇒ holding an object; fully closed (~0.0) ⇒ grasped air. - • WRIST CAM: Read `image_wrist_hi_NN.png` after lifting. The target should + • WRIST CAM: inspect the returned wrist image after lifting. The target should now be raised into the gripper, and the spot it came from should be EMPTY. Compare before/after wrist or agentview evidence; if needed, use `back_project` on wrist pixels to confirm the target surface z jumped up. `pi0_pick.success` (eef-lift + gripper-closure heuristic) is a HINT, not proof — always confirm with the wrist cam before carrying. -Rule 2 — Inspect THEN act. Call `view_driver_state({"step": 0})`, read the - returned high-resolution image path(s), and inspect the relevant memory/guides - BEFORE your first primitive. **Your task is `states.json[0]["task_language"]` - — read it and obey it verbatim.** This is the authoritative instruction (the BDDL +Rule 2 — Inspect THEN act. Call `view_env_state({"step": 0})`, inspect the + returned high-resolution images, and inspect the relevant memory/guides + BEFORE your first primitive. **Your task is the returned `task_language`; + read it and obey it verbatim.** This is the authoritative instruction (the BDDL `:language` tag). Do NOT infer the task from object names, from sibling recipes, or by guessing a task_map index — those caused wrong-task runs in the past. @@ -257,7 +227,7 @@ plate, a stove burner/cook-region, a wooden-cabinet top, and a pot lid all read as "flat disc at table height" in back-projected coordinates. They are only separable in the RGB. So before you carry-and-release onto a surface, look at - `image_cam_hi_NN.png` (and the wrist `image_wrist_hi_NN.png` once close) and + returned agentview image (and the wrist image once close) and NAME each candidate surface: • PLATE ⇒ ceramic disc, usually white, with a clean raised rim (often colored concentric rings). This is the place target for "place it on the plate". @@ -292,8 +262,7 @@ LOCALIZATION = """This is the core of perception-isolated mode. To find where an object is: -1. Look at `image_cam_hi_NN.png` (1024x1024 — PREFER THIS; fall back to the - 256 `image_cam_NN.png` only if the hi file is absent) and find the target +1. Look at the returned agentview image (high resolution when available) and find the target object's pixel (row, col). (row = vertical/y from top, col = horizontal/x from left.) 2. Call `back_project` on that pixel: @@ -316,7 +285,7 @@ metres away. Pick pixels firmly on the object's top surface.) ALWAYS apply the manipulation offsets from memory to the PERCEIVED position -(e.g. BOWL: eef_y = plate_y + 0.045). Verify visually in image_cam after moving.""" +(e.g. BOWL: eef_y = plate_y + 0.045). Verify visually in agentview after moving.""" PERCEPTION_ALGORITHM = """This is the default perception algorithm for EVERY cell (from the 80-task localization sweep: `agentview_identity_wrist_geometry_except_basket`). @@ -335,11 +304,11 @@ ALGORITHM (run this BEFORE manipulating): -1. From `states.json[0]["task_language"]` + `image_cam_hi_00.png` + +1. From the initial `task_language` + returned agentview image + object_names, infer the task-relevant TARGETS and DESTINATIONS (language only; never BDDL/poses). -2. GLOBAL SEMANTIC PASS (agentview hi-res): in `image_cam_hi_NN.png` choose each +2. GLOBAL SEMANTIC PASS (agentview hi-res): in the returned agentview image choose each target/destination candidate by RGB, label/shape, and global spatial relation. For duplicates (two bowls/plates/mugs) pick by RELATION (on stove, on cookie box, left/right/front/back), not `_1/_2`. For sauce/can/box groceries use the @@ -348,7 +317,7 @@ burner vs cabinet/drawer vs basket) semantically in RGB here. 3. COARSE XYZ (agentview): pick 3-8 pixels firmly on the chosen candidate in - `image_cam_hi_NN.png`, call `back_project` on the SAME pixels, take the median. + agentview image, call `back_project` on the SAME pixels, take the median. Avoid edges/holes/shadows/table-gaps. This median is the IDENTITY ANCHOR for that entity. @@ -430,9 +399,9 @@ with older/oracle assumptions; do NOT copy coordinates and do NOT replay stale command lists. Re-derive every coordinate from THIS scene. """, - """INSPECT INITIAL STATE: call `view_driver_state({"step": 0})`; inspect -`task_language`, object_names, eef pose, `image_cam_hi_00.png`, -`image_wrist_hi_00.png` if useful, and `camera_meta.json`. Identify ALL target + """INSPECT INITIAL STATE: call `view_env_state({"step": 0})`; inspect + `task_language`, object_names, eef pose, the returned agentview and wrist images, + and call `view_camera_meta` if needed. Identify ALL target objects, destination surfaces, and relation landmarks named by task_language. """, """RUN THE MANDATORY PRE-TASK PERCEPTION PASS (FIRST-STEP ALGORITHM above) — @@ -450,10 +419,9 @@ pi0_pick({"prompt": "...", "max_chunks": 20, ...}) release({}) -Each primitive tool blocks until the next `states.json` entry is dumped and -returns the new state view + log + image paths. Then inspect the returned state -+ high-resolution image paths (+ `back_project` as needed), decide, repeat -with NN=02, 03, ... +Each primitive tool blocks until the next state record is dumped and returns the +new state view, log, and embedded images. Inspect them, use `back_project` as +needed, decide, and repeat. """, """ALLOWED PRIMITIVES (physics-only; full schemas in the tool list/guides): `move_to`, `pi0_pick`, `pi0_doubled`, `release`, `set_gripper`, @@ -475,16 +443,17 @@ SAM3 localization aid — `segment` (no robot motion): instead of eyeballing a pixel, call `segment({"prompt":"the black bowl on the cookies box", "camera":"agentview"})`. It runs SAM3 on the current image, back-projects the -mask via the matching world map, and writes `segments/segment_NN_XX.json` with a robust -median `world_xyz` (+ a `segments/segment_overlay_NN_XX.png` to confirm the right -object). Use `camera":"wrist"` (after parking the eef ~15–20 cm over the +mask via the matching world map, and returns a robust median `world_xyz` plus +an embedded overlay image for visual confirmation. The logical +`segment_artifact` and `overlay_artifact` names are audit references, not paths +to open manually. Use `camera":"wrist"` (after parking the eef ~15–20 cm over the target) for ±1–2 cm refinement, or `"point":[row,col]` for a point prompt. Text `prompt` and `point` are mutually exclusive; provide exactly one. ⚠ PROMPT PHRASING (SAM3 is sensitive): use a plain colour+shape+RELATION phrase, NEVER the internal/brand name from `object_names`/BDDL. `"the akita black bowl"` scores ~0.03 (SAM3 can't ground "akita") whereas `"the black bowl on the stove"` scores ~0.76. Strip proper nouns (akita, glazed_rim_porcelain_…) — say what it -LOOKS LIKE + where it is. Always inspect the returned overlay path to confirm +LOOKS LIKE + where it is. Always inspect the returned overlay image to confirm the mask landed on the right object before moving. This is a CONVENIENCE alternative to manual back-projection — if it returns `{"error":..., "fallback":...}` (server down / low score / no detection), walk @@ -499,7 +468,7 @@ unrecoverable within this one episode, do NOT reset — write an honest stuck-audit (`libero_terminated:false`) and call `finish`. Never warp. """, - """WHEN state.libero_terminated == True: + """WHEN `libero_terminated == true` in the latest tool result: a. Write audit `{{output_dir}}/{{recipe_tag}}.json` with: suite, task_id, seed, regime:"strict_perception", strategy_notes (incl. how you localized), pick_result, final_state (latest state's `state`), @@ -519,7 +488,7 @@ OUTPUT_DISCIPLINE = """- Brief reasoning before each tool call (1-2 sentences): observation → decision. - Don't re-read files already in this session. -- Don't call `view_driver_state` immediately after a primitive tool already +- Don't call `view_env_state` immediately after a primitive tool already returned the new state. - Save the audit BEFORE calling `finish`. - Stop immediately after writing the audit and calling `finish`. Do not chat further.""" diff --git a/robots/libero/prompts/user.py b/robots/libero/prompts/user.py index 668a4db4..a71b2444 100644 --- a/robots/libero/prompts/user.py +++ b/robots/libero/prompts/user.py @@ -10,9 +10,11 @@ - recipe: {{output_dir}}/recipe_{{recipe_tag}}.jsonl""" -MODE = """Use the high-resolution image paths returned by view_driver_state and -back_project to localize objects before motion.""" +MODE = """Inspect the embedded high-resolution images returned by +view_env_state, then use back_project or segment to localize objects before +motion.""" -BEGIN = """read MEMORY.md, the guides, then `view_driver_state({"step":0})` and the -returned high-resolution images. Localize the target, then plan and execute.""" +BEGIN = """Read MEMORY.md and the guides, then call +`view_env_state({"step": 0})` and inspect its embedded images. Localize every +task-relevant entity before planning and execution.""" diff --git a/robots/libero/toolkit.py b/robots/libero/toolkit.py index 475c1e74..23a315cb 100644 --- a/robots/libero/toolkit.py +++ b/robots/libero/toolkit.py @@ -5,34 +5,32 @@ """ from __future__ import annotations -import shutil -import time from functools import partial from typing import Any from robots.libero import tools as libero_tools -from rpent.dashboard.events import DashboardEventSink, ToolResultEvent -from rpent.tools.toolkit import ToolCancelled, Toolkit +from rpent.dashboard.events import DashboardEventSink +from rpent.tools.state import EnvState +from rpent.tools.toolkit import Toolkit from rpent.utils.logging import get_logger, get_output_dir class LiberoToolkit(Toolkit): """Toolkit for the LIBERO environment.""" - # Tool schemas keyed by name (built once from the canonical ordered list - # in libero_tools.TOOLS_SPEC) so each tool registers with its own spec. - _SPECS = {spec["name"]: spec for spec in libero_tools.TOOLS_SPEC} + _FRAME_ARTIFACTS = { + "camera": "agentview.png", + "wrist": "wrist.png", + } def __init__( self, *, primitives_kwargs: dict[str, Any], dashboard_events: DashboardEventSink, - video_path: str | None = None, ) -> None: - super().__init__(dashboard_events=dashboard_events) - self._next_step: int = 0 - self._video_path: str | None = video_path + state = EnvState(get_output_dir()) + super().__init__(dashboard_events=dashboard_events, state=state) self.init_primitives_clean(primitives_kwargs=primitives_kwargs) self._register_libero_tools() @@ -40,67 +38,66 @@ def __init__( # Registration # ------------------------------------------------------------------ def _register_libero_tools(self) -> None: - specs = self._SPECS - # Inspection tools do not advance environment state. Most are stateless - # module functions; segment is bound to the primitives-owned SAM3 client. - inspection_handlers = { - "view_driver_state": libero_tools.view_driver_state, - "view_camera_meta": libero_tools.view_camera_meta, - "back_project": libero_tools.back_project, - "segment": self._primitives.segment, + # Read-only tools whose handlers aren't primitive methods (they need + # the run's EnvState bound in, or -- like segment -- must stay + # read-only despite being a primitives method). Every other spec binds + # to its primitive-driver method; @updatestate on the method decides + # whether state is captured. + state_handlers = { + "view_env_state": partial( + libero_tools.view_env_state, state=self._state + ), + "view_camera_meta": partial( + libero_tools.view_camera_meta, state=self._state + ), + "back_project": partial(libero_tools.back_project, state=self._state), + "segment": partial(self._primitives.segment, state=self._state), } - for name, handler in inspection_handlers.items(): - self.add_tool(name, specs[name], handler) - # Primitive tools: each goes through _step, which looks up the - # matching primitive method via getattr at call time. - for name in libero_tools.PRIMITIVE_TOOL_NAMES: - self.add_tool(name, specs[name], partial(self._step, name)) - - def _step(self, name: str, **kwargs) -> dict: - """Run ``self._primitives.(**kwargs)``, dump the new step, and - return the rendered state view + log. - """ - command = {"action": name, **kwargs} - t0 = time.time() - start_frame = self._primitives.recorded_frame_count() - try: - result = getattr(self._primitives, name)(**kwargs) - self.raise_if_cancelled() - except ToolCancelled as exc: - result = { - "error": str(exc), - "code": "tool_cancelled", - "interrupted": True, - } - elapsed = round(time.time() - t0, 2) - - if isinstance(result, dict): - result_dict = result - else: - result_dict = {"value": result} - - self._next_step += 1 - step_idx = self._next_step - output_dir = get_output_dir() + for spec in libero_tools.TOOLS_SPEC: + name = spec["name"] + if name in state_handlers: + handler = state_handlers[name] + else: + handler = getattr(self._primitives, name, None) + if handler is None: + continue # spec without a backing primitive method + self.add_tool(name, spec, handler) + + def get_env_state( + self, + *, + command: dict[str, Any], + result: dict[str, Any], + elapsed_s: float, + ) -> dict[str, Any]: + frame_start = self._action_frame_cursor + self._action_frame_cursor = self._primitives.recorded_frame_count() + record = libero_tools.dump_state( + self._primitives, + self._state, + log={"command": command, "result": result, "elapsed_s": elapsed_s}, + ) if self._dashboard_events.enabled: - video_dir = libero_tools.artifact_path(output_dir, "action_videos") - video_path = video_dir / f"step_{step_idx:02d}_{name}.mp4" try: - self._primitives.save_frame_slice(start_frame, str(video_path), fps=20) + frames = self._primitives.frame_slice(frame_start) + if frames: + candidate = f"action_{command['action']}.mp4" + self._state.save( + candidate, + frames, + step=record.step_idx, + fps=20, + ) except Exception as e: get_logger("libero_toolkit").warning( - f"failed to save action clip to {video_path}: {e}" + "failed to save action clip for step %s: %s", + record.step_idx, + e, ) - libero_tools.dump_state( - self._primitives, - str(output_dir), - step_idx=step_idx, - log={"command": command, "result": result_dict, "elapsed_s": elapsed}, - ) - out = libero_tools.view_driver_state(step_idx) - out["agent_elapsed_s"] = elapsed - if result_dict.get("interrupted"): - out.update(result_dict) + out = libero_tools.view_env_state(record.step_idx, state=self._state) + out["agent_elapsed_s"] = elapsed_s + if result.get("interrupted"): + out.update(result) return out def init_primitives_clean( @@ -109,19 +106,7 @@ def init_primitives_clean( primitives_kwargs: dict[str, Any], ) -> None: """Wipe stale run artifacts, build the LiberoPrimitives, dump step 0.""" - out_dir = get_output_dir() - out_dir.mkdir(parents=True, exist_ok=True) - for sub in libero_tools.ARTIFACT_DIRECTORIES: - target = out_dir / sub - if target.exists(): - shutil.rmtree(target) - for target in ( - libero_tools.artifact_path(out_dir, "states"), - libero_tools.artifact_path(out_dir, "metadata", camera="agentview", resolution="low"), - libero_tools.artifact_path(out_dir, "episode_video"), - ): - if target.exists(): - target.unlink() + self._state.reset() primitives = libero_tools.LiberoPrimitives( check_cancelled=self.raise_if_cancelled, @@ -129,30 +114,22 @@ def init_primitives_clean( ) primitives.reset() primitives.start_recording() - libero_tools.dump_state(primitives, str(out_dir), step_idx=0, log=None) - self._dashboard_events.emit( - ToolResultEvent( - name="view_driver_state", - result=libero_tools.view_driver_state(0), - ) - ) - + self._action_frame_cursor = primitives.recorded_frame_count() + record = libero_tools.dump_state(primitives, self._state, log=None) self._primitives = primitives + self._publish_step(record) def close(self) -> None: - """Flush the agent-side video buffer to disk (end-of-run). - """ - if self._video_path is None: - return + """Flush the agent-side video buffer through ``EnvState``.""" try: - self._primitives.stop_recording_and_save(self._video_path) + frames = self._primitives.stop_recording() + if frames: + self._state.save("episode.mp4", frames, step=None, fps=20) except Exception as e: - # The runner is in the cleanup path; never let a video save - # abort it. get_logger("libero_toolkit").warning( - f"failed to save video to {self._video_path}: {e}" + f"failed to save episode video: {e}" ) def write_recipe(self, recipe_tag: str) -> str: """Write the LIBERO recipe JSONL from the dumped state trace.""" - return libero_tools.write_recipe_from_states(str(get_output_dir()), recipe_tag) + return libero_tools.write_recipe_from_states(self._state, recipe_tag) diff --git a/robots/libero/tools.py b/robots/libero/tools.py index 44145227..2f66eb76 100644 --- a/robots/libero/tools.py +++ b/robots/libero/tools.py @@ -1,86 +1,20 @@ """LIBERO + OpenPI tool implementation.""" from __future__ import annotations -import json -import os from collections.abc import Callable -from pathlib import Path from typing import Any -import imageio.v2 as imageio import numpy as np from robots.libero.env_client import LiberoEnvClient -from rpent.utils.logging import get_logger, get_output_dir +from rpent.tools.state import EnvState, StepRecord +from rpent.tools.toolkit import updatestate +from rpent.utils.logging import get_logger from rpent.utils.sam3_client import Sam3Client from rpent.utils.vla_client import VLAClient logger = get_logger("libero") -ARTIFACT_LAYOUT: dict[tuple[str | None, str | None, str], str] = { - ("agentview", "low", "policy_image"): "images/image_{step:02d}.png", - ("agentview", "low", "image"): "images_cam/image_cam_{step:02d}.png", - ("agentview", "low", "depth"): "depths/depth_{step:02d}.npy", - ("agentview", "low", "world"): "world/world_{step:02d}.npy", - ("agentview", "low", "metadata"): "camera_meta.json", - ("agentview", "high", "image"): "images_cam_hi/image_cam_hi_{step:02d}.png", - ("agentview", "high", "world"): "world_hi/world_hi_{step:02d}.npy", - ("wrist", "low", "image"): "images_wrist/image_wrist_{step:02d}.png", - ("wrist", "low", "depth"): "depths_wrist/depth_wrist_{step:02d}.npy", - ("wrist", "low", "world"): "world_wrist/world_wrist_{step:02d}.npy", - ("wrist", "low", "metadata"): "wrist_meta/wrist_meta_{step:02d}.json", - ("wrist", "high", "image"): "images_wrist_hi/image_wrist_hi_{step:02d}.png", - ("wrist", "high", "world"): "world_wrist_hi/world_wrist_hi_{step:02d}.npy", - (None, None, "states"): "states.json", - (None, None, "episode_video"): "episode.mp4", - (None, None, "segments"): "segments", - (None, None, "action_videos"): "action_videos", -} -ARTIFACT_DIRECTORIES: tuple[str, ...] = ( - "images", - "images_cam", - "depths", - "world", - "images_cam_hi", - "world_hi", - "images_wrist", - "depths_wrist", - "world_wrist", - "wrist_meta", - "images_wrist_hi", - "world_wrist_hi", - "segments", - "action_videos", -) - - -def artifact_path( - output_dir: str | os.PathLike[str], - kind: str, - *, - step: int | None = None, - camera: str | None = None, - resolution: str | None = None, -) -> Path: - """Resolve one artifact path from the shared layout. - - ``kind`` identifies the artifact; the remaining fields are optional - qualifiers and must be passed by keyword to avoid mixing them up. - """ - return Path(output_dir) / _artifact_relative_path(step, camera, resolution, kind) - - -def _artifact_relative_path( - step: int | None, - camera: str | None, - resolution: str | None, - kind: str, -) -> str: - template = ARTIFACT_LAYOUT[(camera, resolution, kind)] - if step is None and "{step" in template: - raise ValueError(f"{kind} artifact requires a step") - return template.format(step=step) - def _normalize_xyz(xyz): """Coerce an LLM-supplied xyz into a length-3 list[float].""" @@ -116,7 +50,7 @@ def __init__( self._last_obs_eef_z = None self._last_obs_gripper = None # Per-env-step frame buffer for diagnostic video rendering. - # Toggled via start_recording() / stop_recording_and_save(). + # Toggled via start_recording() / stop_recording(). self._recording = False self._frames = [] @@ -131,22 +65,14 @@ def record_frame(self, obs): def recorded_frame_count(self) -> int: return len(self._frames) - def stop_recording_and_save(self, path: str, fps: int = 20): - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - n = len(self._frames) - if n > 0: - imageio.mimwrite(path, self._frames, fps=fps) + def stop_recording(self) -> list[np.ndarray]: + frames = list(self._frames) self._recording = False self._frames = [] - return {"path": path, "n_frames": n} + return frames - def save_frame_slice(self, start: int, path: str, fps: int = 20): - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - frames = list(self._frames[int(start):]) - n = len(frames) - if n > 0: - imageio.mimwrite(path, frames, fps=fps) - return {"path": path, "n_frames": n, "fps": fps} + def frame_slice(self, start: int) -> list[np.ndarray]: + return list(self._frames[int(start):]) def set_obs(self, obs): self._last_obs = obs @@ -199,6 +125,7 @@ def _vlm_chunk(self, instruction: str): if original_task is not None: self._last_obs["task_descriptions"] = original_task + @updatestate def pi0_pick( self, prompt: str, @@ -274,6 +201,7 @@ def pi0_pick( }, } + @updatestate def pi0_doubled( self, prompt: str, @@ -315,6 +243,7 @@ def pi0_doubled( }, } + @updatestate def move_to( self, xyz, @@ -379,6 +308,7 @@ def move_to( "libero_terminated": self.env.episode_terminated, } + @updatestate def rotate_wrist( self, *, @@ -452,6 +382,7 @@ def _yaw_of(quat_xyzw): "libero_terminated": self.env.episode_terminated, } + @updatestate def rotate_pitch( self, *, @@ -533,6 +464,7 @@ def _pitch_of(quat_xyzw): "libero_terminated": self.env.episode_terminated, } + @updatestate def move_pose( self, xyz, @@ -603,6 +535,7 @@ def _yaw_of(q): "libero_terminated": self.env.episode_terminated, } + @updatestate def release( self, *, @@ -631,6 +564,7 @@ def release( "libero_terminated": self.env.episode_terminated, } + @updatestate def set_gripper( self, *, @@ -659,9 +593,11 @@ def segment( self, prompt: str = "", camera: str = "agentview", - step: int | None = None, + step: int = -1, point: list[int] | None = None, min_score: float = 0.2, + *, + state: EnvState, ) -> dict: """Call SAM3 on an existing image artifact without advancing the env. @@ -669,9 +605,11 @@ def segment( artifacts. Errors are structured so the agent can continue with image inspection and ``back_project``. """ - nn = _latest_step() if step is None else int(step) - if nn is None: - return {"error": "no state entries; cannot select segment image"} + try: + record = state.get(step) + except Exception as exc: + return {"error": f"state step not available: {exc}"} + nn = record.step_idx camera = camera or "agentview" prompt = prompt.strip() @@ -680,26 +618,27 @@ def segment( if has_prompt == has_point: return {"error": "segment needs exactly one of prompt or point"} try: - image_path, world_path, artifact_pairs = _select_segment_artifacts( - nn, camera + image_name, world_name, artifact_pairs = _select_segment_artifacts( + state, record, camera ) except ValueError as e: return {"error": str(e)} - if image_path is None: + if image_name is None or world_name is None: return { "error": "complete segment artifacts not found", "step": nn, "camera": camera, - "checked_paths": [ - str(path) + "checked_artifacts": [ + name for image, world in artifact_pairs - for path in (image, world) + for name in (image, world) + if name ], } try: data = self._sam3_client.segment( - image_path, + state.load_bytes(image_name, step=nn), text_prompt=prompt if has_prompt else None, point=point, min_score=min_score, @@ -709,36 +648,41 @@ def segment( "error": str(e), "step": nn, "camera": camera, - "image_path": str(image_path), + "image_artifact": image_name, } except Exception as e: return { "error": f"segmentation service call failed: {e}", "step": nn, "camera": camera, - "image_path": str(image_path), + "image_artifact": image_name, "fallback": "Use manual visual localization and back_project.", } - out_dir = get_output_dir() - segment_path, overlay_candidate_path, segment_index = ( - _next_segment_artifact_paths(out_dir, nn) - ) - overlay_path = None + segment_index = _next_segment_index(record) + segment_name = f"segment_{segment_index:02d}.json" + overlay_name = f"segment_overlay_{segment_index:02d}.png" + saved_overlay = None mask = data.mask if data.found and isinstance(mask, np.ndarray): - if world_path is None or not world_path.exists(): + try: + world_map = state.load(world_name, step=nn) + except Exception as exc: world_result = { "world_xyz": None, - "world_error": "world map artifact not found for selected image", - "expected_world_path": str(world_path) if world_path else None, + "world_error": f"world map artifact not available: {exc}", + "expected_world_artifact": world_name, } else: - world_result = _mask_to_world(mask, np.load(world_path)) - world_result["world_path"] = str(world_path) - overlay_path = overlay_candidate_path - if not _write_segment_overlay(image_path, mask, overlay_path): - overlay_path = None + world_result = _mask_to_world(mask, world_map) + world_result["world_artifact"] = world_name + overlay = _make_segment_overlay(state.load(image_name, step=nn), mask) + if overlay is not None and state.save( + overlay_name, + overlay, + step=nn, + ): + saved_overlay = overlay_name else: world_result = { "world_xyz": None, @@ -751,7 +695,7 @@ def segment( "camera": camera, "source_step": nn, "segment_index": segment_index, - "image_path": str(image_path), + "image_artifact": image_name, "min_score": min_score, "score": round(float(data.score), 3) if data.score is not None else None, "box": data.box, @@ -764,14 +708,18 @@ def segment( if not data.found: segment_blob["error"] = data.reason or "SAM3 found no mask" segment_blob.update(world_result) - segment_path.write_text(json.dumps(segment_blob, indent=2, default=str)) + state.save( + segment_name, + segment_blob, + step=nn, + ) result = { "found": data.found, "step": nn, "camera": camera, - "image_path": str(image_path), - "segment_path": str(segment_path), + "image_artifact": image_name, + "segment_artifact": segment_name, "score": segment_blob["score"], "box": segment_blob["box"], "world_xyz": segment_blob["world_xyz"], @@ -780,87 +728,70 @@ def segment( if "error" in segment_blob: result["error"] = segment_blob["error"] result["fallback"] = "Use manual visual localization and back_project." - if overlay_path is not None and overlay_path.exists(): - result["overlay_path"] = str(overlay_path) + if saved_overlay is not None: + result["overlay_artifact"] = saved_overlay + result["_image_bytes"] = state.load_bytes(saved_overlay, step=nn) return result -# --------------------------------------------------------------------------- -# State artifacts -# --------------------------------------------------------------------------- - +def _is_primitive_action(name: object) -> bool: + """Whether ``name`` is a state-advancing LIBERO primitive. -def _append_state(output_dir: str, blob: dict) -> None: - """Append *blob* to ``/states.json`` atomically.""" - path = artifact_path(output_dir, "states") - tmp_path = path.parent / f"{path.name}.tmp" - if path.exists(): - with open(path) as f: - states = json.load(f) - else: - states = [] - states.append(blob) - with open(tmp_path, "w") as f: - json.dump(states, f, indent=2) - os.replace(tmp_path, path) + A primitive is any ``@updatestate``-marked method on + :class:`LiberoPrimitives`; read-only tools (``view_env_state``, + ``back_project``, ``segment``, ...) and non-strings read as ``False``. + """ + if not isinstance(name, str): + return False + method = getattr(LiberoPrimitives, name, None) + return method is not None and bool(getattr(method, "_updates_state", False)) -def write_recipe_from_states(output_dir: str, recipe_tag: str) -> str: +def write_recipe_from_states(state: EnvState, recipe_tag: str) -> str: """Find a command sequence that gets ``libero_terminated=True``. - Export non-error LIBERO primitive commands from ``states.json`` and - successful segment calls from ``segments/segment_*.json``. + Export non-error LIBERO primitive commands and successful segment calls. """ - states_path = artifact_path(output_dir, "states") - if states_path.exists(): - with open(states_path) as f: - states = json.load(f) - else: - states = [] - command_events = [] - for step_idx, entry in enumerate(states): - if not entry: - continue - command = entry.get("command") - if command is None: - continue - if command.get("action") not in PRIMITIVE_TOOL_NAMES: - continue - result = entry.get("result") - if isinstance(result, dict) and result.get("error"): - continue - command_events.append(((step_idx, -1), command)) - - for artifact in artifact_path(output_dir, "segments").glob("segment_*.json"): - with artifact.open() as f: - segment = json.load(f) - if segment.get("error"): - continue - if segment["mode"] == "text": - command = { - "action": "segment", - "prompt": segment["prompt"], - "camera": segment["camera"], - } - else: - command = { - "action": "segment", - "point": segment["point"], - "camera": segment["camera"], - } - source_step = int(segment["source_step"]) - event_order = (source_step, int(segment["segment_index"])) - command_events.append((event_order, command)) + for record in state.records(): + command = record.command + result = record.result + if ( + command is not None + and _is_primitive_action(command.get("action")) + and not (isinstance(result, dict) and result.get("error")) + ): + command_events.append(((record.step_idx, -1), command)) + + for name in sorted(record.artifacts): + if not (name.startswith("segment_") and name.endswith(".json")): + continue + segment = state.load(name, step=record.step_idx) + if segment.get("error"): + continue + if segment["mode"] == "text": + segment_command = { + "action": "segment", + "prompt": segment["prompt"], + "camera": segment["camera"], + } + else: + segment_command = { + "action": "segment", + "point": segment["point"], + "camera": segment["camera"], + } + event_order = (record.step_idx, int(segment["segment_index"])) + command_events.append((event_order, segment_command)) - recipe_path = os.path.join(output_dir, f"recipe_{recipe_tag}.jsonl") - tmp_path = recipe_path + ".tmp" command_events.sort(key=lambda event: event[0]) - with open(tmp_path, "w") as f: - for _, command in command_events: - f.write(json.dumps(command, separators=(",", ":")) + "\n") - os.replace(tmp_path, recipe_path) - return recipe_path + recipe_name = f"recipe_{recipe_tag}.jsonl" + state.save( + recipe_name, + [command for _, command in command_events], + step=None, + ) + return recipe_name def _metric_depth(depth: Any, camera_meta: dict) -> np.ndarray: @@ -889,39 +820,13 @@ def _world_from_depth(depth_metric: np.ndarray, camera_meta: dict) -> np.ndarray return (camera_points @ extrinsic.T)[..., :3] -def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, - log: dict | None = None) -> dict: - """Dump state snapshot, images, and depth for step *step_idx*. - - Writes: - - ``/images/image_NN.png`` (Pi0-frame agentview) - - ``/images_cam/image_cam_NN.png`` (calibration-frame agentview) - - ``/depths/depth_NN.npy`` (metric depth, meters) - - ``/world/world_NN.npy`` (agentview world xyz map) - - ``/images_wrist/image_wrist_NN.png`` - - ``/depths_wrist/depth_wrist_NN.npy`` - - ``/world_wrist/world_wrist_NN.npy`` - - ``/wrist_meta/wrist_meta_NN.json`` - - high-res ``images_cam_hi`` / ``world_hi`` artifacts - - high-res ``images_wrist_hi`` / ``world_wrist_hi`` artifacts - - ``/camera_meta.json`` (static, once) - - appends the step blob to ``/states.json`` - - If *log* is provided (the return value of :func:`execute`), its - ``command``, ``result``, and ``elapsed_s`` fields are merged into the - step blob so a single entry captures everything. - """ - for directory in ARTIFACT_DIRECTORIES: - (Path(output_dir) / directory).mkdir(parents=True, exist_ok=True) - - agent_world_map = None - wrist_world_map = None - agent_world_map_hi = None - wrist_world_map_hi = None - # Reuse one raw observation snapshot for state and per-step artifacts. +def dump_state( + primitives: LiberoPrimitives, + env_state: EnvState, + log: dict | None = None, +) -> StepRecord: + """Save one Libero observation through its owned state record.""" raw = primitives.env.raw_obs() - # Expose robot proprioception and object names, but never privileged - # object coordinates; the agent must localize through visual artifacts. state = { "robot0_eef_pos": [float(x) for x in raw["robot0_eef_pos"]], "robot0_eef_quat": [float(x) for x in raw["robot0_eef_quat"]], @@ -932,9 +837,32 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, if k.endswith("_pos") and "robot0" not in k and "to_robot" not in k ), } - imageio.imwrite( - artifact_path(output_dir, "policy_image", step=step_idx, camera="agentview", resolution="low"), + log = log or {} + with env_state.record_step( + state=state, + command=log.get("command"), + result=log.get("result"), + elapsed_s=log.get("elapsed_s"), + extras={ + "terminated": primitives.env.episode_terminated, + "episode_truncated": primitives.env.episode_truncated, + "task_language": primitives.env.get_task_language(), + }, + ) as step_idx: + _save_observation_artifacts(primitives, env_state, step_idx, raw) + return env_state.get(step_idx) + + +def _save_observation_artifacts( + primitives: LiberoPrimitives, + env_state: EnvState, + step_idx: int, + raw: dict[str, Any], +) -> None: + env_state.save( + "agentview_policy.png", primitives._last_obs["main_images"], + step=step_idx, ) # --- camera calibration (static for agentview): fetch metadata as needed --- @@ -943,8 +871,7 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, height=256, width=256, ) or {} - camera_meta_path = artifact_path(output_dir, "metadata", camera="agentview", resolution="low") - if agentview_meta and not camera_meta_path.exists(): + if agentview_meta: cam_meta_out = dict(agentview_meta) cam_meta_out["projection"] = ( "Prefer the back_project(row, col, step=NN) MCP tool; it " @@ -953,24 +880,30 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, "256x256 calibration-frame image." ) cam_meta_out["note"] = ( - "depth_NN.npy is in this camera frame (vertical-flipped raw " - "buffer). image_NN.png is rotated 180deg (Pi0 convention) and " - "is NOT in the same frame as depth/K." + "The agentview_depth.npy observation is aligned with agentview.png. " + "agentview_policy.png uses the Pi0 orientation and must not supply " + "pixels for back-projection." + ) + env_state.save( + "agentview_metadata.json", + cam_meta_out, + step=step_idx, ) - with open(camera_meta_path, "w") as f: - json.dump(cam_meta_out, f, indent=2) # --- per-step RGB in the depth/K frame (vertical-flip of the raw buffer) --- - # The agent picks object pixels HERE (same frame as depth_NN.npy + K), so - # pixel -> depth -> back-project is direct. (image_NN.png is the 180°-rotated - # Pi0-convention frame and must NOT be used for back-projection.) + # Agentview pixels align with the matching depth and calibration. The policy + # image uses Pi0 orientation and must not supply back-projection pixels. try: ci = raw.get("agentview_image") if ci is not None: ci = np.asarray(ci) if ci.dtype != np.uint8: ci = ci.astype(np.uint8) - imageio.imwrite(artifact_path(output_dir, "image", step=step_idx, camera="agentview", resolution="low"), ci[::-1]) + env_state.save( + "agentview.png", + ci[::-1], + step=step_idx, + ) except Exception as e: logger.warning("image_cam dump failed: %s", e) @@ -983,21 +916,19 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, # depth map in this frame. VERIFIED 5/5: projecting each GT object # world pos via M lands on a pixel whose depth_flip[row,col] matches # the object's surface depth (plate Δ6mm, cookies Δ14mm). So - # pixel(row,col) in depth_NN.npy back-projects correctly with - # camera_meta.json (NOT the same frame as the 180°-rotated - # image_NN.png — see camera_meta note). + # Pixels in agentview.png align with agentview_depth.npy and the + # per-step agentview metadata saved with the record artifacts. d = _metric_depth(d, agentview_meta)[::-1] - np.save( - artifact_path(output_dir, "depth", step=step_idx, camera="agentview", resolution="low"), + env_state.save( + "agentview_depth.npy", d.astype(np.float32), + step=step_idx, ) world = _world_from_depth(d, agentview_meta).astype(np.float32) - np.save( - artifact_path(output_dir, "world", step=step_idx, camera="agentview", resolution="low"), + env_state.save( + "agentview_world.npy", world, - ) - agent_world_map = _artifact_relative_path( - step_idx, "agentview", "low", "world" + step=step_idx, ) except Exception as e: logger.warning("depth dump failed: %s", e) @@ -1011,7 +942,11 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, wimg = np.asarray(wimg) if wimg.dtype != np.uint8: wimg = wimg.astype(np.uint8) - imageio.imwrite(artifact_path(output_dir, "image", step=step_idx, camera="wrist", resolution="low"), wimg[::-1]) + env_state.save( + "wrist.png", + wimg[::-1], + step=step_idx, + ) except Exception as e: logger.warning("wrist image dump failed: %s", e) @@ -1031,31 +966,30 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, logger.warning("wrist camera meta missing; skipping wrist depth/world") else: wdpt_metric = _metric_depth(wdpt_arr, wmeta)[::-1] - np.save( - artifact_path(output_dir, "depth", step=step_idx, camera="wrist", resolution="low"), + env_state.save( + "wrist_depth.npy", wdpt_metric.astype(np.float32), + step=step_idx, ) world_w = _world_from_depth(wdpt_metric, wmeta).astype(np.float32) - np.save( - artifact_path(output_dir, "world", step=step_idx, camera="wrist", resolution="low"), + env_state.save( + "wrist_world.npy", world_w, - ) - wrist_world_map = _artifact_relative_path( - step_idx, "wrist", "low", "world" + step=step_idx, ) wmeta_out = dict(wmeta) wmeta_out["note"] = ( "MOVING camera: extrinsic_cam2world is for THIS step " - "only. world_wrist_NN.npy[row,col] gives world " - "(x,y,z) for that pixel, in the SAME world frame as " - "agentview world_NN.npy." + "only. The matching wrist world-map observation gives world " + "(x,y,z) for that pixel in the same world frame as the " + "agentview world-map artifact." + ) + env_state.save( + "wrist_metadata.json", + wmeta_out, + step=step_idx, ) - with open( - artifact_path(output_dir, "metadata", step=step_idx, camera="wrist", resolution="low"), - "w", - ) as f: - json.dump(wmeta_out, f, indent=2) except Exception as e: logger.warning("wrist depth/world dump failed: %s", e) @@ -1069,20 +1003,19 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, meta_hi = primitives.env.get_camera_meta("agentview", 1024, 1024) if meta_hi is None: raise RuntimeError("agentview camera metadata missing") - imageio.imwrite( - artifact_path(output_dir, "image", step=step_idx, camera="agentview", resolution="high"), + env_state.save( + "agentview_high.png", np.asarray(rgb_hi)[::-1], + step=step_idx, ) world_hi = _world_from_depth( _metric_depth(depth_hi, meta_hi)[::-1], meta_hi, ).astype(np.float16) - np.save( - artifact_path(output_dir, "world", step=step_idx, camera="agentview", resolution="high"), + env_state.save( + "agentview_world_high.npy", world_hi, - ) - agent_world_map_hi = _artifact_relative_path( - step_idx, "agentview", "high", "world" + step=step_idx, ) except Exception as e: logger.warning("agentview high-res dump failed: %s", e) @@ -1099,86 +1032,35 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, ) if meta_wrist_hi is None: raise RuntimeError("robot0_eye_in_hand camera metadata missing") - imageio.imwrite( - artifact_path(output_dir, "image", step=step_idx, camera="wrist", resolution="high"), + env_state.save( + "wrist_high.png", np.asarray(rgb_wrist_hi)[::-1], + step=step_idx, ) world_wrist_hi = _world_from_depth( _metric_depth(depth_wrist_hi, meta_wrist_hi)[::-1], meta_wrist_hi, ).astype(np.float16) - np.save( - artifact_path(output_dir, "world", step=step_idx, camera="wrist", resolution="high"), + env_state.save( + "wrist_world_high.npy", world_wrist_hi, - ) - wrist_world_map_hi = _artifact_relative_path( - step_idx, "wrist", "high", "world" + step=step_idx, ) except Exception as e: logger.warning("wrist high-res dump failed: %s", e) - for old_step in range(max(0, int(step_idx) - 4)): - for path in ( - artifact_path(output_dir, "image", step=old_step, camera="agentview", resolution="high"), - artifact_path(output_dir, "world", step=old_step, camera="agentview", resolution="high"), - artifact_path(output_dir, "image", step=old_step, camera="wrist", resolution="high"), - artifact_path(output_dir, "world", step=old_step, camera="wrist", resolution="high"), - ): - try: - os.unlink(path) - except FileNotFoundError: - pass - - blob = { - "step_idx": step_idx, - "libero_terminated": primitives.env.episode_terminated, - "episode_truncated": primitives.env.episode_truncated, - "task_language": primitives.env.get_task_language(), - "state": state, - "world_map": agent_world_map, - "wrist_world_map": wrist_world_map, - "world_map_hi": agent_world_map_hi, - "wrist_world_map_hi": wrist_world_map_hi, - } - # Merge the execution log (command + result + elapsed_s) into the - # state blob so a single entry captures everything for the step. - if log is not None: - blob["command"] = log.get("command") - blob["result"] = log.get("result") - blob["elapsed_s"] = log.get("elapsed_s") - _append_state(output_dir, blob) - return blob - # --------------------------------------------------------------------------- # Tool schema declarations (Anthropic-shaped canonical schema) # --------------------------------------------------------------------------- -PRIMITIVE_TOOL_NAMES: tuple[str, ...] = ( - "move_to", - "pi0_pick", - "pi0_doubled", - "release", - "set_gripper", - "rotate_wrist", - "rotate_pitch", - "move_pose", -) - TOOLS_SPEC = [ { - "name": "view_driver_state", + "name": "view_env_state", "description": ( - "Read step NN from `states.json` + the matching " - "state images in {{output_dir}}. If step is " - "null, returns the latest entry. Each entry contains the robot " - "state, libero_terminated flag, command log, and result. Returns " - "available PNG paths in this stable " - "order: 1) `images/image_NN.png` (Pi0-frame agentview), " - "2) `images_cam/image_cam_NN.png` (calibration-frame agentview), " - "3) `images_wrist/image_wrist_NN.png` (calibration-frame wrist). " - "High-resolution calibration-frame images are returned as file " - "paths, not embedded as image bytes. " + "Read one recorded state and its observation artifacts. Step -1 " + "selects the latest entry. Embeds policy, agentview, and wrist " + "images when available. " "Use the calibration-frame images for pixel back-projection; JSON " "state alone is not enough. Use agentview for global tabletop " "layout and object locations; use wrist for close-range details " @@ -1188,8 +1070,9 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, "type": "object", "properties": { "step": { - "type": ["integer", "null"], - "description": "Step number; 0 = initial. Null = latest.", + "type": "integer", + "default": -1, + "description": "Step number; 0 = initial, -1 = latest.", }, }, }, @@ -1380,9 +1263,7 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, { "name": "view_camera_meta", "description": ( - "Read camera calibration metadata from the output dir. " - "camera='agentview' reads static camera_meta.json. " - "camera='wrist' reads the per-step wrist metadata." + "Read per-step camera calibration metadata from recorded artifacts." ), "input_schema": { "type": "object", @@ -1393,8 +1274,9 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, "description": "Camera metadata to read (default agentview).", }, "step": { - "type": ["integer", "null"], - "description": "Wrist metadata step to use (default latest).", + "type": "integer", + "default": -1, + "description": "Metadata step to use; -1 = latest.", }, }, }, @@ -1420,8 +1302,9 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, "description": "Artifact camera to use (default agentview).", }, "step": { - "type": ["integer", "null"], - "description": "Step NN to segment; null = latest.", + "type": "integer", + "default": -1, + "description": "Step to segment; -1 = latest.", }, "point": { "type": ["array", "null"], @@ -1447,7 +1330,7 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, "selected camera's precomputed world map. Row 0 = top of image, " "col 0 = left. Returns world_xyz in meters.\n\n" "USE THIS to find where an object is in the world — look at " - "the high-resolution paths returned by view_driver_state " + "the embedded high-resolution image returned by view_env_state " "to pick a pixel on the target object, then call back_project. " "The default resolution is high (1024x1024). Pass " "resolution='low' only for pixels from the embedded/standard " @@ -1476,8 +1359,9 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, "description": "Pixel column (0=left) in the selected resolution image.", }, "step": { - "type": ["integer", "null"], - "description": "Depth/world-map step to use (default latest). 0 for initial.", + "type": "integer", + "default": -1, + "description": "Depth/world-map step; 0 = initial, -1 = latest.", }, "camera": { "type": "string", @@ -1517,151 +1401,68 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, ] -# --------------------------------------------------------------------------- -# State trace readers -# --------------------------------------------------------------------------- - - -def _load_states() -> list: - """Return the parsed state trace from the local output dir.""" - path = artifact_path(get_output_dir(), "states") - if not path.exists(): - return [] - with open(path) as f: - return json.load(f) - - -def _latest_step() -> int | None: - states = _load_states() - if not states: - return None - return states[-1]["step_idx"] - - -def _load_step(nn: int) -> dict: - """Look up the state blob for step ``nn`` from states.json.""" - for entry in _load_states(): - if entry.get("step_idx") == nn: - return entry - raise FileNotFoundError(f"step {nn} not present in states.json") - - -def _load_image_path(nn: int, kind: str) -> str | None: - """Return the path to a dumped state image. None if not present.""" - out_dir = get_output_dir() - if kind == "agent": - path = artifact_path(out_dir, "policy_image", step=nn, camera="agentview", resolution="low") - elif kind == "camera": - path = artifact_path(out_dir, "image", step=nn, camera="agentview", resolution="low") - elif kind == "wrist": - path = artifact_path(out_dir, "image", step=nn, camera="wrist", resolution="low") - else: - raise ValueError(f"unknown image kind: {kind}") - if not path.exists(): - return None - return str(path) - - -def _load_camera_meta(camera: str = "agentview", nn: int | None = None) -> dict: - out_dir = get_output_dir() - if camera == "agentview": - path = artifact_path(out_dir, "metadata", camera="agentview", resolution="low") - elif camera == "wrist" and nn is not None: - path = artifact_path(out_dir, "metadata", step=nn, camera="wrist", resolution="low") - else: - raise ValueError("camera must be 'agentview' or 'wrist' with nn") - if not path.exists(): - raise FileNotFoundError(f"{path.name} not found in {out_dir}") - with open(path) as f: - return json.load(f) - - -def _load_depth(camera: str, nn: int) -> np.ndarray: - out_dir = get_output_dir() - if camera not in ("agentview", "wrist"): - raise ValueError("camera must be 'agentview' or 'wrist'") - path = artifact_path(out_dir, "depth", step=nn, camera=camera, resolution="low") - if not path.exists(): - raise FileNotFoundError(f"{path.name} not found in {out_dir}") - depth = np.load(path) - if depth.ndim == 3: - depth = depth[..., 0] - return depth - - -def view_driver_state(step: int | None = None) -> dict: - latest = _latest_step() - if latest is None: - return {"error": "no state entries; env not ready"} - nn = latest if step is None else int(step) +def view_env_state(step: int = -1, *, state: EnvState) -> dict: try: - data = _load_step(nn) - except Exception as e: - return {"error": f"step {nn} not present in state trace: {e}"} - - out: dict = {"step": nn} - out["task_language"] = data.get("task_language") - # Default to {} (not the entire blob) when "state" is missing — otherwise - # command/result/vla_desync would bleed into the "state" field, confusing - # the LLM about what robot state actually contains. - out["state"] = data.get("state", {}) - out["libero_terminated"] = data.get("libero_terminated") - out["episode_truncated"] = data.get("episode_truncated") - out["world_map"] = data.get("world_map") - out["wrist_world_map"] = data.get("wrist_world_map") - out["world_map_hi"] = data.get("world_map_hi") - out["wrist_world_map_hi"] = data.get("wrist_world_map_hi") + record = state.get(step) + except Exception as exc: + return {"error": f"state step not available: {exc}"} + + nn = record.step_idx + extras = record.extras + out: dict = { + "step": nn, + "state": record.state, + "artifacts": sorted(record.artifacts), + } + out["task_language"] = extras.get("task_language") + out["libero_terminated"] = extras.get("terminated") + out["episode_truncated"] = extras.get("episode_truncated") out["log"] = { - "command": data.get("command"), - "result": data.get("result"), - "elapsed_s": data.get("elapsed_s"), + "command": record.command, + "result": record.result, + "elapsed_s": record.elapsed_s, } - for field, kind in ( - ("image_path", "agent"), - ("image_cam_path", "camera"), - ("image_wrist_path", "wrist"), - ): - image_path = _load_image_path(nn, kind) - if image_path: - out[field] = image_path - for field, camera in ( - ("image_cam_hi_path", "agentview"), - ("image_wrist_hi_path", "wrist"), + for slot, names in ( + ("_image_bytes", ("agentview_policy.png",)), + ("_image_cam_bytes", ("agentview_high.png", "agentview.png")), + ("_image_wrist_bytes", ("wrist_high.png", "wrist.png")), ): - image_path = artifact_path(get_output_dir(), "image", step=nn, camera=camera, resolution="high") - if image_path.exists(): - out[field] = str(image_path) + name = next((name for name in names if name in record.artifacts), None) + if name: + try: + out[slot] = state.load_bytes(name, step=nn) + except FileNotFoundError: + pass return out -def _select_segment_artifacts(nn: int, camera: str): - out_dir = get_output_dir() +def _select_segment_artifacts( + state: EnvState, + record: StepRecord, + camera: str, +) -> tuple[str | None, str | None, list[tuple[str | None, str | None]]]: if camera not in ("agentview", "wrist"): raise ValueError(f"unknown segment camera: {camera}") pairs = [ - ( - artifact_path(out_dir, "image", step=nn, camera=camera, resolution=resolution), - artifact_path(out_dir, "world", step=nn, camera=camera, resolution=resolution), - ) - for resolution in ("high", "low") + (f"{camera}_high.png", f"{camera}_world_high.npy"), + (f"{camera}.png", f"{camera}_world.npy"), ] - - for image_path, world_path in pairs: - if image_path.exists() and world_path.exists(): - return image_path, world_path, pairs + for image_name, world_name in pairs: + if ( + image_name in record.artifacts + and world_name in record.artifacts + and state.exists(image_name, step=record.step_idx) + and state.exists(world_name, step=record.step_idx) + ): + return image_name, world_name, pairs return None, None, pairs -def _next_segment_artifact_paths(out_dir: Path, nn: int): - segments_dir = artifact_path(out_dir, "segments") - segments_dir.mkdir(parents=True, exist_ok=True) +def _next_segment_index(record: StepRecord) -> int: idx = 0 - while True: - segment_path = segments_dir / f"segment_{nn:02d}_{idx:02d}.json" - overlay_path = segments_dir / f"segment_overlay_{nn:02d}_{idx:02d}.png" - if not segment_path.exists() and not overlay_path.exists(): - return segment_path, overlay_path, idx + while f"segment_{idx:02d}.json" in record.artifacts: idx += 1 + return idx def _mask_to_world(mask: np.ndarray, world_map: np.ndarray, @@ -1718,56 +1519,58 @@ def _mask_to_world(mask: np.ndarray, world_map: np.ndarray, return result -def _write_segment_overlay(image_path: Path, mask: np.ndarray, - overlay_path: Path) -> bool: - try: - image = imageio.imread(image_path) - if image.ndim != 3 or image.shape[:2] != mask.shape: - return False - overlay = image.copy() - red = np.zeros_like(overlay) - red[..., 0] = 255 - overlay[mask] = ( - 0.55 * overlay[mask].astype(np.float32) - + 0.45 * red[mask].astype(np.float32) - ).astype(np.uint8) - imageio.imwrite(overlay_path, overlay) - return overlay_path.exists() - except Exception: - return False +def _make_segment_overlay( + image: np.ndarray, + mask: np.ndarray, +) -> np.ndarray | None: + if image.ndim != 3 or image.shape[:2] != mask.shape: + return None + overlay = image.copy() + red = np.zeros_like(overlay) + red[..., 0] = 255 + overlay[mask] = ( + 0.55 * overlay[mask].astype(np.float32) + + 0.45 * red[mask].astype(np.float32) + ).astype(np.uint8) + return overlay -def view_camera_meta(camera: str = "agentview", step: int | None = None) -> dict: +def view_camera_meta( + camera: str = "agentview", + step: int = -1, + *, + state: EnvState, +) -> dict: """Read camera calibration metadata for localization.""" if camera not in ("agentview", "wrist"): return {"error": f"bad camera '{camera}' (use 'agentview' or 'wrist')"} - nn = None - if camera == "wrist": - nn = _latest_step() if step is None else int(step) - if nn is None: - return {"error": "no wrist metadata available"} - try: - meta = _load_camera_meta(camera, nn) + record = state.get(step) + metadata_name = f"{camera}_metadata.json" + if metadata_name not in record.artifacts: + raise FileNotFoundError(metadata_name) + meta = state.load(metadata_name, step=record.step_idx) except Exception as e: return {"error": f"{camera} camera metadata not found: {e}"} if camera == "agentview": return {"camera": "agentview", "camera_meta": meta} - return {"camera": "wrist", "step": nn, "camera_meta": meta} + return {"camera": "wrist", "step": record.step_idx, "camera_meta": meta} def back_project( row: int | None = None, col: int | None = None, - step: int | None = None, + step: int = -1, camera: str = "agentview", resolution: str = "high", row_range: list | None = None, col_range: list | None = None, z_min: float | None = None, z_max: float | None = None, + *, + state: EnvState, ) -> dict: """Look up a pixel's world XYZ in the precomputed world map.""" if camera not in ("agentview", "wrist"): @@ -1784,24 +1587,16 @@ def back_project( ) } - latest = _latest_step() - nn = latest if step is None else int(step) - if nn is None: - return {"error": "no depth/world-map files available"} - try: - data = _load_step(nn) + record = state.get(step) except Exception as e: - return {"error": f"step {nn} not present in state trace: {e}"} + return {"error": f"state step not available: {e}"} + nn = record.step_idx - if camera == "agentview": - hi_artifact = data.get("world_map_hi") - low_artifact = data.get("world_map") - else: - hi_artifact = data.get("wrist_world_map_hi") - low_artifact = data.get("wrist_world_map") + hi_artifact = f"{camera}_world_high.npy" + low_artifact = f"{camera}_world.npy" source_artifact = hi_artifact if resolution == "high" else low_artifact - if not source_artifact: + if source_artifact not in record.artifacts: return { "error": ( f"{camera} {resolution}-resolution world map not recorded " @@ -1810,8 +1605,7 @@ def back_project( } try: - world_path = artifact_path(get_output_dir(), "world", step=nn, camera=camera, resolution=resolution) - world_map = np.load(world_path) + world_map = state.load(str(source_artifact), step=nn) except Exception as e: return { "error": ( @@ -1897,7 +1691,12 @@ def back_project( depth_m = None if source_artifact == low_artifact: try: - depth = _load_depth(camera, nn) + depth_artifact = f"{camera}_depth.npy" + if depth_artifact not in record.artifacts: + raise FileNotFoundError(depth_artifact) + depth = state.load(depth_artifact, step=nn) + if depth.ndim == 3: + depth = depth[..., 0] except Exception as e: return {"error": f"{camera} depth not found for step {nn}: {e}"} depth_m = float(depth[row, col]) diff --git a/rpent/cli/dashboard.py b/rpent/cli/dashboard.py index a367820b..3ae39131 100644 --- a/rpent/cli/dashboard.py +++ b/rpent/cli/dashboard.py @@ -150,7 +150,6 @@ def _run_dashboard_task( toolkit = get_toolkit( args.env_name, primitives_kwargs=primitives_kwargs, - video_path=str(output_dir / "episode.mp4"), dashboard_events=state, ) planner = build_planner( diff --git a/rpent/cli/main.py b/rpent/cli/main.py index 2596793d..bc62c165 100644 --- a/rpent/cli/main.py +++ b/rpent/cli/main.py @@ -86,8 +86,9 @@ def _build_argparser() -> argparse.ArgumentParser: description="Standalone hybrid LLM-in-the-loop agent for LIBERO PRO", ) - ap.add_argument("--env", dest="env_name", required=True, choices=["libero"], - help="Environment backend: libero.") + ap.add_argument("--env", dest="env_name", required=True, + choices=["libero", "lerobot", "franka"], + help="Environment backend: libero | lerobot | franka.") # models ap.add_argument("--planner", default="api", @@ -219,7 +220,6 @@ def main() -> int: toolkit = get_toolkit( env_name, primitives_kwargs=primitives_kwargs, - video_path=str(Path(output_dir) / "episode.mp4"), dashboard_events=dashboard_events, ) diff --git a/rpent/context/prompts/prompt.py b/rpent/context/prompts/prompt.py index b3c9e47b..43e6dad0 100644 --- a/rpent/context/prompts/prompt.py +++ b/rpent/context/prompts/prompt.py @@ -5,7 +5,7 @@ OUTPUT = BulletList([ "Brief reasoning before each tool call (1-2 sentences): observation -> decision.", - "Don't re-read files already in this session. Don't view_driver_state right after a primitive tool already returned the state.", + "Don't re-read files already in this session. Don't view_env_state right after a primitive tool already returned the state.", "Numerical coords in 3 decimals are enough.", "Save artifacts BEFORE calling finish. Stop immediately after writing the audit; do not chat further.", ]) diff --git a/rpent/dashboard/events.py b/rpent/dashboard/events.py index 5784624d..0c6cbfae 100644 --- a/rpent/dashboard/events.py +++ b/rpent/dashboard/events.py @@ -39,6 +39,15 @@ class ToolResultEvent: result: Any +@dataclass(frozen=True, slots=True) +class StepRecordEvent: + """Publish one recorded environment step and its artifact context.""" + + record: Any + env_state: Any + frame_artifacts: dict[str, str] + + @dataclass(frozen=True, slots=True) class RunStartedEvent: """Mark startup complete and the agent run active.""" @@ -49,6 +58,7 @@ class RunStartedEvent: | UsageEvent | RuntimeStatusEvent | ToolResultEvent + | StepRecordEvent | RunStartedEvent ) diff --git a/rpent/dashboard/server.py b/rpent/dashboard/server.py index e485d950..f76c5f06 100644 --- a/rpent/dashboard/server.py +++ b/rpent/dashboard/server.py @@ -259,10 +259,11 @@ def api_frame( @app.get("/api/run/video") def api_video(run: str) -> Response: live = self._resolve(run) - if live is None or not live.has_video(): + video = live.video() if live else None + if video is None: return Response(status_code=404) - return FileResponse( - live.video_path, + return Response( + video, media_type="video/mp4", headers={"Cache-Control": "no-store, max-age=0"}, ) @@ -270,11 +271,11 @@ def api_video(run: str) -> Response: @app.get("/api/run/action-video") def api_action_video(run: str, step: int) -> Response: live = self._resolve(run) - path = live.action_video_path(step) if live else None - if path is None: + video = live.action_video(step) if live else None + if video is None: return Response(status_code=404) - return FileResponse( - path, + return Response( + video, media_type="video/mp4", headers={"Cache-Control": "no-store, max-age=0"}, ) diff --git a/rpent/dashboard/state.py b/rpent/dashboard/state.py index 7617cdea..0a1179b8 100644 --- a/rpent/dashboard/state.py +++ b/rpent/dashboard/state.py @@ -11,6 +11,7 @@ DashboardEvent, RunStartedEvent, RuntimeStatusEvent, + StepRecordEvent, ToolResultEvent, TranscriptEvent, UsageEvent, @@ -26,6 +27,7 @@ if TYPE_CHECKING: from rpent.dashboard.commands import TaskCommand + from rpent.tools.state import EnvState, StepRecord RUNTIME_COMPONENTS = ("env", "vla", "sam3") RUNTIME_STATUSES = {"pending", "starting", "ready", "failed"} @@ -78,6 +80,8 @@ def __init__( self._timeline: list[dict[str, Any]] = [] self._frames: dict[str, bytes] = {} self._frame_idx = -1 + self.env_state: EnvState | None = None + self.frame_artifacts: dict[str, str] = {} self._accepting_input = False self._planner_activity: PlannerActivity = "starting" self._interrupt_requested = False @@ -268,6 +272,8 @@ def _begin_task_locked( self._timeline = [] self._frames = {} self._frame_idx = -1 + self.env_state = None + self.frame_artifacts = {} self._accepting_input = False self._planner_activity = "starting" self._interrupt_requested = False @@ -477,6 +483,11 @@ def emit(self, event: DashboardEvent) -> None: if isinstance(event, ToolResultEvent): self._apply_tool_result(event) return + if isinstance(event, StepRecordEvent): + self.env_state = event.env_state + self.frame_artifacts = dict(event.frame_artifacts) + self.on_step(event.record) + return if isinstance(event, RunStartedEvent): self._start() return @@ -502,7 +513,14 @@ def _apply_tool_result(self, event: ToolResultEvent) -> None: result = event.result if not isinstance(result, dict): return - self._apply_frame_paths(result) + frames = { + "camera": result.get("_image_cam_bytes") or result.get("_image_bytes"), + "wrist": result.get("_image_wrist_bytes"), + } + self._update_frames( + step=result.get("step"), + frames={kind: data for kind, data in frames.items() if data}, + ) log = result.get("log") if not isinstance(log, dict): return @@ -521,16 +539,53 @@ def _apply_tool_result(self, event: ToolResultEvent) -> None: "result": log.get("result"), "elapsed_s": log.get("elapsed_s"), "terminated": terminated, - "has_action_video": ( - self.output_dir - / "action_videos" - / f"step_{step:02d}_{command.get('action', name)}.mp4" - ).exists(), + "action_video_artifact": result.get("action_video_artifact"), + "has_action_video": bool(result.get("action_video_artifact")), } with self._lock: self._timeline.append(item) self._terminated = self._terminated or terminated + def on_step(self, record: StepRecord) -> None: + """Project one recorded environment step into frames and timeline.""" + self._update_step_frames(record) + command = record.command + if not isinstance(command, dict) or not command.get("action"): + return + terminated = bool(record.extras.get("terminated")) + action_video = next( + (name for name in sorted(record.artifacts) if name.endswith(".mp4")), + None, + ) + item = { + "step": record.step_idx, + "action": str(command.get("action")), + "args": {key: value for key, value in command.items() if key != "action"}, + "result": record.result, + "elapsed_s": record.elapsed_s, + "terminated": terminated, + "action_video_artifact": action_video, + "has_action_video": action_video is not None, + } + with self._lock: + self._timeline.append(item) + self._terminated = self._terminated or terminated + + def _update_step_frames(self, record: StepRecord) -> None: + """Load dashboard frame bytes from the step's canonical artifacts.""" + env_state = self.env_state + if env_state is None: + return + frames: dict[str, bytes] = {} + for kind, artifact in self.frame_artifacts.items(): + if kind not in FRAME_KINDS or artifact not in record.artifacts: + continue + try: + frames[kind] = env_state.load_bytes(artifact, step=record.step_idx) + except FileNotFoundError: + continue + self._update_frames(step=record.step_idx, frames=frames) + def _apply_frame_paths(self, result: dict[str, Any]) -> None: path_keys = { "camera": "image_cam_path", @@ -675,18 +730,26 @@ def frame(self, kind: str) -> bytes | None: with self._lock: return self._frames.get(kind) - def action_video_path(self, step: int) -> Path | None: + def action_video(self, step: int) -> bytes | None: + env_state = self.env_state + if env_state is None: + return None with self._lock: + artifact = None for item in self._timeline: if int(item.get("step", -1)) != int(step): continue - video_path = ( - self.output_dir - / "action_videos" - / f"step_{int(step):02d}_{item.get('action', '')}.mp4" - ) - return video_path if video_path.exists() else None - return None + artifact = item.get("action_video_artifact") + break + if not artifact: + return None + try: + return env_state.load_bytes(artifact, step=int(step)) + except FileNotFoundError: + return None + + def video(self) -> bytes | None: + return self.video_path.read_bytes() if self.video_path.exists() else None def has_video(self) -> bool: with self._lock: diff --git a/rpent/planner/api_loop.py b/rpent/planner/api_loop.py index 4e904666..6470dadf 100644 --- a/rpent/planner/api_loop.py +++ b/rpent/planner/api_loop.py @@ -49,12 +49,10 @@ _ARGS_LOG_LIMIT = 250 _TOOL_LOG_LIMIT = 350 -#: Cap on cumulative decoded image bytes kept in the resent request history. -_MAX_HISTORY_IMAGE_BYTES = 4 * 1024 * 1024 - -#: Always retain at least this many of the most recent images, even if a single -#: frame exceeds the byte budget, so the model never loses its current view. -_MIN_RECENT_IMAGES = 2 +#: Maximum number of recent images kept in the resent request history. +#: Franka action results usually return scene + wrist, so 4 preserves the +#: latest visual observation while bounding multimodal context growth. +_MAX_HISTORY_IMAGES = 4 class ApiAgentLoop: @@ -335,8 +333,8 @@ def _build_model_settings(model: Model, max_tokens: int) -> ModelSettings: def _prune_history_images(messages: list[ModelMessage]) -> list[ModelMessage]: """Drop old camera images so the resent request body stays bounded.""" - # Every image in history, oldest -> newest: (msg_idx, part_idx, item_idx, nbytes). - located: list[tuple[int, int, int, int]] = [] + # Every image in history, oldest -> newest: (msg_idx, part_idx, item_idx). + located: list[tuple[int, int, int]] = [] for mi, message in enumerate(messages): for pi, part in enumerate(getattr(message, "parts", ()) or ()): if not isinstance(part, UserPromptPart) or not isinstance( @@ -347,24 +345,15 @@ def _prune_history_images(messages: list[ModelMessage]) -> list[ModelMessage]: if isinstance(item, BinaryContent) and item.media_type.startswith( "image/" ): - located.append((mi, pi, ii, len(item.data))) + located.append((mi, pi, ii)) - if not located: + if len(located) <= _MAX_HISTORY_IMAGES: return messages - # Walk newest -> oldest, keeping images while under the byte budget. - keep: set[tuple[int, int, int]] = set() - total = 0 - for rank, (mi, pi, ii, nbytes) in enumerate(reversed(located)): - if rank < _MIN_RECENT_IMAGES or total + nbytes <= _MAX_HISTORY_IMAGE_BYTES: - keep.add((mi, pi, ii)) - total += nbytes - - if len(keep) == len(located): - return messages + keep = set(located[-_MAX_HISTORY_IMAGES:]) drop_items_by_part: dict[tuple[int, int], set[int]] = {} - for mi, pi, ii, _ in located: + for mi, pi, ii in located: if (mi, pi, ii) not in keep: drop_items_by_part.setdefault((mi, pi), set()).add(ii) @@ -419,7 +408,11 @@ def _build_tools(toolkit: Toolkit, *, no_images: bool = False) -> list[Tool]: def read_image(path: str) -> ToolReturn: - """Read a local image path returned by an RPent tool as visual input.""" + """Read an explicitly provided local image file as visual input. + + Environment observations are already embedded by ``view_env_state``; + this helper is only for other user-selected local files. + """ return ToolReturn( return_value=path, content=[BinaryContent.from_path(path)], @@ -430,7 +423,7 @@ def read_image_text_only(path: str) -> str: """``read_image`` stub for ``--no-images``: acknowledge, send no bytes.""" return ( f"{path} exists, but image input is disabled (--no-images, text-only " - "model). Reason from textual state instead: view_driver_state, " + "model). Reason from textual state instead: view_env_state, " "back_project, and the numeric fields in tool results." ) diff --git a/rpent/tools/common.py b/rpent/tools/common.py index fd959566..54e02ce9 100644 --- a/rpent/tools/common.py +++ b/rpent/tools/common.py @@ -3,10 +3,16 @@ import os from pathlib import Path +from typing import Any + +import numpy as np from rpent.utils.config import get_repo_root from rpent.utils.logging import get_output_dir +_BACKPROJECT_RADIUS = 6 +_DEPTH_BAND_M = 0.02 + TOOLS_SPEC: list[dict] = [ { "name": "read_text_file", @@ -134,6 +140,88 @@ def finish(status: str, summary: str) -> dict: return {"_finish": True, "status": status, "summary": summary} +def backproject_points(K, rows, cols, depths) -> np.ndarray: + """Back-project pixel coords + depths to camera-frame XYZ, shape (N, 3).""" + K = np.asarray(K, dtype=np.float64) + rows = np.asarray(rows, dtype=np.float64) + cols = np.asarray(cols, dtype=np.float64) + depths = np.asarray(depths, dtype=np.float64) + fx, fy = K[0, 0], K[1, 1] + cx, cy = K[0, 2], K[1, 2] + return np.stack( + [(cols - cx) * depths / fx, (rows - cy) * depths / fy, depths], axis=1 + ) + + +def robust_surface_centroid( + depth: np.ndarray, + K, + T_base_cam, + row: int, + col: int, + *, + radius: int = _BACKPROJECT_RADIUS, + band: float = _DEPTH_BAND_M, +) -> dict: + """Back-project a pixel neighbourhood to a robust 3D point. + + Back-projects every valid pixel in a ``(2*radius+1)`` window, keeps those on + the dominant surface (depth within ``band`` of the window median, rejecting + background / table / dropouts), and returns the median point + diagnostics. + Returns world ``xyz`` when ``T_base_cam`` is given, otherwise camera-frame + ``xyz_cam``. + """ + row, col = int(row), int(col) + radius = max(0, int(radius)) + height, width = depth.shape[:2] + if not (0 <= row < height and 0 <= col < width): + return { + "error": ( + f"pixel ({row},{col}) out of bounds; image is {height}x{width}" + ) + } + row_start, row_end = max(0, row - radius), min(height, row + radius + 1) + col_start, col_end = max(0, col - radius), min(width, col + radius + 1) + rows, cols = np.mgrid[row_start:row_end, col_start:col_end] + depths = depth[row_start:row_end, col_start:col_end].reshape(-1).astype( + np.float64 + ) + rows = rows.reshape(-1).astype(np.float64) + cols = cols.reshape(-1).astype(np.float64) + valid = np.isfinite(depths) & (depths > 0) + if not np.any(valid): + return {"error": f"no valid depth near ({row},{col}); pick another pixel"} + depths, rows, cols = depths[valid], rows[valid], cols[valid] + median_depth = float(np.median(depths)) + surface = np.abs(depths - median_depth) <= band + depths, rows, cols = depths[surface], rows[surface], cols[surface] + if depths.size == 0: + return {"error": f"no dominant surface depth near ({row},{col})"} + + camera_points = backproject_points(K, rows, cols, depths) + camera_point = np.median(camera_points, axis=0) + out: dict[str, Any] = { + "pixel": [row, col], + "radius": radius, + "n_points": int(camera_points.shape[0]), + "depth_m": round(median_depth, 4), + "xyz_cam": [round(float(value), 4) for value in camera_point], + } + if T_base_cam is not None: + transform = np.asarray(T_base_cam, dtype=np.float64) + base_points = camera_points @ transform[:3, :3].T + transform[:3, 3] + base_point = np.median(base_points, axis=0) + out["xyz"] = [round(float(value), 4) for value in base_point] + out["xy_spread_m"] = round( + float(np.hypot(*base_points[:, :2].std(axis=0))), 4 + ) + else: + out["xy_spread_m"] = round( + float(np.hypot(*camera_points[:, :2].std(axis=0))), 4 + ) + return out + + TOOL_HANDLERS: dict = { "read_text_file": read_text_file, "write_text_file": write_text_file, diff --git a/rpent/tools/state.py b/rpent/tools/state.py new file mode 100644 index 00000000..a059cb77 --- /dev/null +++ b/rpent/tools/state.py @@ -0,0 +1,460 @@ +"""Per-run environment state and artifact storage.""" +from __future__ import annotations + +import copy +import fnmatch +import json +import os +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import imageio.v2 as imageio +import numpy as np + +from rpent.utils.logging import get_logger + +logger = get_logger("env_state") + +_MANIFEST_NAME = "states.json" +_MANIFEST_VERSION = 2 +_IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg"} +_TEXT_SUFFIXES = {".txt", ".md"} +_SUPPORTED_SUFFIXES = _IMAGE_SUFFIXES | { + ".npy", + ".json", + ".jsonl", + ".mp4", + ".bin", +} | _TEXT_SUFFIXES + + +def _json_default(value: Any) -> Any: + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, Path): + return str(value) + raise TypeError(f"{type(value).__name__} is not JSON serializable") + + +@dataclass +class StepRecord: + """One motion step and the artifact base names captured for it.""" + + step_idx: int + state: dict[str, Any] + artifacts: set[str] = field(default_factory=set) + command: dict | None = None + result: dict | None = None + elapsed_s: float | None = None + extras: dict[str, Any] = field(default_factory=dict) + + def to_blob(self) -> dict[str, Any]: + blob: dict[str, Any] = { + "step_idx": self.step_idx, + "state": self.state, + "artifacts": sorted(self.artifacts), + } + if self.command is not None: + blob["command"] = self.command + if self.result is not None: + blob["result"] = self.result + if self.elapsed_s is not None: + blob["elapsed_s"] = self.elapsed_s + if self.extras: + blob["extras"] = self.extras + return blob + + @classmethod + def from_blob(cls, blob: dict[str, Any]) -> "StepRecord": + return cls( + step_idx=int(blob["step_idx"]), + state=dict(blob.get("state") or {}), + artifacts={str(name) for name in blob.get("artifacts") or []}, + command=blob.get("command"), + result=blob.get("result"), + elapsed_s=blob.get("elapsed_s"), + extras=dict(blob.get("extras") or {}), + ) + + +class EnvState: + """Own a run's step trace and all state-related files in its output root.""" + + def __init__(self, output_dir: Path | str): + self._output_dir = Path(output_dir) + self._output_dir.mkdir(parents=True, exist_ok=True) + self._steps: list[StepRecord] = [] + self._run_artifacts: set[str] = set() + self._open_count = 0 + self._next_step = 0 + self.reset() + + # -- private file resolution ----------------------------------------- + + @staticmethod + def _validate_name(name: str) -> str: + if not isinstance(name, str) or not name: + raise ValueError("artifact name must be a non-empty string") + path = Path(name) + if path.name != name or path.is_absolute() or name in {".", ".."}: + raise ValueError(f"artifact name must be a base filename: {name!r}") + if name == _MANIFEST_NAME: + raise ValueError(f"{_MANIFEST_NAME!r} is reserved") + if path.suffix.lower() not in _SUPPORTED_SUFFIXES: + raise ValueError(f"unsupported artifact suffix: {path.suffix or ''}") + return name + + def _artifact_file(self, name: str, step: int | None) -> Path: + name = self._validate_name(name) + if step is None: + return self._output_dir / name + if step < 0: + raise ValueError("artifact writes require a nonnegative step") + return self._output_dir / f"{step:02d}_{name}" + + def _manifest_file(self) -> Path: + return self._output_dir / _MANIFEST_NAME + + def _temporary_file(self, destination: Path) -> Path: + return destination.with_name( + f".{destination.stem}.tmp{destination.suffix}" + ) + + def _record_for(self, step: int) -> StepRecord: + """Return the live step record at ``step`` for in-place updates.""" + for record in self._steps: + if record.step_idx == step: + return record + raise KeyError(f"step {step} not present in state trace") + + # -- manifest -------------------------------------------------------- + + def _write_manifest(self) -> None: + destination = self._manifest_file() + temporary = self._temporary_file(destination) + manifest = { + "version": _MANIFEST_VERSION, + "run_artifacts": sorted(self._run_artifacts), + "steps": [record.to_blob() for record in self._steps], + } + try: + with temporary.open("w") as file: + json.dump(manifest, file, indent=2, default=_json_default) + os.replace(temporary, destination) + finally: + temporary.unlink(missing_ok=True) + + # -- lifecycle and counters ----------------------------------------- + + def reset(self) -> None: + """Remove state-owned artifacts and start a fresh trace.""" + self._output_dir.mkdir(parents=True, exist_ok=True) + for record in self._steps: + for name in record.artifacts: + self._artifact_file(name, record.step_idx).unlink(missing_ok=True) + for name in self._run_artifacts: + self._artifact_file(name, None).unlink(missing_ok=True) + self._manifest_file().unlink(missing_ok=True) + self._steps = [] + self._run_artifacts = set() + self._open_count = 0 + self._next_step = 0 + + @property + def next_step_idx(self) -> int: + return self._next_step + + @property + def latest_step(self) -> int | None: + if not self._steps: + return None + return self._steps[-1].step_idx + + def latest_record(self) -> StepRecord | None: + """Return the most recently recorded step (live reference, no copy).""" + return self._steps[-1] if self._steps else None + + def _resolve_read_step(self, step: int | None) -> int | None: + if step is None: + return None + if step == -1: + latest = self.latest_step + if latest is None: + raise LookupError("no steps available") + return latest + if step < 0: + raise ValueError("step must be -1, None, or nonnegative") + return step + + # -- generic artifact operations ------------------------------------ + + def save( + self, + name: str, + value: Any, + *, + step: int | None = -1, + **options: Any, + ) -> str | None: + """Serialize ``value`` according to ``name`` and return its base name. + + ``step`` defaults to ``-1`` (the most recently recorded step), so calls + made inside (or right after) a :meth:`record_step` block attach to that + step without an explicit index. Pass an ``int`` to target a specific + step, or ``None`` for a run-level artifact such as an episode video. + """ + step = self._resolve_read_step(step) + destination = self._artifact_file(name, step) + temporary = self._temporary_file(destination) + suffix = destination.suffix.lower() + record: StepRecord | None = ( + self._record_for(step) if step is not None else None + ) + try: + if suffix in _IMAGE_SUFFIXES: + array = np.asarray(value) + if array.dtype != np.uint8: + array = array.astype(np.uint8) + imageio.imwrite(temporary, array) + elif suffix == ".npy": + np.save(temporary, np.asarray(value)) + elif suffix == ".json": + with temporary.open("w") as file: + json.dump(value, file, indent=2, default=_json_default) + elif suffix == ".jsonl": + with temporary.open("w") as file: + if isinstance(value, str): + file.write(value) + else: + for item in value: + file.write( + json.dumps(item, default=_json_default) + "\n" + ) + elif suffix == ".mp4": + if isinstance(value, (bytes, bytearray, memoryview)): + temporary.write_bytes(bytes(value)) + else: + imageio.mimwrite( + temporary, + list(value), + fps=int(options.get("fps", 20)), + ) + elif suffix in _TEXT_SUFFIXES: + temporary.write_text(str(value)) + elif suffix == ".bin": + temporary.write_bytes(bytes(value)) + else: + raise ValueError(f"unsupported artifact suffix: {suffix}") + os.replace(temporary, destination) + if record is not None: + record.artifacts.add(name) + else: + self._run_artifacts.add(name) + self._write_manifest() + return name + except Exception as exc: + logger.warning("failed to save artifact %s: %s", name, exc) + return None + finally: + temporary.unlink(missing_ok=True) + + def load(self, name: str, *, step: int | None = -1) -> Any: + """Load an artifact; ``step=-1`` selects the latest recorded step.""" + resolved_step = self._resolve_read_step(step) + source = self._artifact_file(name, resolved_step) + suffix = source.suffix.lower() + if suffix in _IMAGE_SUFFIXES: + return imageio.imread(source) + if suffix == ".npy": + return np.load(source) + if suffix == ".json": + with source.open() as file: + return json.load(file) + if suffix == ".jsonl": + with source.open() as file: + return [json.loads(line) for line in file if line.strip()] + if suffix in _TEXT_SUFFIXES: + return source.read_text() + return source.read_bytes() + + def load_bytes(self, name: str, *, step: int | None = -1) -> bytes: + resolved_step = self._resolve_read_step(step) + return self._artifact_file(name, resolved_step).read_bytes() + + def exists(self, name: str, *, step: int | None = -1) -> bool: + try: + resolved_step = self._resolve_read_step(step) + except LookupError: + return False + return self._artifact_file(name, resolved_step).exists() + + def remove(self, name: str, *, step: int | None) -> bool: + destination = self._artifact_file(name, step) + if not destination.exists(): + return False + if step is None: + self._run_artifacts.discard(name) + else: + self._record_for(step).artifacts.discard(name) + destination.unlink() + self._write_manifest() + return True + + def list( + self, + pattern: str = "*", + *, + step: int | None = -1, + ) -> list[str]: + resolved_step = self._resolve_read_step(step) + if resolved_step is None: + names = [ + name + for name in self._run_artifacts + if self._artifact_file(name, None).exists() + ] + else: + record = self.get(resolved_step) + names = [ + name + for name in record.artifacts + if self._artifact_file(name, resolved_step).exists() + ] + return sorted(name for name in names if fnmatch.fnmatch(name, pattern)) + + def list_all(self, pattern: str = "*") -> list[tuple[int | None, str]]: + artifacts: list[tuple[int | None, str]] = [ + (None, name) for name in self.list(pattern, step=None) + ] + for record in self._steps: + artifacts.extend( + (record.step_idx, name) + for name in record.artifacts + if fnmatch.fnmatch(name, pattern) + and self._artifact_file(name, record.step_idx).exists() + ) + return sorted( + artifacts, + key=lambda item: (-1 if item[0] is None else item[0], item[1]), + ) + + # -- step records ---------------------------------------------------- + + @contextmanager + def record_step( + self, + *, + state: dict[str, Any], + command: dict | None = None, + result: dict | None = None, + elapsed_s: float | None = None, + extras: dict[str, Any] | None = None, + ) -> Iterator[int]: + """Append a new step record and yield its index. + + The step is committed to the trace immediately, so any subsequent + :meth:`save` without an explicit ``step`` attaches to it. If the block + raises, the step and any artifacts written to it are rolled back. + """ + if self._open_count: + raise RuntimeError("a step record is already open") + record = StepRecord( + step_idx=self._next_step, + state=copy.deepcopy(state), + command=copy.deepcopy(command), + result=copy.deepcopy(result), + elapsed_s=elapsed_s, + extras=copy.deepcopy(extras or {}), + ) + self._steps.append(record) + self._next_step = record.step_idx + 1 + self._open_count += 1 + self._write_manifest() + try: + yield record.step_idx + except BaseException: + for name in list(record.artifacts): + self._artifact_file(name, record.step_idx).unlink(missing_ok=True) + if self._steps and self._steps[-1] is record: + self._steps.pop() + self._next_step = record.step_idx + try: + self._write_manifest() + except Exception as exc: + logger.warning( + "failed to rewrite manifest after step rollback: %s", exc + ) + raise + finally: + self._open_count -= 1 + + def get(self, step: int = -1) -> StepRecord: + resolved_step = self._resolve_read_step(step) + if resolved_step is None: + raise ValueError("step records are not run-level artifacts") + for record in self._steps: + if record.step_idx == resolved_step: + return copy.deepcopy(record) + raise KeyError(f"step {resolved_step} not present in state trace") + + def records(self) -> list[StepRecord]: + return copy.deepcopy(self._steps) + + # -- LLM-facing view ------------------------------------------------- + + def view( + self, + step: int = -1, + *, + image_slots: dict[str, str] | None = None, + ) -> dict[str, Any]: + try: + record = self.get(step) + except Exception as exc: + return {"error": f"state step not available: {exc}"} + + metadata: dict[str, Any] = {} + metadata_suffix = "_metadata.json" + for name in sorted(record.artifacts): + if not name.endswith(metadata_suffix): + continue + key = name.removesuffix(metadata_suffix) + try: + loaded = self.load(name, step=record.step_idx) + if isinstance(loaded, dict): + loaded = { + field: value + for field, value in loaded.items() + if field not in {"K", "T_base_cam"} + } + metadata[key] = loaded + except Exception as exc: + metadata[key] = {"error": str(exc)} + + out: dict[str, Any] = { + "step": record.step_idx, + "state": record.state, + "artifacts": sorted(record.artifacts), + "camera_meta": metadata, + "log": { + "command": record.command, + "result": record.result, + "elapsed_s": record.elapsed_s, + }, + } + if record.extras: + out["extras"] = record.extras + if image_slots: + for slot, name in image_slots.items(): + if name not in record.artifacts: + continue + try: + out[slot] = self.load_bytes(name, step=record.step_idx) + except FileNotFoundError: + continue + return out diff --git a/rpent/tools/toolkit.py b/rpent/tools/toolkit.py index 31a7cbca..76882f81 100644 --- a/rpent/tools/toolkit.py +++ b/rpent/tools/toolkit.py @@ -9,14 +9,18 @@ import base64 import json import threading +import time import traceback from collections.abc import Callable from dataclasses import dataclass, field -from typing import Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar -from rpent.dashboard.events import DashboardEventSink, ToolResultEvent +from rpent.dashboard.events import DashboardEventSink, StepRecordEvent from rpent.utils.templates import substitute +if TYPE_CHECKING: + from rpent.tools.state import StepRecord + @dataclass(slots=True) class _ToolOperation: @@ -28,6 +32,37 @@ class ToolCancelled(Exception): """Raised when an environment reaches a safe cancellation boundary.""" +def updatestate(func): + """Mark a tool handler as one that advances environment state. + + The toolkit captures a fresh observation (:meth:`Toolkit.get_env_state`) + after a handler carrying this marker runs (or raises). Apply it to + primitive-driver methods that move the robot; read-only tools (file IO, + ``view_env_state``, ``back_project``, ``segment``, ...) are left + unmarked and skip state capture. + + The marker is read off the underlying function, so registering a bound + method (``getattr(self._driver, name)``) inherits it automatically. + ``functools.partial`` does not delegate attribute access, so a partial + wrapping a marked method is treated as read-only (intentional -- see the + LIBERO ``segment`` tool). + """ + func._updates_state = True + return func + + +def _updates_state(handler: Callable[..., Any]) -> bool: + """Whether ``handler`` was marked with :func:`updatestate`. + + Resolves through ``__func__`` so bound methods (the common case for + primitive-driver tools) report the marker set on their underlying + function. Plain functions, closures, and ``functools.partial`` objects + do not delegate, so undecorated read-only handlers read as ``False``. + """ + target = getattr(handler, "__func__", handler) + return bool(getattr(target, "_updates_state", False)) + + @dataclass class ToolResult: """Result of executing one tool call. @@ -106,10 +141,18 @@ class Toolkit: :meth:`close` to release env-side primitives / servers at the end of the run. """ - def __init__(self, *, dashboard_events: DashboardEventSink) -> None: - # name -> (spec, handler) - self._tools: dict[str, tuple[dict[str, Any], Callable[..., dict[str, Any]]]] = {} + def __init__( + self, + *, + dashboard_events: DashboardEventSink, + state: Any = None, + ) -> None: + self._tools: dict[ + str, + tuple[dict[str, Any], Callable[..., Any], bool], + ] = {} self._dashboard_events = dashboard_events + self._state = state self._operation_lock = threading.Lock() self._active_operation: _ToolOperation | None = None self._register_common_tools() @@ -122,7 +165,7 @@ def add_tool( self, name: str, spec: dict[str, Any], - handler: Callable[..., dict[str, Any]], + handler: Callable[..., Any], ) -> None: """Register one tool under ``name`` with its schema and handler. @@ -131,9 +174,10 @@ def add_tool( spec: Anthropic-shaped tool schema dict (``name``, ``description``, ``input_schema``). handler: Callable invoked with the tool's input kwargs; returns - a result dict. + a result dict. Decorate state-advancing primitives with + :func:`updatestate`. """ - self._tools[name] = (spec, handler) + self._tools[name] = (spec, handler, _updates_state(handler)) def _register_common_tools(self) -> None: """Register the file/IO tools shared by every run.""" @@ -150,39 +194,103 @@ def _register_common_tools(self) -> None: def get_tools_spec(self) -> list[dict[str, Any]]: """Return the tool schemas the LLM sees.""" return substitute( - [spec for spec, _ in self._tools.values()] + [spec for spec, _, _ in self._tools.values()] ) def execute_tool(self, name: str, input_dict: dict[str, Any]) -> ToolResult: """Dispatch a tool call to its registered handler.""" entry = self._tools.get(name) if entry is None: - return ToolResult(name=name, result={"error": f"unknown tool: {name}"}) - handler = entry[1] + return self._finish_tool(name, {"error": f"unknown tool: {name}"}) + _, handler, captures_state = entry with self._operation_lock: if self._active_operation is not None: - return ToolResult( - name=name, - result={"error": "another tool operation is still active"}, + return self._finish_tool( + name, + {"error": "another tool operation is still active"}, ) operation = _ToolOperation() self._active_operation = operation try: + started = time.perf_counter() + failed = False try: result = handler(**input_dict) except TypeError as e: - result = {"error": f"bad arguments for {name}: {e}", "got": input_dict} + return self._finish_tool( + name, + {"error": f"bad arguments for {name}: {e}", "got": input_dict}, + ) + except ToolCancelled as e: + result = { + "error": str(e), + "code": "tool_cancelled", + "interrupted": True, + } + failed = True except Exception as e: result = {"error": str(e), "traceback": traceback.format_exc()} - self._dashboard_events.emit(ToolResultEvent(name=name, result=result)) - return ToolResult(name=name, result=result) + failed = True + + if captures_state: + elapsed_s = round(time.perf_counter() - started, 2) + result_dict = result if isinstance(result, dict) else {"value": result} + command = {"action": name, **input_dict} + record: StepRecord | None = None + try: + captured = self.get_env_state( + command=command, + result=result_dict, + elapsed_s=elapsed_s, + ) + except Exception as e: + captured = result_dict + captured["state_capture_error"] = str(e) + captured.setdefault( + "error", f"failed to capture state after {name}: {e}" + ) + captured.setdefault("traceback", traceback.format_exc()) + else: + record = self._state.latest_record() + result = captured + if failed: + result.setdefault("error", result_dict["error"]) + if "traceback" in result_dict: + result.setdefault("traceback", result_dict["traceback"]) + if record is not None: + self._publish_step(record) + + return self._finish_tool(name, result) finally: with self._operation_lock: self._active_operation = None operation.done_event.set() + def _publish_step(self, record: StepRecord) -> None: + """Publish one recorded environment step to the dashboard sink.""" + self._dashboard_events.emit( + StepRecordEvent( + record=record, + env_state=self._state, + frame_artifacts=dict(getattr(type(self), "_FRAME_ARTIFACTS", {})), + ) + ) + + def _finish_tool(self, name: str, result: Any) -> ToolResult: + return ToolResult(name=name, result=result) + + def get_env_state( + self, + *, + command: dict[str, Any], + result: dict[str, Any], + elapsed_s: float, + ) -> dict[str, Any]: + """Capture and return the observation produced by a stateful tool.""" + raise NotImplementedError + # ------------------------------------------------------------------ # Server lifecycle hooks (overridden by env toolkits) # ------------------------------------------------------------------ diff --git a/rpent/utils/sam3_client.py b/rpent/utils/sam3_client.py index f8a80037..22c23977 100644 --- a/rpent/utils/sam3_client.py +++ b/rpent/utils/sam3_client.py @@ -5,7 +5,6 @@ import base64 import io from dataclasses import dataclass -from pathlib import Path from typing import Any import imageio.v2 as imageio @@ -35,7 +34,7 @@ def __init__(self, client: RpcClient, *, timeout_s: float = 120.0) -> None: def segment( self, - image_path: str | Path, + image: bytes | bytearray | memoryview | np.ndarray, *, text_prompt: str | None = None, point: list[int] | None = None, @@ -55,7 +54,12 @@ def segment( if not 0.0 <= float(min_score) <= 1.0: raise ValueError("min_score must be between 0 and 1") - image_bytes = Path(image_path).read_bytes() + if isinstance(image, np.ndarray): + buffer = io.BytesIO() + imageio.imwrite(buffer, image, format="png") + image_bytes = buffer.getvalue() + else: + image_bytes = bytes(image) body: dict[str, Any] = { "image_base64": base64.b64encode(image_bytes).decode("ascii"), "min_score": float(min_score), diff --git a/scripts/codex_proxy/litellm_callbacks.py b/scripts/codex_proxy/litellm_callbacks.py index 23aca5f2..5d54e05b 100644 --- a/scripts/codex_proxy/litellm_callbacks.py +++ b/scripts/codex_proxy/litellm_callbacks.py @@ -92,7 +92,7 @@ def _extract_mcp_namespace_map(data: dict) -> dict[str, str]: Example:: {"back_project": "mcp__rpent", - "view_driver_state": "mcp__rpent"} + "view_env_state": "mcp__rpent"} """ tools = data.get("tools") if not isinstance(tools, list):