From ae68d673c9506b4a540fa2c07f8c2a310925705d Mon Sep 17 00:00:00 2001 From: user Date: Thu, 17 Sep 2026 14:37:16 +0800 Subject: [PATCH 1/2] refactor: migrate tools and planners to the native tool protocol Preserve the complete implementation from RLinf/RPent#173, including follow-up fixes and merge resolutions. Original head: da4abb0c0ee34f97c6ffb7b911f1e53fff6fb4e6. --- .../rst_source/development/add_primitive.rst | 175 +- .../rst_source/development/add_robot.rst | 124 +- .../rst_source/development/architecture.rst | 51 +- .../rst_source/development/interfaces.rst | 124 +- .../rst_source/development/memory.rst | 16 + .../rst_source/usage/configure_planner.rst | 24 +- .../rst_source/development/add_primitive.rst | 151 +- .../rst_source/development/add_robot.rst | 116 +- .../rst_source/development/architecture.rst | 45 +- .../rst_source/development/interfaces.rst | 93 +- .../rst_source/development/memory.rst | 9 + .../rst_source/usage/configure_planner.rst | 22 +- .../rst_source/usage/dual_franka.rst | 2 +- pyproject.toml | 3 +- robots/dual_franka/dual_franka_manual_call.py | 93 +- robots/dual_franka/perception.py | 3 - robots/dual_franka/toolkit.py | 181 +- robots/dual_franka/tools.py | 1267 ++++---- robots/franka/perception.py | 4 - robots/franka/toolkit.py | 118 +- robots/franka/tools.py | 429 ++- robots/libero/flash/replay.py | 83 +- robots/libero/robot_spec.py | 1 + robots/libero/toolkit.py | 603 +++- robots/libero/tools.py | 2684 ++++++----------- robots/robocasa/primitives.py | 553 ---- robots/robocasa/rldx_skill.py | 42 +- robots/robocasa/robot_spec.py | 6 +- robots/robocasa/toolkit.py | 394 ++- robots/robocasa/tools.py | 1480 +++++---- robots/robotwin/primitives.py | 428 --- robots/robotwin/robot_spec.py | 3 +- robots/robotwin/toolkit.py | 423 +-- robots/robotwin/tools.py | 947 +++--- rpent/dashboard/server.py | 16 +- rpent/dashboard/state.py | 73 +- rpent/memory/manager.py | 77 +- rpent/memory/tools.py | 167 - rpent/planner/api_loop.py | 205 +- rpent/planner/base.py | 10 +- rpent/planner/claude_code.py | 81 +- rpent/planner/codex.py | 66 +- rpent/planner/utils/http_mcp_server.py | 88 +- rpent/session/base.py | 2 +- rpent/tools/__init__.py | 19 +- rpent/tools/base.py | 230 ++ rpent/tools/common.py | 168 -- rpent/tools/common_tools.py | 104 + rpent/tools/toolkit.py | 449 +-- tests/README.md | 4 + tests/e2e_tests/common.py | 2 +- tests/e2e_tests/robocasa/scenario.py | 4 +- tests/unit_tests/robots/conftest.py | 93 - .../fixtures/pre_native_tool_contracts.json | 527 ++++ .../dual_franka/test_dual_franka_tools.py | 177 +- .../robots/dual_franka/test_exploration.py | 52 +- .../fixtures/common_tool_contracts.json | 60 + tests/unit_tests/robots/franka/_fakes.py | 74 + .../fixtures/pre_native_tool_contracts.json | 408 +++ tests/unit_tests/robots/franka/test_tools.py | 223 +- tests/unit_tests/robots/libero/conftest.py | 138 + .../fixtures/pre_native_tool_contracts.json | 417 +++ tests/unit_tests/robots/libero/test_flash.py | 370 +++ .../robots/libero/test_libero_integration.py | 105 +- .../libero/test_libero_toolkit_contracts.py | 478 ++- tests/unit_tests/robots/robocasa/conftest.py | 199 ++ .../fixtures/pre_native_tool_contracts.json | 482 +++ .../robots/robocasa/test_memory_contracts.py | 26 +- .../test_robocasa_toolkit_contracts.py | 341 ++- .../robocasa/test_vla_protocol_contracts.py | 331 +- tests/unit_tests/robots/robotwin/conftest.py | 158 + .../fixtures/pre_native_tool_contracts.json | 252 ++ .../test_robotwin_tool_schema_contracts.py | 39 +- .../test_robotwin_toolkit_contracts.py | 553 ++-- .../robots/test_tool_schema_contracts.py | 182 +- .../robots/test_toolkit_contracts.py | 73 +- .../rpent/cli/test_main_contracts.py | 19 +- .../rpent/dashboard/test_state_contracts.py | 114 +- .../rpent/memory/test_authorization.py | 103 + .../rpent/planner/_native_helpers.py | 121 + tests/unit_tests/rpent/planner/conftest.py | 31 + .../rpent/planner/test_api_contracts.py | 170 +- .../rpent/planner/test_claude_contracts.py | 212 +- .../rpent/planner/test_codex_contracts.py | 121 +- .../rpent/planner/test_http_mcp_server.py | 198 +- .../rpent/planner/test_native_adapters.py | 116 + .../rpent/tools/test_common_tools.py | 211 ++ .../rpent/tools/test_human_in_the_loop.py | 2 +- .../rpent/tools/test_native_protocol.py | 209 ++ .../unit_tests/rpent/tools/test_scheduling.py | 61 + .../rpent/tools/test_toolkit_contracts.py | 709 ++--- 91 files changed, 11591 insertions(+), 9026 deletions(-) delete mode 100644 robots/robocasa/primitives.py delete mode 100644 robots/robotwin/primitives.py delete mode 100644 rpent/memory/tools.py create mode 100644 rpent/tools/base.py delete mode 100644 rpent/tools/common.py create mode 100644 rpent/tools/common_tools.py create mode 100644 tests/unit_tests/robots/dual_franka/fixtures/pre_native_tool_contracts.json create mode 100644 tests/unit_tests/robots/fixtures/common_tool_contracts.json create mode 100644 tests/unit_tests/robots/franka/_fakes.py create mode 100644 tests/unit_tests/robots/franka/fixtures/pre_native_tool_contracts.json create mode 100644 tests/unit_tests/robots/libero/conftest.py create mode 100644 tests/unit_tests/robots/libero/fixtures/pre_native_tool_contracts.json create mode 100644 tests/unit_tests/robots/libero/test_flash.py create mode 100644 tests/unit_tests/robots/robocasa/conftest.py create mode 100644 tests/unit_tests/robots/robocasa/fixtures/pre_native_tool_contracts.json create mode 100644 tests/unit_tests/robots/robotwin/conftest.py create mode 100644 tests/unit_tests/robots/robotwin/fixtures/pre_native_tool_contracts.json create mode 100644 tests/unit_tests/rpent/memory/test_authorization.py create mode 100644 tests/unit_tests/rpent/planner/_native_helpers.py create mode 100644 tests/unit_tests/rpent/planner/conftest.py create mode 100644 tests/unit_tests/rpent/planner/test_native_adapters.py create mode 100644 tests/unit_tests/rpent/tools/test_common_tools.py create mode 100644 tests/unit_tests/rpent/tools/test_native_protocol.py create mode 100644 tests/unit_tests/rpent/tools/test_scheduling.py diff --git a/docs/source-en/rst_source/development/add_primitive.rst b/docs/source-en/rst_source/development/add_primitive.rst index b3b68971e..4d92527c7 100644 --- a/docs/source-en/rst_source/development/add_primitive.rst +++ b/docs/source-en/rst_source/development/add_primitive.rst @@ -28,58 +28,66 @@ Two types of primitives - ``move_to``, ``rotate_wrist``, ``release``, ``back_project`` -From the LLM's perspective, both types expose the same interface: a -tool schema, a primitives method, and a state dump after the -call. They differ only in how the method is implemented. +Both types are native tools: a typed handler receives the session's resources +through ``ToolContext`` and returns ``ToolResult``. The toolkit validates +arguments and captures the post-action observation. Add a scripted primitive ------------------------ -Adding a scripted primitive usually involves two steps: +Define a module-level handler in ``robots//tools.py`` and add the +resulting ``Tool`` to the robot's tool tuple. For example, this LIBERO handler +holds the current pose for a bounded number of environment steps: -1. **Add a method to the primitives.** Add the method to the - current robot's primitives class, such as - ``LiberoPrimitives`` or ``MyRobotPrimitives``. The method accepts - the tool-call arguments, performs the work, usually through one or - more ``self._env.step(...)`` calls, and returns a small log ``dict``. - - Primitive methods capture and re-render state (``get_env_state``) - automatically after they run: - - .. code-block:: python - - 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} - - You can mark read-only tools (``view_env_state``, ``back_project``, ``segment``, - ...) with :func:`~rpent.tools.toolkit.readonly` so the toolkit skips state - capture for them, improving performance. - -2. **Add the tool schema.** Add an entry to ``TOOLS_SPEC`` in - ``robots//tools.py``: - - .. code-block:: python - - { - "name": "open_drawer", - "description": "Pull the currently-grasped drawer handle " - "backwards by ``dx`` meters.", - "input_schema": { - "type": "object", - "properties": {"dx": {"type": "number"}}, - "required": [], - }, - } - -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)``). +.. code-block:: python -After these steps, the ``api``, ``claude_code``, and ``codex`` planners -can all call the primitive without any other code changes. + from typing import Annotated + + from pydantic import Field + + from rpent.tools import ToolContext, ToolResult, tool + + @tool + def hold_pose( + steps: Annotated[int, Field(ge=1, le=100)] = 10, + *, + ctx: ToolContext, + ) -> ToolResult: + """Hold the current pose with the gripper closed. + + Args: + steps: Number of environment steps. + """ + runtime = ctx.robot + for _ in range(steps): + ctx.check_cancelled() + obs, _, terminated, truncated, _ = runtime.env.step( + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0] + ) + runtime.executed_steps += 1 + runtime.set_obs(obs) + ctx.record_frame(obs["main_images"]) + if terminated or truncated: + break + return ToolResult(data={"steps_requested": steps}) + + # Add hold_pose to the existing LIBERO_TOOLS tuple. + +Use the concrete runtime type in ``ToolContext[LiberoRuntime]`` in robot code. +``@tool`` generates the parameter model and JSON schema from the same function +signature. Google-style docstrings provide descriptions; use ``Annotated`` / +``Field`` for constraints on model-supplied arguments. ``ctx`` is injected by +the executor and is not part of the published schema. + +After adding the declaration to ``LIBERO_TOOLS``, all three planners can call +it. The toolkit handles capture through ``_capture_observation``; handlers +submit frames but do not save their own episode video or duplicate the state dump. + +For a tool that reads existing observations, place ``@readonly`` below +``@tool`` to skip automatic capture. Calls still execute one at a time. +``write_text_file`` and ``finish`` have no readonly marker and run exclusively; +the executor skips observation capture for common tools and ``finish``. +See :doc:`interfaces` for execution and cancellation. .. _add-primitive-model-based: @@ -113,44 +121,26 @@ primitive requires a few additional components: ``rpent.robots.components.pi05_vla_client.Pi05VLAClient`` for the LIBERO implementation. -3. **Add a method to the primitives.** In the current - robot's primitives class, call the model client, pass - the returned action chunk to the environment, and return a log - ``dict``. The model client API is - :meth:`rpent.robots.components.pi05_vla_client.Pi05VLAClient.predict`, - which reads the instruction from ``env_obs["task_descriptions"]`` and - returns a ``[chunk, action_dim]`` numpy action chunk (batch dim already - stripped): - - .. code-block:: python - - def mymodel_pick(self, target: str) -> dict: - env_obs = self._env.get_obs() - env_obs["task_descriptions"] = f"pick {target}" - chunk = self._model.predict(env_obs) - self._env.chunk_step(chunk) - return {"model": "mymodel", "target": target} - -4. **Add the tool schema and register it in the toolkit.** Follow the - same pattern as for a scripted primitive. - -5. **Wire the components together in ``robot_spec.py``.** The - robot's ``get_toolkit`` builds the toolkit with - ``runtime_kwargs``: - - .. code-block:: python - - def get_toolkit(*, runtime_kwargs, dashboard_events): - from robots.myrobot.toolkit import MyRobotToolkit - return MyRobotToolkit( - runtime_kwargs=runtime_kwargs, - dashboard_events=dashboard_events, - ) - - The robot package's ``_init_runtime`` builds - ``runtime_kwargs``, for example - ``{"env": MyRobotEnvClient(...), "model": MyModelClient(...)}``. - The toolkit constructor then forwards it to the primitives. +3. **Write a native tool handler.** Use ``ctx.robot`` to access the model and + environment clients, request a prediction, execute the returned actions, + and return ``ToolResult(data={...})``. Follow the robot's observation and + action conventions: Pi0.5 reads the instruction from + ``env_obs["task_descriptions"]`` and returns a ``[chunk, action_dim]`` + NumPy array. Check cancellation before inference and at safe action + boundaries, and submit environment frames through ``ctx.record_frame``. + See ``pi0_pick`` in ``robots/libero/tools.py`` and ``rldx_skill`` in + ``robots/robocasa/tools.py`` for concrete implementations. + +4. **Add the tool to the robot's tuple.** The toolkit receives that tuple in + its constructor and handles observation capture after execution, just as + for a scripted primitive. + +5. **Wire the clients in ``robot_spec.py``.** ``_init_runtime`` returns + ``(owned_daemons, runtime_kwargs)``, with entries such as ``env`` and + ``model``. ``get_toolkit(*, runtime_kwargs, dashboard_events, config)`` + passes those inputs, ``config.output_dir``, and a ``MemoryManager`` to the + robot toolkit. The toolkit constructs the session runtime from + ``runtime_kwargs``. See :doc:`add_robot` for the complete factory. Reuse an existing vla_server across runs ---------------------------------------- @@ -190,7 +180,7 @@ parts: the server-side handler by the facade — the client does **not** pass it, and must not forge ``session_ids`` inside ``predict``'s ``options``. -- **Primitives side**: call ``reset_session`` before a task starts to clear +- **Runtime / tool side**: call ``reset_session`` before a task starts to clear policy state left over from the previous episode, so consecutive runs do not leak state into each other. @@ -234,11 +224,12 @@ Design principles for a new primitive - **Tools describe intent, not motion.** A good tool name is ``pi0_pick``, not ``execute_action_chunk_of_length_20``. -- **Every tool ends with a state dump.** The next turn depends on - 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. Save larger observations through ``EnvState.save``; ``EnvState`` +- **Action calls include a fresh observation.** The toolkit captures it after + the handler finishes and before returning to the planner. Readonly tools + reuse recorded observations. +- **Return small ``ToolResult.data`` payloads.** The planner serializes them + as text and sends ``ToolResult.images`` as PNG content. 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. @@ -262,5 +253,5 @@ The same pattern extends to non-VLA model primitives: head to call via a ``model`` kwarg on ``predict``. Regardless of the implementation, the framework contract remains -unchanged: model process → model client → primitives method → -tool schema → ``Toolkit.add_tool``. +unchanged: model process → model client → native ``@tool`` handler → +robot tool tuple → ``Toolkit.execute_tool``. diff --git a/docs/source-en/rst_source/development/add_robot.rst b/docs/source-en/rst_source/development/add_robot.rst index 89e7ac194..23742d422 100644 --- a/docs/source-en/rst_source/development/add_robot.rst +++ b/docs/source-en/rst_source/development/add_robot.rst @@ -57,7 +57,8 @@ For a new robot named ``myrobot``, use the following directory layout: robot_spec.py # RobotSpec, factories, Dashboard spec, runtime hooks env_client.py # MyEnvClient — agent-side RPC stub (§1) prompt_bundle.py # system()/user() prompt factories (§2) - toolkit.py # MyRobotToolkit + primitives + tool definitions (§3) + toolkit.py # Runtime, MyRobotToolkit, observations (§3) + tools.py # Native tool declarations and handlers (§3) env_server.py # server-side facade + RPC server (§1) vla_server.py # (optional) VLA model server @@ -100,6 +101,7 @@ these two functions: from robots.myrobot.toolkit import MyRobotToolkit return MyRobotToolkit( runtime_kwargs=runtime_kwargs, + output_dir=config.output_dir, dashboard_events=dashboard_events, memory=MemoryManager( root=config.prompt_vars.get("memory_dir") or get_memory_dir("myrobot"), @@ -279,69 +281,69 @@ render time. 3. ``toolkit.py`` ------------------ -This module owns everything the LLM can call: the tool schemas, the primitives, -the per-step state dump, and the MCP allowlist. (In the LIBERO robot these -are split between ``tools.py`` and ``toolkit.py`` for historical reasons; for a -new robot it is fine to keep them all in ``toolkit.py``.) - -A toolkit module typically contains four pieces: - -**Primitives class** (e.g. ``MyRobotPrimitives``) — a Python object owned -by the toolkit. It holds the ``EnvClient``, the VLA ``model`` client, and any -state needed for the current run. It exposes one method per primitive tool -(``move_to``, ``pi0_pick``, ``release``, …), with each method returning a -``dict`` log. - -**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_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``: - -- forward ``memory`` (a :class:`~rpent.memory.MemoryManager`) and ``state`` to - ``super().__init__(...)``. Configure ``memory_access`` and - ``inbox_cell_tag`` on the ``MemoryManager``; eval uses read-only access by - default. -- build the primitives in ``__init__`` through a custom initialization - helper (named ``init_primitives`` 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_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 save remaining agent-side artifacts through - ``EnvState`` (for example ``state.save("episode.mp4", frames, step=None)``). - -``runtime_kwargs`` (forwarded from ``robot_spec.py:get_toolkit``) is the dict -the toolkit passes verbatim to your primitives' ``__init__`` — typically -``{"env": MyEnvClient(...), "model": VLAClient(...), ...}``. +Split robot behavior between ``tools.py`` and ``toolkit.py``, following +``robots/libero/``, ``robots/robocasa/``, or ``robots/robotwin/``: + +**Runtime** — a per-session object such as ``MyRobotRuntime`` holds the env +and model clients, cached observations, and task progress. Construct it from +``runtime_kwargs`` returned by ``init_runtime``. Tools access it through +``ctx.robot``. + +**Native tools** — define typed functions in ``tools.py`` with ``@tool`` and a +required keyword-only ``ctx: ToolContext[MyRobotRuntime]``. Return +``ToolResult`` and collect the declarations in a tuple such as +``MYROBOT_TOOLS``. Tool schemas come from signatures and Google-style +docstrings; planner adapters handle SDK and MCP serialization. See +:doc:`add_primitive` for an example. + +**Observations** — ``dump_state(runtime, env_state, log)`` opens +``env_state.record_step(...)`` to allocate and commit a ``StepRecord``. +Save arrays, images, and metadata through ``env_state.save(...)``: inside +the block, omitting ``step`` targets the new step; an explicit integer targets +another step and ``step=None`` creates a run-level artifact. ``EnvState`` +adds each saved logical name to the record's ``artifacts`` set. +``build_observation(state, record)`` returns JSON data and ordered PNG bytes +for both action responses and ``view_env_state``. + +**Toolkit** — subclass ``Toolkit[MyRobotRuntime]``: + +- Construct the runtime and a fresh ``EnvState(output_dir)``; pass ``state``, + ``memory``, ``robot=runtime``, ``output_dir``, ``tools=MYROBOT_TOOLS``, and + ``dashboard_events`` to ``super().__init__``. Configure ``memory_access`` and + ``inbox_cell_tag`` on ``MemoryManager``; evaluation defaults to read-only memory. +- Initialize the environment and record step 0. Publish its ``StepRecordEvent`` + so the Dashboard can display the initial observation. +- Implement ``_capture_observation(*, command, result, elapsed_s)`` using + ``dump_state`` and ``build_observation``. Save the original ``result.to_dict()`` + in the step log. Return an observation containing the step, artifact names, + state, and action log, together with its PNG images. The base executor calls + this after non-readonly robot handlers, including failed handlers, except + ``finish``. Common tools also skip capture. +- Implement ``solved()`` using the environment's success criterion. Provide a + robot-specific ``finish`` tool returning ``status`` and ``summary``; the + executor stores accepted results in ``toolkit.finish_result``. + +Tools submit environment-step RGB frames via ``ctx.record_frame(rgb)``. +The base toolkit collects the frames and signals cancellation. Robot toolkits +save Dashboard action clips during observation capture, override ``close()`` +to save the episode video, and implement ``write_recipe()`` from their trace. Conventions worth keeping ------------------------- -- ``output_dir`` is the working directory that the runner creates for each - 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_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. +- ``output_dir`` is the runner-created working directory. Environment artifacts + are managed by ``EnvState`` through logical names; transcripts share the same + run directory. +- ``@readonly`` below ``@tool`` skips automatic observation capture. Calls + execute one at a time. Call + ``ctx.check_cancelled()`` at safe boundaries in long loops. +- Server-side values use transport-supported Python / NumPy types and remain + torch-free. +- Add new observation modalities in ``dump_state`` and expose them through + ``build_observation``. ``view_env_state`` reads a saved step without stepping + or re-rendering the environment. + +See :doc:`interfaces` for the full tool execution and lifecycle contract. .. _add-robot-config: @@ -404,7 +406,7 @@ validates those fields and returns a by this process. The active runner stops them during cleanup. A client for an external endpoint must not add that external service to this list. - ``runtime_kwargs: dict`` is passed to the toolkit constructor, which - forwards it to the primitives' ``__init__``. A complete set commonly + uses it to construct the robot runtime. A complete set commonly contains ``{"env": MyEnvClient(...), "model": VLAClient(...)}`` plus any supporting clients. diff --git a/docs/source-en/rst_source/development/architecture.rst b/docs/source-en/rst_source/development/architecture.rst index b38485ecc..680e7433d 100644 --- a/docs/source-en/rst_source/development/architecture.rst +++ b/docs/source-en/rst_source/development/architecture.rst @@ -57,7 +57,7 @@ A single run is an LLM-in-the-loop cycle: 1. The LLM reasons about the task and calls a tool (e.g. ``pi0_pick``). -2. The tool's primitives requests an action from the ``vla_server`` +2. The native tool handler requests an action from the ``vla_server`` (``predict``). 3. The ``env_server`` executes the action. 4. The environment returns updated observations and camera frames. @@ -81,12 +81,15 @@ The framework code is organized by responsibility: context/ # Prompt utilities and shared prompt sections. dashboard/ # FastAPI monitor + SSE streams (optional). robots/ # RobotSpec, PromptBundle, and on-demand robot loading. - tools/ # Toolkit base class and shared tool helpers. + tools/ # Native tool protocol, executor, and common tools. + session/ # EnvState, step records, and artifact storage. + memory/ # Memory sync, access control, and exploration merge. utils/ # Config, logging, RPC, and VLA client helpers. robots/ libero/ # LIBERO env_client / env_server / vla_server / # toolkit / prompt_bundle. The reference robot. robocasa/ # RoboCasa robot (RLDX-1 VLA, kitchen tasks). + robotwin/ # RoboTwin robot (LingBot-VLA, dual-arm tasks). (franka/) # Franka robot — in progress. (so101/) # SO-101 robot — in progress. scripts/ @@ -152,7 +155,7 @@ two factories exposed by that package: # robots/myrobot/__init__.py def get_robot_spec() -> RobotSpec: ... # identity, prompt bundle, and runner hooks def get_toolkit( - *, runtime_kwargs, dashboard_events + *, runtime_kwargs, dashboard_events, config ): ... ``RobotSpec`` gathers the robot's identity, prompt templates, optional @@ -160,25 +163,27 @@ Dashboard description, and three runner hooks (``add_cli_args`` / ``parse_config`` / ``init_runtime``). See :doc:`interfaces` for what each field must provide. -The loader itself does not maintain a list of robot names. The -current CLI restricts ``--robot`` to ``libero`` and ``robocasa``; adding a -new name therefore also requires updating the CLI choices. See -:doc:`add_robot` for the complete procedure. +The loader discovers robot packages on disk. The CLI obtains its ``--robot`` +choices from ``enumerate_robots()``; LIBERO, RoboCasa, and RoboTwin are available. +See :doc:`add_robot` for the complete procedure. Planner, Toolkit, and RPC transports ------------------------------------- -These three layers stay decoupled, each owning one segment of the path. The -planner only pulls the tool list via ``get_tools_spec`` and invokes tools with -``execute_tool``, indifferent to whether a tool is scripted or a VLA. The -toolkit translates each tool call into a primitive call, and the primitives -issues ``reset`` / ``step`` / ``predict`` requests to ``env_server`` / -``vla_server`` over RPC. The RPC transport (HTTP or socket) only ferries those -calls and their NumPy observations across processes, transparent to the layers -above. That is why swapping the planner leaves the tools untouched, and -swapping the transport leaves the planner untouched. The concrete interface -contracts (``Planner.solve``, ``Toolkit.add_tool``, ``RpcFacade._dispatch``) -are collected in :doc:`interfaces`. +The planner reads native ``Tool`` declarations through ``list_tools()`` and +invokes them through ``execute_tool``. It adapts schemas and ``ToolResult`` +text / PNG images to its SDK; Claude Code and Codex use MCP adapters at this +boundary. Robot handlers are independent of the planner transport. + +``Toolkit`` validates arguments, admits one call, injects ``ToolContext``, and +captures post-action observations. Handlers use ``ctx.robot`` to access their +session runtime and issue ``reset`` / ``step`` / ``predict`` requests through +environment and model clients. HTTP or socket RPC carries those calls and +NumPy observations between processes. + +The toolkit also owns cancellation, the frame buffer, and ``finish_result``. +Robot-specific subclasses build observations, save videos and recipes, and +report native success through ``solved()``. See :doc:`interfaces` for the contracts. Dashboard (optional) -------------------- @@ -190,12 +195,12 @@ frontend. With ``--dashboard``, ``rpent/cli/main.py`` hands control to the CLI before it calls ``robot_spec.init_runtime`` once with the shared component names. The environment must provide ``robot_spec.dashboard``; it defines the task -command and fields, runtime components, and frame channels exposed by the -frontend. The Session controller waits for that robot-defined command +command and fields, runtime components, and allowed primitive controls. +Camera tabs use the ``frame_channels`` mapping to recorded image artifacts. The Session controller waits for that robot-defined command (``/rpent-task`` for LIBERO). For every claimed TaskRun, the Dashboard calls ``parse_config`` and the same ``robot_spec.init_runtime`` hook with the unique component names, merges the shared and unique -primitive inputs, and creates a fresh toolkit and planner conversation. Both +runtime inputs, and creates a fresh toolkit and planner conversation. Both subsets come from explicit ``shared`` / ``unique`` scope values in the environment's Dashboard spec. In LIBERO, VLA and SAM3 are reused while the Dashboard is running, while every TaskRun gets a separate environment runtime @@ -211,8 +216,8 @@ During a TaskRun, the Dashboard shows: The page accepts ordinary planner messages, new task commands, and interrupt requests. It also exposes the primitives listed by the environment's Dashboard -spec and validates their arguments against the Toolkit input schemas before -executing them directly. Planners, toolkits, and robot runtimes publish display +spec and executes them through the Toolkit, which validates their arguments +using the same parameter models as planner calls. Planners, toolkits, and robot runtimes publish display updates through a ``dashboard_events`` sink. The server sends state summaries over SSE, and the frontend fetches detailed events, timeline data, and images as needed. diff --git a/docs/source-en/rst_source/development/interfaces.rst b/docs/source-en/rst_source/development/interfaces.rst index 109abde2a..137645446 100644 --- a/docs/source-en/rst_source/development/interfaces.rst +++ b/docs/source-en/rst_source/development/interfaces.rst @@ -80,38 +80,104 @@ Most users pick a built-in ``api``, ``claude_code``, or ``codex`` planner — se dashboard_interaction=None, ) -> PlannerResult: ... -Contract: pass ``toolkit.get_tools_spec()`` to the model; dispatch each call via -``toolkit.execute_tool(name, input_dict)``; feed results back to the model; return -``PlannerResult`` on the ``finish`` tool or when turns are exhausted. - -Toolkit -------- - -Subclass ``Toolkit`` in ``robots//toolkit.py`` and register robot tools with -``add_tool``: +Contract: read ``toolkit.list_tools()`` and adapt each ``Tool``'s ``name``, +``description``, and ``input_schema`` to the model SDK. Dispatch through +``toolkit.execute_tool(name, arguments)`` and return ``PlannerResult`` when +``toolkit.finish_result`` is set or a run limit is reached. For asynchronous +adapters, use ``rpent.planner.base.execute_tool`` to run the synchronous executor +in a worker. The API and MCP adapters serialize tool calls. + +Native tools and Toolkit +------------------------ + +Import ``Tool``, ``ToolContext``, ``ToolResult``, ``Toolkit``, ``tool``, +and ``readonly`` from ``rpent.tools``. + +- ``@tool`` turns a function into a ``Tool``. Its name and Google-style + docstring describe the tool; typed parameters and Pydantic ``Field`` + constraints generate both the validation model (``args_schema``) and + the published JSON schema (``input_schema``). Unknown top-level arguments, + including a caller-supplied ``ctx``, are rejected before execution; + the schema advertises ``additionalProperties: false``. +- Every handler takes a required keyword-only ``ctx: ToolContext[RobotRuntime]``. + The executor injects it and excludes it from the model-facing schema. + It provides ``state``, ``memory``, ``robot``, ``output_dir``, + ``record_frame(rgb)``, and ``check_cancelled()``. +- Handlers return ``ToolResult(data={...}, images=[png_bytes], error=None)``. + ``to_dict()`` combines data and any error; ``to_text()`` serializes that + payload, truncating only the model-facing text to 60,000 bytes. PNG bytes + remain separate in ``images``; ``is_error`` indicates an error. + After a successful ``finish`` call, ``Toolkit.finish_result`` retains the full + data payload except the internal ``_finish`` marker. Planners use this accepted + result, including robot-specific fields such as ``operator_aborted`` and + ``operator_notes``. + +Construct the robot subclass with a fixed tuple of native tools: .. code-block:: python - def add_tool(self, name: str, spec: dict, handler) -> None: ... - -.. list-table:: - :header-rows: 1 - :widths: 22 78 - - * - Argument - - Meaning - * - ``name`` - - Tool name the LLM sees. - * - ``spec`` - - Tool description and parameter schema (``name``, ``description``, - ``input_schema``). - * - ``handler`` - - Implementation; **must return a ``dict``**. Set ``_finish`` when the task - ends; optional ``_image_bytes`` (etc.) to return camera images. - -The base class already registers common file tools; call ``super().__init__()`` then -``add_tool`` for robot tools. Per-step state and ``view_env_state`` are in -:doc:`add_primitive`. + super().__init__( + state=state, + memory=memory, + robot=runtime, + output_dir=output_dir, + tools=MYROBOT_TOOLS, + dashboard_events=dashboard_events, + ) + +The base class adds ``read_text_file``, ``write_text_file``, ``list_dir``, and +``read_image`` from ``rpent.tools.common_tools``. The MCP adapters omit +``read_image`` because Claude Code and Codex use their built-in image readers. +Memory file access goes through ``MemoryManager.authorize_read`` and +``authorize_write``. + +Franka compatibility +~~~~~~~~~~~~~~~~~~~~ + +``FrankaToolkit`` and ``DualFrankaToolkit`` preserve main robot tool parameters, +normal result fields, image order and path fields, and ``finish``. Focused tests +compare full input schemas, tool descriptions and normal return fields with a +historical baseline, retaining the published defaults and parameter descriptions. +Schema differences are explicitly listed optional inputs that now also accept +``null`` and the shared rejection of unknown top-level arguments. File tools, +error handling, and validated argument logs follow +the shared native executor. Arm normalization is declared in the Pydantic +parameter type. + +Execution and lifecycle +~~~~~~~~~~~~~~~~~~~~~~~ + +Each toolkit permits one active call. Overlapping direct calls return an error; +API and MCP adapters serialize their calls. Place ``@readonly`` below ``@tool`` +to skip automatic observation capture. + +Non-readonly robot tools capture a new observation after execution. Common tools +and ``finish`` are excluded from capture: ``write_text_file`` and ``finish`` run +exclusively without adding an observation. LIBERO ``segment`` uses ``@readonly``: +it saves segmentation artifacts on the source step and returns the segmentation +result directly, without capturing a new observation. + +Override ``_capture_observation(*, command, result, elapsed_s)`` to save a +``StepRecord`` and return ``(observation_data, png_images)``. The executor +replaces action data with observation data, appends the images, and retains +any action error. Include the action log in the observation when needed. +Capture also runs after handler errors; the call remains active until capture +and Dashboard publication finish. + +Long-running handlers call ``ctx.check_cancelled()`` at safe boundaries. +``cancel_active_and_wait()`` signals the active call and waits for it to exit. +Subsequent calls receive a fresh cancellation signal. Tools submit RGB frames +with ``ctx.record_frame``; robot toolkits save per-action clips during capture +and override ``close()`` to save their episode video. + +Each robot supplies its own ``finish`` tool. A successful call stores its +``status`` and ``summary`` in ``toolkit.finish_result``; it does not close +admission. ``solved()`` reports environment success independently of the +planner's requested finish status. Robot toolkits implement +``write_recipe(recipe_tag)`` using their recorded state trace. LIBERO exports +the successful attempt after the last reset; RoboCasa and RoboTwin retain +their action filters. The runner decides whether the run qualifies for memory +publication. Inter-process communication --------------------------- diff --git a/docs/source-en/rst_source/development/memory.rst b/docs/source-en/rst_source/development/memory.rst index 2edfdb0b2..09fff8df3 100644 --- a/docs/source-en/rst_source/development/memory.rst +++ b/docs/source-en/rst_source/development/memory.rst @@ -73,3 +73,19 @@ repository ships no self-serve upload path. To contribute a new or updated memory note, open an RPent issue with the proposed memory file and its provenance, and a maintainer will review and publish accepted files to ``RLinf/RPent-memory``. + +Tool access and recipe export +----------------------------- + +Shared file tools live in ``rpent.tools.common_tools`` and use the current +``ToolContext.memory``. ``MemoryManager.authorize_read(path)`` and +``authorize_write(path)`` resolve paths relative to the repository root and +apply the current robot's memory permissions. Published memory is read-only; +exploration may write to its configured ``_internal/inbox//``. Access to +another robot's repository memory is denied. These checks govern memory +access; they do not restrict all files to the output directory. + +Each robot toolkit implements ``write_recipe(recipe_tag)`` from its state +trace. LIBERO exports the successful attempt after the last reset; RoboCasa +and RoboTwin filter recorded actions using their existing recipe rules. The +runner decides whether the audit and recipe qualify for publication to memory. diff --git a/docs/source-en/rst_source/usage/configure_planner.rst b/docs/source-en/rst_source/usage/configure_planner.rst index 5a495a418..ead12df35 100644 --- a/docs/source-en/rst_source/usage/configure_planner.rst +++ b/docs/source-en/rst_source/usage/configure_planner.rst @@ -281,14 +281,20 @@ construction branch to ``rpent.planner.base.build_planner``: toolkit, max_turns, input_queue=None, + dashboard_interaction=None, ): - tool_specs = toolkit.get_tools_spec() + tools = toolkit.list_tools() + tool_specs = [ + {"name": tool.name, "description": tool.description, + "input_schema": tool.input_schema} + for tool in tools + ] # Call the model with system_prompt, user_message, and tool_specs. # Execute each tool call through this interface: tool_result = toolkit.execute_tool(tool_name, arguments) ... return PlannerResult( - finish_result=finish_result, + finish_result=toolkit.finish_result, messages=messages, stats=stats, error=error, @@ -297,12 +303,14 @@ construction branch to ``rpent.planner.base.build_planner``: Any planner must: 1. Accept the rendered ``system_prompt`` and ``user_message``. -2. Read the tool schemas from ``toolkit.get_tools_spec()`` and execute - tools with ``toolkit.execute_tool(name, arguments)``. -3. Convert the text and images in ``ToolResult.content_blocks`` to the - format expected by the model SDK. -4. Detect ``ToolResult.is_finish`` and stop according to - ``max_turns`` and any other limits. +2. Read native tools from ``toolkit.list_tools()`` and adapt their ``name``, + ``description``, and ``input_schema`` to the SDK. Execute calls through + ``toolkit.execute_tool(name, arguments)``; asynchronous adapters use + ``rpent.planner.base.execute_tool`` to execute tools in a worker thread. +3. Convert ``ToolResult.to_text()`` and the PNG bytes in ``ToolResult.images`` + to the SDK format, preserving ``ToolResult.is_error``. +4. Check ``toolkit.finish_result`` and stop according to ``max_turns`` and + other limits. A finish result does not itself close the toolkit. 5. Return a ``PlannerResult`` containing the finish state, messages, statistics, and an optional error. diff --git a/docs/source-zh/rst_source/development/add_primitive.rst b/docs/source-zh/rst_source/development/add_primitive.rst index 4d39960a3..3047ef1f8 100644 --- a/docs/source-zh/rst_source/development/add_primitive.rst +++ b/docs/source-zh/rst_source/development/add_primitive.rst @@ -27,54 +27,67 @@ - ``move_to``、``rotate_wrist``、``release``、 ``back_project`` -从 LLM 的视角看，两类原语采用相同的接口：一份工具定义、一个 -primitives 方法，以及调用完成后的状态快照。区别仅在于方法的具体实现。 +从 LLM 的视角看，两类原语采用相同的接口：一份工具定义、一个执行函数， +以及调用完成后的状态快照。区别仅在于函数内部是调用模型，还是执行脚本化动作。 添加一个脚本化原语 ------------------ 添加脚本化原语通常需要以下两个步骤： -1. **在 primitives 中添加方法。** 在当前机器人的 primitives - 类（如 ``LiberoPrimitives``、``MyRobotPrimitives``）中添加 - 一个方法。该方法接收工具调用的参数，执行一次或多次 - ``self._env.step(...)``，并返回一个简短的日志字典。 - - primitive 方法执行后默认会自动捕获并重新渲染状态 - （``get_env_state``）： +1. **编写工具函数。** 在 ``robots//tools.py`` 中添加函数，用 ``@tool`` + 将它声明为工具。函数通过 ``ctx.robot`` 访问当前机器人的环境和模型客户端， + 执行动作后返回 ``ToolResult``。例如，下面的 LIBERO 工具会保持当前位姿， + 并在指定步数后停止： .. code-block:: python - 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} - - 只读工具（``view_env_state``、``back_project``、``segment`` 等） - 可以使用 :func:`~rpent.tools.toolkit.readonly` 标记，toolkit 会跳过 - 它们的状态捕获，提升性能。 - -2. **添加工具定义。** 在 ``robots//tools.py`` 的 ``TOOLS_SPEC`` 中新增一项： - - .. code-block:: python - - { - "name": "open_drawer", - "description": "Pull the currently-grasped drawer handle " - "backwards by ``dx`` meters.", - "input_schema": { - "type": "object", - "properties": {"dx": {"type": "number"}}, - "required": [], - }, - } - -两者就位后，toolkit 会自动注册该工具：它遍历 ``TOOLS_SPEC``，把每个定义 -绑定到对应的 primitive 方法（如 ``getattr(self._primitives, name)``）。 + from typing import Annotated + + from pydantic import Field + + from rpent.tools import ToolContext, ToolResult, tool + + @tool + def hold_pose( + steps: Annotated[int, Field(ge=1, le=100)] = 10, + *, + ctx: ToolContext, + ) -> ToolResult: + """Hold the current pose with the gripper closed. + + Args: + steps: Number of environment steps. + """ + runtime = ctx.robot + for _ in range(steps): + ctx.check_cancelled() + obs, _, terminated, truncated, _ = runtime.env.step( + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0] + ) + runtime.executed_steps += 1 + runtime.set_obs(obs) + ctx.record_frame(obs["main_images"]) + if terminated or truncated: + break + return ToolResult(data={"steps_requested": steps}) + + ``@tool`` 根据函数签名生成参数 schema，并从 Google 风格 docstring 中读取 + 工具说明。示例中的 ``Field`` 限定了模型可传入的步数；``ctx`` 则由 toolkit + 提供，不需要模型填写。在机器人代码中，可以用 ``ToolContext[LiberoRuntime]`` + 进一步标明上下文中的机器人类型。 + + 工具执行后，toolkit 会自动保存新的状态快照。对于 ``view_env_state``、 + ``back_project`` 等读取已有观测的工具，可以在 ``@tool`` 下方添加 + ``@readonly``，省去这次状态捕获；调用仍按顺序执行。公共工具和 ``finish`` + 不触发状态捕获；``write_text_file`` 和 ``finish`` 不设置 readonly，独占执行。 + +2. **将工具加入 toolkit。** 把函数声明加入该机器人的工具集合，例如 LIBERO 的 + ``LIBERO_TOOLS``。Toolkit 在构造时接收这组工具，并统一处理参数校验和调用。 完成以上步骤后，``api``、``claude_code`` 和 ``codex`` 三种 planner -都可以调用该工具，无需修改其他代码。 +都可以调用该工具，无需分别编写适配代码。执行和取消的约定参见 +:doc:`interfaces` 中的工具集说明。 .. _add-primitive-model-based: @@ -102,39 +115,21 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法 LIBERO 的实现可参考 ``rpent.robots.components.pi05_vla_client.Pi05VLAClient``。 -3. **在 primitives 中添加方法。** 在当前机器人的 primitives - 类中调用 model client，将其返回的动作块交给环境执行，并返回日志字典。 - model client 的接口是 - :meth:`rpent.robots.components.pi05_vla_client.Pi05VLAClient.predict`， - 指令从 ``env_obs["task_descriptions"]`` 中读取；返回 ``[chunk, action_dim]`` - 的 numpy 动作块（已剥掉 batch 维）： - - .. code-block:: python - - def mymodel_pick(self, target: str) -> dict: - env_obs = self._env.get_obs() - env_obs["task_descriptions"] = f"pick {target}" - chunk = self._model.predict(env_obs) - self._env.chunk_step(chunk) - return {"model": "mymodel", "target": target} - -4. **添加工具定义并在 toolkit 中注册。** 具体做法与脚本化原语相同。 - -5. **在 ``robot_spec.py`` 中连接各组件。** 机器人的 ``get_toolkit`` 使用 - ``runtime_kwargs`` 构造 toolkit： - - .. code-block:: python +3. **编写工具函数。** 在函数中通过 ``ctx.robot`` 调用 model client，将返回的 + 动作块交给环境执行，并用 ``ToolResult`` 返回执行结果。以 Pi0.5 为例， + 指令从 ``env_obs["task_descriptions"]`` 中读取，模型返回 + ``[chunk, action_dim]`` 的 NumPy 动作块（已去掉 batch 维）。具体实现可参考 + ``robots/libero/tools.py`` 中的 ``pi0_pick``，以及 + ``robots/robocasa/tools.py`` 中的 ``rldx_skill``。 - def get_toolkit(*, runtime_kwargs, dashboard_events): - from robots.myrobot.toolkit import MyRobotToolkit - return MyRobotToolkit( - runtime_kwargs=runtime_kwargs, - dashboard_events=dashboard_events, - ) +4. **将工具加入 toolkit。** 和脚本化原语一样，使用 ``@tool`` 声明工具， + 再将它加入机器人的工具集合。执行后的状态捕获仍由 toolkit 负责。 - 机器人包中的 ``_init_runtime`` 则负责构造 ``runtime_kwargs``，例如 - ``{"env": MyRobotEnvClient(...), "model": MyModelClient(...)}``，再由 - toolkit 构造器将其转发给 primitives。 +5. **在 ``robot_spec.py`` 中连接各组件。** 机器人包中的 ``_init_runtime`` + 负责创建环境和模型客户端，并通过 ``runtime_kwargs`` 返回，例如 + ``{"env": MyRobotEnvClient(...), "model": MyModelClient(...)}``。 + ``get_toolkit`` 将这些客户端连同输出目录和 ``MemoryManager`` 传给 toolkit， + 由后者构造本次会话的运行时对象。完整的工厂示例见 :doc:`add_robot`。 在多次运行之间复用 vla_server ----------------------------- @@ -169,8 +164,8 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法 ``session_id`` 由 facade 从连接派生并注入 server 端 handler，客户端 **不传**，也不应在 ``predict`` 的 ``options`` 里伪造 ``session_ids``。 -- **primitives 侧**：任务开始前调用 ``reset_session`` 清空上一回合残留 - 的策略状态，保证连续多次运行之间状态不串。 +- **工具侧**：在任务开始或环境重置时调用 ``reset_session``，清空上一回合残留 + 的策略状态，避免影响后续任务。 单线程 serve（EGL 渲染后端） ---------------------------- @@ -209,12 +204,12 @@ mixin 覆盖的 ``serve`` 与 :class:`~rpent.utils.rpc.RpcFacade` 的 - **工具名称应描述意图，而非底层动作序列。** 例如使用 ``pi0_pick``， 而不是 ``execute_action_chunk_of_length_20``。 -- **每个工具执行结束后都要保存新的状态快照。** 下一轮需要读取动作执行后的 - 环境状态，因此原语不能在渲染完成前返回。 -- **工具只返回简短的字典。** 返回值会以文本形式提供给 LLM；图像、深度数据和 - 其他大型观测应通过 ``EnvState.save`` 保存；``EnvState`` 会把每个逻辑基础 - 文件名自动加入其持有的 ``StepRecord.artifacts`` 集合。图像通过 - ``view_env_state`` 提供，几何数据通过环境工具访问，不返回原始路径。 +- **动作执行后要有新的状态快照。** Toolkit 会在工具函数执行完毕后捕获观测， + 再将结果交给 planner，让下一轮推理能看到动作后的环境。只读工具可以复用已有观测。 +- **返回简短的执行结果。** 将动作摘要放在 ``ToolResult.data`` 中，供 planner + 以文本形式读取。图像、深度等大型观测通过 ``EnvState.save`` 保存，文件名会 + 自动记入 ``StepRecord.artifacts``。需要向模型展示图片时，使用 + ``ToolResult.images``；历史观测仍可通过 ``view_env_state`` 读取。 - **安全限制由 ``env_server`` 强制执行。** LLM 可能使用任意参数调用工具， 因此工作空间边界和安全限制不能只依赖 toolkit。 @@ -231,5 +226,5 @@ mixin 覆盖的 ``serve`` 与 :class:`~rpent.utils.rpc.RpcFacade` 的 多个模型，由工具通过 ``predict`` 的 ``model`` kwarg 选择要调用的模型 或输出 head。 -无论具体实现如何，框架的契约都保持不变：模型进程 → model client → -primitives 方法 → 工具定义 → ``Toolkit.add_tool``。 +无论使用哪种模型，都可以沿用上述接入方式：由客户端连接模型服务，再把调用模型 +和执行动作的过程写成工具函数，交给 toolkit 调用。 diff --git a/docs/source-zh/rst_source/development/add_robot.rst b/docs/source-zh/rst_source/development/add_robot.rst index fccc52bae..1344d710b 100644 --- a/docs/source-zh/rst_source/development/add_robot.rst +++ b/docs/source-zh/rst_source/development/add_robot.rst @@ -28,7 +28,7 @@ RPent 的整体进程划分、服务职责和通信方式见 :doc:`系统说明 服务和 model client，参见 :ref:`添加一个 VLA（或其他基于模型的原语）`。 3. :ref:`定义 prompt `。 -4. :ref:`实现 toolkit 和 primitives `。 +4. :ref:`实现 toolkit 和工具函数 `。 5. :ref:`注册环境参数并生成 RunConfig `。 6. 实现 :ref:`runtime 钩子 `：同一个钩子既能为普通 CLI 初始化完整 runtime，也能为 Dashboard 初始化指定的 component 子集。 @@ -48,7 +48,8 @@ RPent 的整体进程划分、服务职责和通信方式见 :doc:`系统说明 robot_spec.py # RobotSpec、工厂、Dashboard 描述和 runtime 钩子 env_client.py # MyEnvClient —— agent 侧 RPC client (§1) prompt_bundle.py # system()/user() prompt 工厂 (§2) - toolkit.py # MyRobotToolkit + primitives + 工具定义 (§3) + toolkit.py # Runtime、MyRobotToolkit 和观测处理 (§3) + tools.py # 原生工具声明与 handler (§3) env_server.py # 环境侧 facade + RPC 服务 (§1) vla_server.py # （可选）VLA 模型服务 @@ -90,6 +91,7 @@ RPent 的整体进程划分、服务职责和通信方式见 :doc:`系统说明 from robots.myrobot.toolkit import MyRobotToolkit return MyRobotToolkit( runtime_kwargs=runtime_kwargs, + output_dir=config.output_dir, dashboard_events=dashboard_events, memory=MemoryManager( root=config.prompt_vars.get("memory_dir") or get_memory_dir("myrobot"), @@ -259,63 +261,63 @@ API 版本。 3. ``toolkit.py`` ------------------ -这个模块持有 LLM 能调用的一切: 工具 schema、primitives、每步状态 dump 以及 -MCP allowlist。(LIBERO 中由于历史原因把这些拆到了 ``tools.py`` 和 ``toolkit.py`` -两个文件; 新增 robot 时全部放在 ``toolkit.py`` 里没问题。) - -toolkit 模块通常包含四部分： - -**Primitives 类**\ （例如 ``MyRobotPrimitives``）是 toolkit 持有的 Python -对象。它保存 ``EnvClient``、VLA ``model`` client 和单次运行所需的状态。每个 -原语工具（``move_to``、``pi0_pick``、``release`` 等）对应一个方法，并返回 -日志字典。 - -**工具定义和处理函数** 包括模块级的 ``TOOLS_SPEC`` 列表（列表元素采用 -Anthropic API 的工具定义格式，包含 ``name``、``description`` 和 -``input_schema``），以及 toolkit 引用的模块级函数，例如 -``view_env_state``、``back_project`` 和 ``finish``。 - -**每步状态 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``： - -- 在 ``super().__init__(...)`` 中传入 ``memory`` （一个 - :class:`~rpent.memory.MemoryManager`）和 ``state``。``memory_access`` 和 - ``inbox_cell_tag`` 在构造 ``MemoryManager`` 时配置；eval 默认只读。 -- 在 ``__init__`` 中通过自定义的初始化辅助方法构建 primitives（LIBERO - 中的方法名为 ``init_primitives``；它会调用 ``EnvState.reset()``、构造 - 原语并 dump 第 0 步）, -- 用 ``self.add_tool(name, spec, handler)`` 注册每个工具。无状态的读取工具 - （如 ``view_env_state``、``finish``）直接绑定模块级函数；原语工具通过 - ``_step(name, **kwargs)`` 调用。``_step`` 使用 - ``getattr(self._primitives, name)(**kwargs)`` 调用 driver 方法并重新渲染状态； -- 重写 ``close()``，通过 ``EnvState`` 保存 agent 侧剩余工件（例如 - ``state.save("episode.mp4", frames, step=None)``）。 - -``runtime_kwargs`` 由 ``robot_spec.py:get_toolkit`` 转发给 toolkit，再原样传入 -primitives 的 ``__init__``。其中通常包含 -``{"env": MyEnvClient(...), "model": VLAClient(...), ...}``。 +这个模块负责组织 LLM 可以调用的工具，以及这些工具需要的客户端和环境状态。 +现有的 LIBERO、RoboCasa 和 RoboTwin 实现将工具函数放在 ``tools.py``， +其余部分放在 ``toolkit.py``，新增机器人时可以沿用这一划分。 + +通常需要实现以下四部分： + +**运行时对象**\ （例如 ``MyRobotRuntime``）由 toolkit 持有，保存 ``EnvClient``、 +VLA 客户端和本次会话的状态。工具函数通过 ``ctx.robot`` 访问它，从而共用同一个 +环境、缓存观测和任务进度。 + +**工具定义和处理函数** 放在 ``tools.py`` 中，用 ``@tool`` 声明。函数签名描述 +工具参数，Google 风格 docstring 提供工具说明，函数本身负责执行动作并返回 +``ToolResult``。将这些声明收集到 ``MYROBOT_TOOLS`` 元组中，就得到了该机器人 +提供的工具集合。具体写法见 :doc:`add_primitive`。 + +**每步状态保存** 由 ``dump_state(runtime, env_state, log)`` 完成。它通过 +``env_state.record_step(...)`` 创建步骤记录，并用 ``env_state.save(...)`` +保存观测。在 ``record_step`` 块内省略 ``step`` 时，文件属于当前步骤；指定 +``step=`` 可以保存到其他步骤，``step=None`` 则用于整次运行的文件。 +每次保存成功后，文件名会自动加入该 ``StepRecord`` 的 ``artifacts`` 集合。 +再由 ``build_observation(state, record)`` 将记录整理成文本数据和图片，供动作 +响应与 ``view_env_state`` 共用。 + +**Toolkit 类** 继承 ``Toolkit[MyRobotRuntime]``，将前面几部分连接起来： + +- 在 ``__init__`` 中构造运行时对象和 ``EnvState``，并将它们连同 ``memory``、 + ``output_dir``、``tools=MYROBOT_TOOLS`` 和 ``dashboard_events`` 传给基类。 + 运行时对象使用 ``robot`` 参数传入；memory 的访问权限在 ``MemoryManager`` + 上配置，评测时默认为只读。 +- 初始化环境并保存第 0 步状态。如果支持 Dashboard，同时发布 + ``StepRecordEvent``，让页面显示初始观测。 +- 实现 ``_capture_observation(*, command, result, elapsed_s)``，调用上述 + 状态保存与观测整理函数，返回观测数据和 PNG 图片。Toolkit 会在动作执行后 + 自动调用它；原始执行结果可通过 ``result.to_dict()`` 保存到步骤日志中。 +- 实现 ``solved()``，根据环境状态判断任务是否成功。工具调用与取消由基类处理。 + 机器人 toolkit 在捕获观测时保存动作视频，重写 ``close()`` 保存回合录像， + 并通过 ``write_recipe()`` 从记录的状态导出 recipe。 + +``runtime_kwargs`` 由 ``robot_spec.py:get_toolkit`` 转发给 toolkit，用于构造 +运行时对象。其中通常包含 ``{"env": MyEnvClient(...), "model": VLAClient(...)}`` +及其他辅助客户端。 建议遵循的约定 -------------- - ``output_dir`` 是 runner 为单次运行创建的工作目录。环境观测由 - ``EnvState`` 管理；调用方只使用逻辑基础文件名，不自行拼接存储路径。 - transcript 等运行管理输出与环境工件共享该目录。 -- 工具定义使用 Anthropic API 格式（``name`` / ``description`` / - ``input_schema``）。 - 每个用 ``self.add_tool(...)`` 注册的工具都会暴露给所有 planner。 -- 环境侧的返回值必须可 pickle，且不包含 torch 对象。 -- 每个原语工具执行后要 dump 一次新的状态快照, 这样下一次 - ``view_env_state`` 看到的是动作后的世界。 -- ``dump_state`` 是 Agent 获取环境状态的唯一数据来源；任何新的模态 - （例如触觉、力）都通过它提供。 + ``EnvState`` 管理；调用方使用逻辑文件名，不自行拼接存储路径。 + transcript 等运行记录与环境文件共享该目录。 +- 工具定义来自带 ``@tool`` 的函数，planner 负责将它们转换为各自 SDK 所需的格式。 + 机器人工具只需实现一次，就能供不同 planner 调用。 +- 环境侧返回传输层支持的 Python / NumPy 数据，不包含 torch 对象。 +- 动作工具执行后，由 toolkit 保存新的状态快照，让下一次 ``view_env_state`` + 能读到动作后的环境。工具需要录制过程画面时，通过 ``ctx.record_frame`` 提交帧。 +- 新增触觉、力等观测模态时，仍通过 ``dump_state`` 保存，再由 + ``build_observation`` 提供给 planner。 + +工具的参数、返回值和执行约定见 :doc:`interfaces`。 .. _add-robot-config: @@ -373,9 +375,9 @@ main.py 已创建的共享 parser。``use_dashboard`` 决定原本必填的参 - ``owned_daemons: list[ProcessDaemon]`` 只包含当前进程实际启动的子进程， 当前 runner 会在清理阶段停止它们。连接外部 endpoint 时，不能把外部服务加入 该列表。 -- ``runtime_kwargs: dict`` 会传给 toolkit 构造器，再由后者传入 primitives - 的 ``__init__``。完整参数通常包含 - ``{"env": MyEnvClient(...), "model": VLAClient(...)}``，以及其他辅助 client。 +- ``runtime_kwargs: dict`` 包含本次运行所需的客户端，通常是 + ``{"env": MyEnvClient(...), "model": VLAClient(...)}`` 及其他辅助客户端。 + Toolkit 使用这些参数构造运行时对象，供工具函数共用。 第四个参数 ``components`` 指定要初始化的服务名称。``None`` 表示全部服务，普通 CLI 会传入这个值。Dashboard 根据 ``dashboard.runtime_components`` 得到两个子集， diff --git a/docs/source-zh/rst_source/development/architecture.rst b/docs/source-zh/rst_source/development/architecture.rst index 3cfe7784c..484107e4a 100644 --- a/docs/source-zh/rst_source/development/architecture.rst +++ b/docs/source-zh/rst_source/development/architecture.rst @@ -46,7 +46,7 @@ LLM-in-the-loop 运行流程 一次运行就是一段 LLM-in-the-loop 循环： 1. LLM 分析任务、调一个工具 (如 ``pi0_pick``)。 -2. 工具的底层 primitives 向 ``vla_server`` 请求动作 (``predict``)。 +2. 工具函数通过模型客户端向 ``vla_server`` 请求动作（``predict``）。 3. ``env_server`` 执行动作。 4. 环境返回更新后的观测数据和相机画面。 5. 执行结果会整理成由文本和图像组成的上下文，返回给 LLM 进行下一轮推理。 @@ -67,12 +67,15 @@ LLM-in-the-loop 运行流程 context/ # 提示词工具和共享提示词片段。 dashboard/ # FastAPI 监控页面和 SSE 事件流（可选）。 robots/ # RobotSpec、PromptBundle 和按需加载机器人的逻辑。 - tools/ # Toolkit 基类和共享 tool 辅助函数。 + tools/ # 原生工具协议、执行器和公共工具。 + session/ # EnvState、步骤记录和工件存储。 + memory/ # Memory 同步、访问控制与探索结果合并。 utils/ # 配置、日志、RPC 客户端/服务端和 VLA 客户端。 robots/ libero/ # LIBERO 的 env_client / env_server / vla_server / # toolkit / prompt_bundle。参考实现。 robocasa/ # RoboCasa 机器人 (RLDX-1 VLA，厨房任务)。 + robotwin/ # RoboTwin 机器人 (LingBot-VLA，双臂任务)。 (franka/) # Franka 机器人——研发中。 (so101/) # SO-101 机器人——研发中。 scripts/ @@ -131,28 +134,31 @@ planner 后端集中在 ``rpent/planner/``， # robots/myrobot/__init__.py def get_robot_spec() -> RobotSpec: ... # 机器人标识、提示词模板与 Runner 钩子 def get_toolkit( - *, runtime_kwargs, dashboard_events + *, runtime_kwargs, dashboard_events, config ): ... ``RobotSpec`` 汇集了机器人标识、prompt 模板、可选的 Dashboard 描述与三个 Runner 钩子（``add_cli_args`` / ``parse_config`` / ``init_runtime``）。各字段要填什么见 :doc:`interfaces`。 -加载器本身不维护机器人名称列表。当前 CLI 将 ``--robot`` 限定为 ``libero`` -和 ``robocasa``；接入新的机器人名称时，还需要同步更新 CLI 的可选值。完整步骤见 -:doc:`add_robot`。 +加载器从磁盘发现机器人包，CLI 通过 ``enumerate_robots()`` 得到 ``--robot`` +可选值。目前提供 LIBERO、RoboCasa 和 RoboTwin；完整接入步骤见 :doc:`add_robot`。 Planner、Toolkit 与 RPC 传输层 ------------------------------ -这三层各管一段、层层解耦。planner 只通过 ``get_tools_spec`` 拿到工具清单、 -用 ``execute_tool`` 逐个调用，并不关心工具背后是脚本还是 VLA； -toolkit 把每次工具调用翻译成对 primitive 的调用，再由 primitives -经 RPC 向 ``env_server`` / ``vla_server`` 发起 ``reset`` / ``step`` / -``predict`` 请求；RPC 传输层（HTTP 或 socket）只负责把这些调用和 NumPy -观测在进程间搬运，对上层透明。正因如此，换 planner 不影响工具， -换传输协议也不影响 planner。三者的具体接口契约（``Planner.solve``、 -``Toolkit.add_tool``、``RpcFacade._dispatch``）集中在 :doc:`interfaces`。 +这三层分别负责工具调用中的不同环节。Planner 通过 ``list_tools()`` 获取工具 +定义，再用 ``execute_tool`` 执行调用，不需要关心工具背后是脚本还是 VLA。 +它只需将工具定义和返回的文本、图片转换成模型 SDK 所需的格式；Claude Code +和 Codex 使用的 MCP 适配也在这一层完成。 + +Toolkit 负责校验参数、安排工具执行，并在动作完成后保存新的观测。工具函数 +通过 ``ctx.robot`` 获取环境和模型客户端，向 ``env_server`` / ``vla_server`` +发起 ``reset``、``step`` 或 ``predict`` 请求。RPC 传输层（HTTP 或 socket） +再将这些请求和 NumPy 观测传递到对应进程。 + +这样的分工使更换 planner 不必修改工具，更换传输协议也不影响 planner。 +三者的具体接口约定见 :doc:`interfaces`。 Dashboard（可选） ----------------- @@ -162,11 +168,12 @@ Dashboard（可选） ``--dashboard-host`` 和 ``--dashboard-port`` 启动 Dashboard。Session 配置全部来自 命令行，然后用共享 component 名称调用一次 ``robot_spec.init_runtime``。环境必须 提供 ``robot_spec.dashboard``，由它定义 -前端使用的任务命令与字段、runtime components 和 frame channels。Session +前端使用的任务命令与字段、runtime components 和允许执行的原语。相机标签通过 +``frame_channels`` 映射到每步记录的图片工件。Session controller 随后等待该环境定义的命令（LIBERO 使用 ``/rpent-task``）；每次取得一个 TaskRun 后，Dashboard 会调用 ``parse_config``，再用 unique component 名称调用 -同一个 ``robot_spec.init_runtime``，合并 shared 与 unique primitive 参数，并新建 toolkit -和 planner conversation。两个子集都来自环境 Dashboard spec 中显式声明的 +同一个 ``robot_spec.init_runtime``，合并两次返回的客户端参数，并为本次任务新建 toolkit +和 planner 会话。两个子集都来自环境 Dashboard spec 中显式声明的 ``shared`` / ``unique`` scope。在 LIBERO 中，VLA 和 SAM3 会在 Dashboard 运行期间 复用，每个 TaskRun 使用独立环境并按顺序执行。 @@ -179,8 +186,8 @@ TaskRun 运行期间，Dashboard 页面提供： - 运行结束后的完整回合录像（如果已生成）。 页面支持提交 planner 消息、新任务命令和中断请求，也会展示环境在仪表盘配置 -（``DashboardSpec``）中列出的原语。执行原语前，参数会根据工具包（``Toolkit``） -定义的输入结构进行校验。 +（``DashboardSpec``）中列出的原语。原语通过工具包（``Toolkit``）执行，并复用 +planner 工具调用所用的参数模型进行校验。 planner、toolkit 和机器人运行时通过 ``dashboard_events`` 事件接收器发布展示更新。 服务端通过 SSE 推送运行状态摘要，前端再按需读取详细事件、时间线和图像。 diff --git a/docs/source-zh/rst_source/development/interfaces.rst b/docs/source-zh/rst_source/development/interfaces.rst index 76bea2122..dfba02f0e 100644 --- a/docs/source-zh/rst_source/development/interfaces.rst +++ b/docs/source-zh/rst_source/development/interfaces.rst @@ -46,8 +46,8 @@ ``runtime_kwargs``。普通 CLI 传 ``None``；Dashboard 从 spec 得到显式声明 的 shared 和 unique 子集后分别传入。``DashboardEventSink`` 用于上报运行时状态。 -``get_toolkit`` 一般只需把 ``runtime_kwargs`` 传给机器人子类； -``dashboard_events`` 和 ``config`` 由当前 runner 传入。它需要构造一个 +``get_toolkit`` 用 ``runtime_kwargs`` 中的客户端构造机器人 toolkit； +运行配置 ``config`` 和 Dashboard 事件接收器 ``dashboard_events`` 也由 runner 提供。它需要构造一个 :class:`~rpent.memory.MemoryManager`（root 取自 ``config.prompt_vars["memory_dir"]``，未设置时回退到 ``get_memory_dir(robot_name)``）并传给 toolkit。Memory 访问权限在 @@ -77,34 +77,91 @@ Planner dashboard_interaction=None, ) -> PlannerResult: ... -约定：用 ``toolkit.get_tools_spec()`` 把工具交给模型；每次调用 ``toolkit.execute_tool(name, input_dict)``； -把结果喂回模型；在 ``finish`` 工具或轮次用尽时返回 ``PlannerResult``。 +Planner 先通过 ``toolkit.list_tools()`` 获取工具定义，将每个工具的 ``name``、 +``description`` 和 ``input_schema`` 交给模型。模型发起调用后，使用 +``toolkit.execute_tool(name, arguments)`` 执行，再将返回的文本和图片送回模型， +继续下一轮推理。当 ``toolkit.finish_result`` 有值，或达到运行限制时，返回 +``PlannerResult``。 + +如果 planner 使用异步调用，可以通过 ``rpent.planner.base.execute_tool`` +在线程中执行工具。这个辅助函数也会处理取消，等待工具完成清理后再退出。 工具集 ------ -在 ``robots//toolkit.py`` 里继承 ``Toolkit``，用 ``add_tool`` 注册机器人工具： +在 ``robots//toolkit.py`` 中继承 ``Toolkit``，构造时传入机器人使用的 +客户端、状态和工具集合： .. code-block:: python - def add_tool(self, name: str, spec: dict, handler) -> None: ... + super().__init__( + state=state, + memory=memory, + robot=runtime, + output_dir=output_dir, + tools=MYROBOT_TOOLS, + dashboard_events=dashboard_events, + ) + +``MYROBOT_TOOLS`` 是由 ``@tool`` 声明组成的元组。编写工具时，主要需要了解 +以下三部分，它们都可以从 ``rpent.tools`` 导入： .. list-table:: :header-rows: 1 :widths: 22 78 - * - 参数 - - 含义 - * - ``name`` - - LLM 看到的工具名。 - * - ``spec`` - - 工具说明与参数 schema（``name``、``description``、``input_schema``）。 - * - ``handler`` - - 执行逻辑，须返回 ``dict``。任务结束时在该 ``dict`` 里设 ``_finish``； - 需要回传相机图时可设 ``_image_bytes`` 等字段。 - -基类已注册公共文件工具；子类 ``super().__init__()`` 后追加本机器人工具即可。逐步状态与 -``view_env_state`` 见 :doc:`add_primitive`。 + * - 接口 + - 用法 + * - ``@tool`` + - 将函数声明为工具。函数名就是工具名，Google 风格 docstring 提供说明， + 参数类型和 ``Field`` 约束用于生成校验模型及 JSON schema。 + * - ``ToolContext`` + - 工具通过必填的仅限关键字参数 ``ctx`` 获取上下文。其中 ``robot`` 是 + 运行时对象，``state`` 和 ``memory`` 分别管理观测与记忆，``output_dir`` + 指向输出目录。``ctx`` 由 toolkit 提供，不出现在模型可见的参数中。 + * - ``ToolResult`` + - 工具函数的返回值。``data`` 保存执行结果，``images`` 保存 PNG 字节， + 出错时填写 ``error``。Planner 使用 ``to_text()`` 和 ``images`` 读取 + 文本与图片，通过 ``is_error`` 判断调用是否失败。 + +基类会自动加入公共文件工具，文件访问权限由 ``MemoryManager`` 检查。 +其中 ``read_image`` 供 API planner 使用；Claude Code 和 Codex 使用各自内置的 +图片读取工具。工具函数的完整示例见 :doc:`add_primitive`。 + +``FrankaToolkit`` 和 ``DualFrankaToolkit`` 保留机器人工具的主要参数、正常返回字段、 +图片顺序和路径字段，以及 ``finish``。测试完整对比历史输入 schema、工具说明和 +正常返回字段，保留公开的默认值和参数说明。schema 仅允许明确列出的可选参数 +额外接受 ``null``，并统一增加 ``additionalProperties: false``，在执行前拒绝 +未知的顶层参数（包括调用者传入的 ``ctx``）。文件工具、错误处理和校验后的参数 +日志沿用公共 native 执行器。 +arm 的归一化声明在 Pydantic 参数类型中。 + +``finish`` 成功后，``Toolkit.finish_result`` 保留完整的业务数据，仅去掉内部的 +``_finish`` 标记。Planner 读取该结果，包括 ``operator_aborted``、``operator_notes`` +等机器人特有字段，供会话控制与运行记录使用。 + +默认情况下，工具独占执行，完成后由机器人子类的 ``_capture_observation`` +保存状态并返回新的观测。观测会替换动作返回的数据，因此需要保留的执行详情 +应写入观测中的日志；动作错误仍会保留。即使工具函数出错，toolkit 也会尝试 +捕获观测，让 planner 了解当前环境。 + +读取已有观测的工具可以在 ``@tool`` 下方添加 ``@readonly``，跳过自动捕获。 +每个 toolkit 同时只允许一个调用；重叠的直接调用会返回错误。API 和 MCP +适配器按顺序执行工具调用。 + +公共工具和 ``finish`` 不触发观测捕获。``write_text_file`` 和 ``finish`` +不设置 readonly，因此独占执行但不新增观测。LIBERO 的 ``segment`` 使用 +``@readonly``，将分割附件保存到源 step 并直接返回分割结果，不新增观测。 + +长时间运行的工具应在安全的动作边界调用 ``ctx.check_cancelled()``。收到中断后， +``cancel_active_and_wait()`` 向当前调用发送取消信号，并等待它退出。后续调用 +使用新的取消信号。工具通过 ``ctx.record_frame`` 提交录像帧；机器人 toolkit +在捕获观测时保存动作片段，并重写 ``close()`` 保存回合录像。 + +每个机器人提供自己的 ``finish`` 工具。调用成功后，toolkit 将其中的 ``status`` +和 ``summary`` 保存到 ``finish_result``，供 planner 结束循环；环境是否真正成功， +则由 ``solved()`` 判断。Recipe 的导出由 ``write_recipe(recipe_tag)`` 完成， +memory 的使用与发布见 :doc:`memory`。 进程间通信 ---------- diff --git a/docs/source-zh/rst_source/development/memory.rst b/docs/source-zh/rst_source/development/memory.rst index 6e5235217..f45b428e0 100644 --- a/docs/source-zh/rst_source/development/memory.rst +++ b/docs/source-zh/rst_source/development/memory.rst @@ -56,6 +56,15 @@ memory 同步到 ``memory//``。数据集是公开的，无需 token 即 本地 memory 配置使用。Hugging Face memory 和本地 memory 使用相同的目录规范，区别只 在于来源。 +工具读写 memory 时，会使用当前会话的 ``MemoryManager`` 检查权限：评测只能读取 +当前机器人的已发布内容，探索时可以向自己的 ``_internal/inbox//`` 写入草稿。 +公共文件工具已接入这些检查；新增需要访问 memory 的工具时，可通过 +``ctx.memory.authorize_read(path)`` 或 ``authorize_write(path)`` 获取允许访问的路径。 + +各机器人 toolkit 通过 ``write_recipe(recipe_tag)`` 从状态记录导出 recipe。 +LIBERO 只导出最后一次 reset 后的成功尝试；RoboCasa 和 RoboTwin 保留各自的 +动作筛选规则。是否将生成的 audit 和 recipe 纳入 memory，由运行流程根据结果决定。 + 贡献 memory ----------- diff --git a/docs/source-zh/rst_source/usage/configure_planner.rst b/docs/source-zh/rst_source/usage/configure_planner.rst index 3ca3196f1..4417a7bf2 100644 --- a/docs/source-zh/rst_source/usage/configure_planner.rst +++ b/docs/source-zh/rst_source/usage/configure_planner.rst @@ -255,14 +255,20 @@ agent SDK，可以实现 ``rpent.planner.base.Planner`` 协议，并在 toolkit, max_turns, input_queue=None, + dashboard_interaction=None, ): - tool_specs = toolkit.get_tools_spec() + tools = toolkit.list_tools() + tool_specs = [ + {"name": tool.name, "description": tool.description, + "input_schema": tool.input_schema} + for tool in tools + ] # 使用 system_prompt、user_message 和 tool_specs 调用模型。 # 每次工具调用都通过下面的接口执行： tool_result = toolkit.execute_tool(tool_name, arguments) ... return PlannerResult( - finish_result=finish_result, + finish_result=toolkit.finish_result, messages=messages, stats=stats, error=error, @@ -271,11 +277,13 @@ agent SDK，可以实现 ``rpent.planner.base.Planner`` 协议，并在 任何 planner 必须： 1. 接收已经渲染好的 ``system_prompt`` 和 ``user_message``。 -2. 从 ``toolkit.get_tools_spec()`` 取得工具定义，并通过 - ``toolkit.execute_tool(name, arguments)`` 执行工具。 -3. 将 ``ToolResult.content_blocks`` 中的文本和图片转换成模型 SDK - 所需的格式。 -4. 识别 ``ToolResult.is_finish``，并按 ``max_turns`` 等限制终止循环。 +2. 从 ``toolkit.list_tools()`` 读取原生工具，并将其 ``name``、``description`` + 和 ``input_schema`` 转为 SDK 格式。通过 ``toolkit.execute_tool(name, arguments)`` + 执行调用；异步适配器通过 ``rpent.planner.base.execute_tool`` 在线程中执行工具。 +3. 将 ``ToolResult.to_text()`` 和 ``ToolResult.images`` 中的 PNG 字节转换为 + SDK 格式，并保留 ``ToolResult.is_error``。 +4. 检查 ``toolkit.finish_result``，按 ``max_turns`` 等限制终止循环； + 结束结果本身不会关闭 toolkit。 5. 返回包含结束状态、消息、统计信息和可选错误的 ``PlannerResult``。 由于 RPent 工具定义和 prompt 渲染流程保持不变，新增 planner 不需要修改 diff --git a/docs/source-zh/rst_source/usage/dual_franka.rst b/docs/source-zh/rst_source/usage/dual_franka.rst index 044b75233..e31a822b1 100644 --- a/docs/source-zh/rst_source/usage/dual_franka.rst +++ b/docs/source-zh/rst_source/usage/dual_franka.rst @@ -101,7 +101,7 @@ VLA 抓取 DEMO ------------- RPent 提供了一个使用 VLA 抓取物品的 DEMO。task-id ``1`` 会暴露 ``vla_right_grasp`` / ``vla_handoff`` / ``vla_left_place``， -并可在本地启动双臂 Franka VLA 服务。``PI05_CHECKPOINT_PATH`` 指向 +并可在本地启动双臂 Franka VLA 服务。``PI05_CHECKPOINT_PATH`` 指向 训练好的 Pi-05 checkpoint，``DUAL_FRANKA_REPO_ID`` 是用于查找对应归一化统计的数据集 ID： .. code-block:: bash diff --git a/pyproject.toml b/pyproject.toml index c772b7db1..d6d2771d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,10 +29,10 @@ classifiers = [ dependencies = [ "pydantic-ai-slim[anthropic,openai]>=2.1", "pydantic>=2", + "docstring-parser>=0.16,<1", "fastapi>=0.110", "uvicorn>=0.27", "httpx>=0.27", - "jsonschema>=4.18", "claude-agent-sdk>=0.1.60", "openai-codex>=0.1.0b3", "mcp>=1.23.0,<2.0.0", @@ -61,6 +61,7 @@ test = [ "coverage", "pytest", "pytest-timeout", + "scipy", "omegaconf", "gymnasium", ] diff --git a/robots/dual_franka/dual_franka_manual_call.py b/robots/dual_franka/dual_franka_manual_call.py index d9301bc9b..c0e2be6f3 100755 --- a/robots/dual_franka/dual_franka_manual_call.py +++ b/robots/dual_franka/dual_franka_manual_call.py @@ -30,6 +30,7 @@ import sys import time from pathlib import Path +from threading import Event from typing import Any import numpy as np @@ -41,20 +42,18 @@ from robots.dual_franka import perception as dual_franka_perception from robots.dual_franka.runtime_config import DEFAULT_CONFIG from robots.dual_franka.tasks import get_dual_franka_task -from robots.dual_franka.toolkit import DualFrankaToolkit +from robots.dual_franka.toolkit import DualFrankaRuntime, DualFrankaToolkit from robots.dual_franka.tools import ( - TOOLS_SPEC, - DualFrankaPrimitives, + DUAL_FRANKA_TOOLS, dump_state, - view_env_state, ) from robots.franka.runtime_config import set_calibration_path, set_robot_config_path -from robots.franka.tools import view_camera_meta from rpent.dashboard.events import NullDashboardEventSink from rpent.memory.manager import MemoryManager from rpent.robots.components.pi05_vla_client import Pi05VLAClient from rpent.robots.components.sam3_client import Sam3Client from rpent.session import EnvState +from rpent.tools import ToolContext from rpent.utils.config import get_repo_root from rpent.utils.rpc import make_rpc_client, wait_for_ready @@ -69,7 +68,14 @@ def _tool_spec_map(*, include_manual: bool = False) -> dict[str, dict[str, Any]]: - specs = {str(spec["name"]): spec for spec in TOOLS_SPEC} + specs = { + item.name: { + "name": item.name, + "description": item.description, + "input_schema": item.input_schema, + } + for item in DUAL_FRANKA_TOOLS + } if include_manual: specs[str(_RESET_SPEC["name"])] = _RESET_SPEC return specs @@ -334,31 +340,50 @@ def _default_output_dir(primitive: str) -> Path: return get_repo_root() / "logs" / f"{stamp}-manual-{primitive}" +def _call_native_tool( + name: str, + params: dict[str, Any], + runtime: DualFrankaRuntime, + state: EnvState | None = None, +) -> dict[str, Any]: + item = next(item for item in DUAL_FRANKA_TOOLS if item.name == name) + args = item.args_schema.model_validate(params) + ctx = ToolContext( + robot=runtime, + state=state, + memory=None, + output_dir=Path("."), + record_frame=lambda frame: None, + _cancel_event=Event(), + ) + return item.handler(**args.model_dump(), ctx=ctx).to_dict() + + def _call_readonly_tool( primitive: str, params: dict[str, Any], *, - primitives: DualFrankaPrimitives, + runtime: DualFrankaRuntime, state: EnvState, sam3_client: Sam3Client | None, dump_state_enabled: bool, ) -> dict[str, Any]: if primitive == "describe_dual_franka_setup": - return primitives.describe_dual_franka_setup() + return _call_native_tool(primitive, params, runtime, state) if primitive == "view_env_state": if dump_state_enabled: dump_state( - primitives, + runtime, state, command={"action": "view_env_state", "params": params}, result=None, elapsed_s=None, ) - return view_env_state(state=state, **params) + return _call_native_tool(primitive, params, runtime, state) return { "step_idx": None, - "state": primitives.env.get_robot_state(), - "camera_meta": primitives.env.get_camera_meta() or {}, + "state": runtime.env.get_robot_state(), + "camera_meta": runtime.env.get_camera_meta() or {}, "images": [], "artifact_images": [], "note": ( @@ -369,14 +394,14 @@ def _call_readonly_tool( if primitive == "view_camera_meta": if dump_state_enabled: dump_state( - primitives, + runtime, state, command={"action": "view_camera_meta", "params": params}, result=None, elapsed_s=None, ) - return view_camera_meta(state=state, **params) - return {"step": None, "camera_meta": primitives.env.get_camera_meta() or {}} + return _call_native_tool(primitive, params, runtime, state) + return {"step": None, "camera_meta": runtime.env.get_camera_meta() or {}} if primitive == "back_project": if state.latest_step is None: if not dump_state_enabled: @@ -385,7 +410,7 @@ def _call_readonly_tool( "--no-dump-state so the manual tool can capture one first." ) dump_state( - primitives, + runtime, state, command={"action": "snapshot_before_back_project", "params": {}}, result=None, @@ -400,7 +425,7 @@ def _call_readonly_tool( "--no-dump-state so the manual tool can capture one first." ) dump_state( - primitives, + runtime, state, command={"action": "snapshot_before_segment", "params": {}}, result=None, @@ -419,23 +444,12 @@ def _call_mutating_primitive( params: dict[str, Any], *, env: ManualDualFrankaEnv, - primitives: DualFrankaPrimitives, + runtime: DualFrankaRuntime, ) -> dict[str, Any]: if primitive == "reset": return env.reset() - if primitive == "move_delta": - return primitives.move_delta(**params) - if primitive == "rotate_delta": - return primitives.rotate_delta(**params) - if primitive == "open_gripper": - return primitives.open_gripper(**params) - if primitive == "close_gripper": - return primitives.close_gripper(**params) - if primitive == "recover_joint_posture": - return primitives.recover_joint_posture(**params) if primitive in _registered_tool_names(): - handler = getattr(primitives, primitive) - return handler(**params) + return _call_native_tool(primitive, params, runtime) raise KeyError(primitive) @@ -468,7 +482,7 @@ def _call_toolkit_tool( state_output_dir=output_dir, ) try: - return toolkit.execute_tool(primitive, params).result + return toolkit.execute_tool(primitive, params).to_dict() finally: toolkit.close() @@ -523,12 +537,11 @@ def main() -> int: sam3_client = Sam3Client(sam3_rpc) state = EnvState(output_dir) - primitives = DualFrankaPrimitives( + runtime = DualFrankaRuntime( env=env, model=model, task_description=task.instruction, vla_instruction=task.vla_instruction, - check_cancelled=lambda: None, sam3_client=sam3_client, ) @@ -546,15 +559,17 @@ def main() -> int: result = _call_readonly_tool( primitive, params, - primitives=primitives, + runtime=runtime, state=state, sam3_client=sam3_client, dump_state_enabled=not args.no_dump_state, ) else: - if primitive in _registered_tool_names() and not hasattr( - primitives, primitive - ): + if primitive in { + "request_scene_reset", + "request_operator_verdict", + "finish", + }: used_toolkit_tool = True result = _call_toolkit_tool( primitive, @@ -570,7 +585,7 @@ def main() -> int: primitive, params, env=env, - primitives=primitives, + runtime=runtime, ) except KeyError: known = sorted(_manual_primitive_names()) @@ -590,7 +605,7 @@ def main() -> int: ): try: dump_state( - primitives, + runtime, state, command={"action": primitive, "params": params}, result=result, diff --git a/robots/dual_franka/perception.py b/robots/dual_franka/perception.py index e6b4f6f2e..d2aab425a 100644 --- a/robots/dual_franka/perception.py +++ b/robots/dual_franka/perception.py @@ -30,7 +30,6 @@ load_mapping, ) from rpent.session import EnvState, StepRecord -from rpent.tools.toolkit import readonly from rpent.utils.transforms import ( invert_transform, transform_points, @@ -44,7 +43,6 @@ class DualFrankaPerceptionError(ValueError): """Raised when a dual-Franka perception artifact is missing or invalid.""" -@readonly def back_project( *, # PhysicalAgent alignment note: D455 is the deployed clean-desk primary @@ -71,7 +69,6 @@ def back_project( ) -@readonly def segment( *, # Same deployment default as back_project: SAM3 can run on any registered diff --git a/robots/dual_franka/toolkit.py b/robots/dual_franka/toolkit.py index cd6c670de..52c8ed489 100644 --- a/robots/dual_franka/toolkit.py +++ b/robots/dual_franka/toolkit.py @@ -21,16 +21,16 @@ import threading import time from collections.abc import Callable +from dataclasses import dataclass, field, replace from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any -from robots.dual_franka import perception as dual_franka_perception from robots.dual_franka import tools as dual_franka_tools -from robots.franka import tools as franka_tools -from robots.franka.toolkit import FrankaToolkit -from rpent.dashboard.events import DashboardEventSink -from rpent.tools.toolkit import ToolCancelled, ToolResult, readonly +from robots.franka.toolkit import FrankaRuntime, FrankaToolkit +from rpent.dashboard.events import DashboardEventSink, StepRecordEvent +from rpent.session import StepRecord +from rpent.tools import Tool, ToolCancelled, ToolContext, ToolResult from rpent.utils.logging import get_output_dir if TYPE_CHECKING: @@ -52,6 +52,13 @@ } +@dataclass +class DualFrankaRuntime(FrankaRuntime): + sam3_client: Any | None = None + vla_instruction: str | None = None + session: DualFrankaToolkit | None = field(default=None, init=False, repr=False) + + class DualFrankaToolkit(FrankaToolkit): """Dual-arm tools and the attended exploration lifecycle from PR #176. @@ -59,8 +66,8 @@ class DualFrankaToolkit(FrankaToolkit): attempt boundaries are additional artifacts, not simulator termination flags. """ - _tools_module = dual_franka_tools - _primitives_cls = dual_franka_tools.DualFrankaPrimitives + _runtime_type = DualFrankaRuntime + _robot_tools = dual_franka_tools.DUAL_FRANKA_TOOLS def __init__( self, @@ -98,6 +105,31 @@ def __init__( memory=memory, state_output_dir=state_output_dir, ) + self._robot.session = self + self._tools = { + name: replace(item, handler=partial(self._invoke_guarded, item)) + if self._mode == "exploration" + and name in _MOTION_TOOLS | {"back_project", "segment"} + else item + for name, item in self._tools.items() + if self._mode == "exploration" or name not in _EXPLORATION_ONLY_TOOLS + } + + def _invoke_guarded( + self, tool: Tool, *, ctx: ToolContext, **kwargs: Any + ) -> ToolResult: + handler = partial(tool.handler, ctx=ctx) + guard = ( + self._guard_motion + if tool.name in _MOTION_TOOLS + else self._current_perception + ) + result = guard(handler, **kwargs) + return ( + result + if isinstance(result, ToolResult) + else dual_franka_tools._result(result) + ) @property def direct_verdict_requested(self) -> bool: @@ -127,16 +159,17 @@ def raise_if_cancelled(self) -> None: raise ToolCancelled( "operator submitted a terminal verdict; stopping exploration" ) - super().raise_if_cancelled() + with self._operation_lock: + operation = self._active_operation + cancelled = operation is not None and operation.cancel_event.is_set() + if cancelled: + raise ToolCancelled("Tool call cancelled.") def execute_tool(self, name: str, input_dict: dict[str, Any]) -> ToolResult: if self._direct_verdict_event.is_set(): return ToolResult( - name=name, - result={ - "error": "operator submitted a terminal verdict; exploration is closing", - "motion_refused": True, - }, + error="operator submitted a terminal verdict; exploration is closing", + data={"motion_refused": True}, ) return super().execute_tool(name, input_dict) @@ -159,12 +192,16 @@ def finalize_direct_verdict(self) -> dict[str, Any]: self._event("verdict", verdict="abort", source="interactive_command") self._event("finish", **result) return result - self.get_env_state( - command={"action": "observe_for_verdict"}, result={}, elapsed_s=0.0 + self._capture_observation( + command={"action": "observe_for_verdict"}, + result=ToolResult(), + elapsed_s=0.0, ) self._validate_observation() record = self.state.latest_record() - self._publish_step(record) + self._dashboard_events.emit( + StepRecordEvent(record=record, env_state=self.state) + ) self._scene_ready = True self._operator_verdict = self._direct_verdict self._operator_notes = ( @@ -187,17 +224,6 @@ def finalize_direct_verdict(self) -> dict[str, Any]: self._event("finish", **result) return result - @readonly - def _describe_exploration_setup(self, inner): - result = inner() - result["phase"] = "exploration" - result["reset_policy"] = ( - "No automatic reset was performed by the client. Before motion in " - "each session call request_scene_reset and wait for operator confirmation." - ) - result["scene_ready"] = self._scene_ready - return result - def _clear_verdict(self) -> None: self._operator_verdict = None self._operator_notes = "" @@ -216,7 +242,6 @@ def _guard_motion(self, inner, **kwargs): self._clear_verdict() return inner(**kwargs) - @readonly def _current_perception(self, inner, **kwargs): step = kwargs.get("step") if not self._scene_ready or ( @@ -285,7 +310,7 @@ def _request_scene_reset( "operator_aborted": self._operator_aborted, } # Never count a failed reset or a failed post-reset observation as a new attempt. - result = self._primitives.reset() + result = self._robot.env.reset() if ( not isinstance(result, dict) or result.get("ok") is not True @@ -300,7 +325,6 @@ def _request_scene_reset( "notice": "Scene restored by operator; robot posture reset. Re-localize from the new images.", } - @readonly def _request_operator_verdict( self, question="Does the current scene satisfy the task success criteria?" ): @@ -312,12 +336,16 @@ def _request_operator_verdict( ): return {"error": "verdict refused; no active confirmed attempt"} # Save the evidence being judged using the existing camera/state logger. - self.get_env_state( - command={"action": "observe_for_verdict"}, result={}, elapsed_s=0.0 + self._capture_observation( + command={"action": "observe_for_verdict"}, + result=ToolResult(), + elapsed_s=0.0, ) self._validate_observation() record = self.state.latest_record() - self._publish_step(record) + self._dashboard_events.emit( + StepRecordEvent(record=record, env_state=self.state) + ) response = self._ask_operator( f"{question}\nAttempt {self._attempt}, observation step {record.step_idx}. " "Reply success, failure, continue, or abort; optional notes may follow.", @@ -345,7 +373,6 @@ def _request_operator_verdict( "operator_aborted": self._operator_aborted, } - @readonly def _guarded_finish(self, inner, **kwargs): if not self._operator_aborted: if self._operator_verdict is None: @@ -366,16 +393,17 @@ def _guarded_finish(self, inner, **kwargs): self._event("finish", **result) return result - def get_env_state(self, *, command, result, elapsed_s): + def _capture_observation(self, *, command, result: ToolResult, elapsed_s): if self._mode != "exploration": - return super().get_env_state( + return super()._capture_observation( command=command, result=result, elapsed_s=elapsed_s ) try: - output = super().get_env_state( + output, images = super()._capture_observation( command=command, result=result, elapsed_s=elapsed_s ) - if command["action"] == "request_scene_reset" and result.get( + payload = result.to_dict() + if command["action"] == "request_scene_reset" and payload.get( "scene_reset_confirmed" ): self._validate_observation() @@ -393,9 +421,9 @@ def get_env_state(self, *, command, result, elapsed_s): self.state.save("exploration.json", status) output["exploration"] = status # Keep the original record layout, and expose lifecycle errors to the planner. - if result.get("error"): - output["error"] = result["error"] - return output + if result.is_error: + output["error"] = result.error + return output, images except Exception: self._scene_ready = False self._clear_verdict() @@ -476,57 +504,6 @@ def write_recipe(self, recipe_tag: str) -> str | None: audit_path.write_text(json.dumps(audit, indent=2) + "\n") return str(recipe) - def _register_tools(self) -> None: - state_handlers = { - "view_env_state": partial( - dual_franka_tools.view_env_state, state=self._state - ), - "view_camera_meta": partial( - franka_tools.view_camera_meta, - state=self._state, - ), - "back_project": partial( - dual_franka_perception.back_project, - state=self._state, - ), - "segment": partial( - dual_franka_perception.segment, - state=self._state, - sam3_client=getattr(self._primitives, "_sam3_client", None), - ), - "request_scene_reset": self._request_scene_reset, - "request_operator_verdict": self._request_operator_verdict - if self._mode == "exploration" - else self._evaluation_verdict, - } - for spec in self._tools_module.TOOLS_SPEC: - name = spec["name"] - if name in _EXPLORATION_ONLY_TOOLS and self._mode != "exploration": - continue - handler = state_handlers.get(name) or getattr(self._primitives, name, None) - if handler is None: - continue - if self._mode == "exploration": - if name in _MOTION_TOOLS: - handler = partial(self._guard_motion, handler) - elif name in {"back_project", "segment"}: - handler = partial(self._current_perception, handler) - elif name == "describe_dual_franka_setup": - handler = partial(self._describe_exploration_setup, handler) - self.add_tool(name, spec, handler) - if self._mode != "exploration": - finish_spec, finish_handler = self._tools["finish"] - self.add_tool( - "finish", finish_spec, partial(self._evaluation_finish, finish_handler) - ) - if self._mode == "exploration": - finish_spec, finish_handler = self._tools["finish"] - self.add_tool( - "finish", - finish_spec, - partial(self._guarded_finish, finish_handler), - ) - def _read_operator_line(self, prompt: str) -> str | None: if sys.stdin is None or not sys.stdin.isatty(): return None @@ -535,7 +512,6 @@ def _read_operator_line(self, prompt: str) -> str | None: except EOFError: return None - @readonly def _evaluation_verdict( self, question: str = "Does the current real-robot scene satisfy the task?", @@ -581,7 +557,6 @@ def _evaluation_verdict( "attempt": self._attempt, } - @readonly def _evaluation_finish(self, inner: Any, **kwargs: Any) -> dict[str, Any]: """Require real-robot operator feedback before finishing a task.""" if self._operator_verdict is None: @@ -615,3 +590,21 @@ def _evaluation_finish(self, inner: Any, **kwargs: Any) -> dict[str, Any]: if self._operator_notes: result.setdefault("operator_notes", self._operator_notes) return result + + def _dump_state( + self, + *, + command: dict[str, Any] | None, + result: dict[str, Any] | None, + elapsed_s: float | None, + ) -> StepRecord: + return dual_franka_tools.dump_state( + self._robot, + self.state, + command=command, + result=result, + elapsed_s=elapsed_s, + ) + + def _build_observation(self, record: StepRecord) -> ToolResult: + return dual_franka_tools.build_observation(self.state, record) diff --git a/robots/dual_franka/tools.py b/robots/dual_franka/tools.py index fc38c974a..0f1cca385 100644 --- a/robots/dual_franka/tools.py +++ b/robots/dual_franka/tools.py @@ -12,319 +12,101 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Dual-Franka planner tools, primitives, and canonical state capture.""" +"""Native dual-Franka tools and canonical RGB-D state capture.""" from __future__ import annotations import time -from collections.abc import Callable, Sequence -from typing import Any +from dataclasses import replace +from typing import TYPE_CHECKING, Annotated, Any, Literal import numpy as np - -from robots.franka.tools import FrankaPrimitives, coerce_vec3 +from pydantic import BeforeValidator, Field, FiniteFloat + +from robots.dual_franka import perception +from robots.franka.tools import ( + _result, +) +from robots.franka.tools import ( + view_camera_meta as franka_view_camera_meta, +) from rpent.session import EnvState, StepRecord -from rpent.tools.toolkit import readonly - -_ARM_PROPERTY = { - "type": "string", - "enum": ["left", "right"], - "description": "Which arm to command; the other arm is left uncommanded.", -} - -TOOLS_SPEC = [ - { - "name": "describe_dual_franka_setup", - "description": ( - "Read the dual-Franka runtime conventions, camera aliases, VLA " - "policy conditioning text, semantic stop rules, and available " - "primitive names before acting. This is read-only." - ), - "input_schema": {"type": "object", "properties": {}}, - }, - { - "name": "view_env_state", - "description": ( - # This wording documents the current lab observation policy used for - # PhysicalAgent log alignment. The implementation itself reads - # cameras.agent_observation from robot config, so non-D455 setups - # should update the config and this deployment-oriented prose. - "Read a dual-Franka state snapshot. The D455 image is returned inline; " - "left_wrist, base, and right_wrist are returned as artifact paths for " - "targeted read_image inspection." - ), - "input_schema": { - "type": "object", - "properties": {"step": {"type": "integer", "default": -1}}, - }, - }, - { - "name": "view_camera_meta", - "description": "Read camera intrinsics, serials, and projection metadata for the dual-Franka rig.", - "input_schema": { - "type": "object", - "properties": {"step": {"type": "integer", "default": -1}}, - }, - }, - { - "name": "back_project", - "description": ( - "Back-project one pixel from a registered RGBD camera view into " - "shared right-base coordinates. Use a camera listed by " - "view_env_state/view_camera_meta; default is the configured primary " - "metric localization camera." - ), - "input_schema": { - "type": "object", - "properties": { - "camera": { - "type": "string", - "default": "d455", - "description": ( - "Registered projection view name, e.g. d455 or base. " - "Valid names come from perception.projection_views and " - "the current state's saved artifacts." - ), - }, - "row": {"type": "integer", "minimum": 0}, - "col": {"type": "integer", "minimum": 0}, - "target_name": {"type": "string", "default": "target"}, - "step": {"type": "integer"}, - "window_radius": {"type": "integer", "minimum": 0, "default": 2}, - }, - "required": ["row", "col"], - }, - }, - { - "name": "segment", - "description": ( - "Use SAM3 on a registered RGB image with either a text prompt or one " - "positive [row, col] point, return a mask overlay for verification, " - "and estimate the mask median point in shared right-base coordinates." - ), - "input_schema": { - "type": "object", - "properties": { - "camera": { - "type": "string", - "default": "d455", - "description": ( - "Registered projection view name, e.g. d455 or base. " - "Valid names come from perception.projection_views and " - "the current state's saved artifacts." - ), - }, - "prompt": { - "type": "string", - "default": "", - "description": ( - # Phrase hints are intentionally copied from the live - # clean-desk debugging runs where SAM3 often confused - # the cardboard-box interior with nearby table/metal - # basket surfaces. They are task/deployment hints, not - # an inherent property of the generic segment tool. - "Text prompt for SAM3. Prefer short object/relation phrases; " - "for the clean-desk box use 'white interior of the black " - "cardboard box' or 'cardboard box'. Avoid over-specific " - "surface words such as 'floor' when grounding is weak. " - "Provide exactly one of prompt or point." - ), - }, - "point": { - "type": "array", - "items": {"type": "integer"}, - "minItems": 2, - "maxItems": 2, - "description": ( - "Positive SAM3 point in camera image coordinates [row, col]. " - "Provide exactly one of prompt or point." - ), - }, - "target_name": {"type": "string", "default": "target"}, - "step": {"type": "integer"}, - "min_score": { - "type": "number", - "minimum": 0.0, - "maximum": 1.0, - "default": 0.2, - }, - "min_valid_depth_pixels": { - "type": "integer", - "minimum": 1, - "default": 25, - }, - }, - }, - }, - { - "name": "move_delta", - "description": "Move one Franka TCP by a bounded world-frame xyz delta in meters.", - "input_schema": { - "type": "object", - "properties": { - "arm": _ARM_PROPERTY, - "delta_xyz": { - "type": "array", - "items": {"type": "number"}, - "minItems": 3, - "maxItems": 3, - }, - }, - "required": ["arm", "delta_xyz"], - }, - }, - { - "name": "rotate_delta", - "description": "Rotate one Franka TCP by a bounded world-frame rpy delta in radians.", - "input_schema": { - "type": "object", - "properties": { - "arm": _ARM_PROPERTY, - "delta_rpy": { - "type": "array", - "items": {"type": "number"}, - "minItems": 3, - "maxItems": 3, - }, - }, - "required": ["arm", "delta_rpy"], - }, - }, - { - "name": "open_gripper", - "description": "Open one Franka gripper and wait for the command to settle.", - "input_schema": { - "type": "object", - "properties": {"arm": _ARM_PROPERTY}, - "required": ["arm"], - }, - }, - { - "name": "close_gripper", - "description": "Close one Franka gripper and wait for the command to settle.", - "input_schema": { - "type": "object", - "properties": {"arm": _ARM_PROPERTY}, - "required": ["arm"], - }, - }, - { - "name": "recover_joint_posture", - "description": ( - "Reset both arms to their healthy configured joint posture while " - "preserving each gripper's open/closed state. Closed grippers are " - "re-commanded before/after the joint reset so held objects stay clamped, " - "then both TCPs return near their prior poses." - ), - "input_schema": { - "type": "object", - "properties": { - "reason": {"type": "string", "default": ""}, - "return_to_start": {"type": "boolean", "default": True}, - }, - }, - }, - { - "name": "request_scene_reset", - "description": ( - "Exploration-only real-robot reset gate. Ask the human operator to " - "remove/secure held objects and restore the tabletop scene for another " - "attempt, wait for terminal confirmation, then reset the robot posture. " - "This does not automatically restore physical objects like a simulator." - ), - "input_schema": { - "type": "object", - "properties": { - "reason": { - "type": "string", - "description": "Why the scene needs to be restored.", - }, - "expected_scene_state": { - "type": "string", - "default": "", - "description": ( - "Short instruction for the operator describing the " - "desired restored layout." - ), - }, - }, - "required": ["reason"], - }, - }, - { - "name": "request_operator_verdict", - "description": ( - "Exploration-only human feedback gate. Ask the operator to mark " - "the current physical task state as success, failure, or continue " - "before the planner finishes or starts another attempt." - ), - "input_schema": { - "type": "object", - "properties": { - "question": { - "type": "string", - "default": "Does the current real-robot scene satisfy the task?", - }, - }, - }, - }, - *[ - { - "name": name, - "description": description, - "input_schema": { - "type": "object", - "properties": { - "prompt": { - "type": "string", - "description": ( - "Planner-facing segment intent. This is recorded in " - "the tool result; the current live clean-desk " - "checkpoint still receives its fixed training " - "instruction during policy inference." - ), - }, - "max_chunks": { - "type": "integer", - "minimum": 1, - "maximum": 20, - "default": 20, - }, - }, - "required": ["prompt"], - }, - } - for name, description in ( - ( - "vla_right_grasp", - "Run the learned right-grasp VLA segment. The active task prompt " - "decides which object is currently allowed; this tool only " - "defines the capability boundary: right gripper closes and the " - "right TCP lifts.", - ), - ( - "vla_handoff", - "Run the learned bimanual handoff VLA segment. The capability " - "boundary is right-gripper release followed by the configured " - "settle delay; do not rule-base pre-position either arm for it.", - ), - ( - "vla_left_place", - "Run the learned left-placement VLA segment. The active task " - "decides the destination; this tool only defines the capability " - "boundary: left gripper opens and the left TCP lifts.", - ), - ) - ], -] +from rpent.tools import ToolContext, ToolResult, readonly, tool + +if TYPE_CHECKING: + from robots.dual_franka.toolkit import DualFrankaRuntime -def coerce_arm(value: Any) -> str: - """Return exactly ``'left'`` or ``'right'`` or raise a useful error.""" - arm = str(value).strip().lower() - if arm not in {"left", "right"}: +def _normalize_arm(value: str) -> str: + """Normalize string input before Literal validation; reject other types.""" + if not isinstance(value, str): raise ValueError("arm must be exactly 'left' or 'right'") - return arm + return value.strip().lower() + + +Arm = Annotated[Literal["left", "right"], BeforeValidator(_normalize_arm)] +view_camera_meta = replace( + franka_view_camera_meta, + description="Read camera intrinsics, serials, and projection metadata for the dual-Franka rig.", +) + + +@tool +def move_delta( + arm: Arm, + delta_xyz: Annotated[list[FiniteFloat], Field(min_length=3, max_length=3)], + *, + ctx: ToolContext[DualFrankaRuntime], +) -> ToolResult: + """Move one Franka TCP by a bounded world-frame xyz delta in meters. + + Args: + arm: Which arm to command; the other arm is left uncommanded. + """ + ctx.check_cancelled() + return _result( + ctx.robot.env.move_delta(arm, np.asarray(delta_xyz, dtype=np.float32)) + ) + + +@tool +def rotate_delta( + arm: Arm, + delta_rpy: Annotated[list[FiniteFloat], Field(min_length=3, max_length=3)], + *, + ctx: ToolContext[DualFrankaRuntime], +) -> ToolResult: + """Rotate one Franka TCP by a bounded world-frame rpy delta in radians. + + Args: + arm: Which arm to command; the other arm is left uncommanded. + """ + ctx.check_cancelled() + return _result( + ctx.robot.env.rotate_delta(arm, np.asarray(delta_rpy, dtype=np.float32)) + ) + + +@tool +def open_gripper(arm: Arm, *, ctx: ToolContext[DualFrankaRuntime]) -> ToolResult: + """Open one Franka gripper and wait for the command to settle. + + Args: + arm: Which arm to command; the other arm is left uncommanded. + """ + ctx.check_cancelled() + return _result(ctx.robot.env.set_gripper(arm, open=True)) + + +@tool +def close_gripper(arm: Arm, *, ctx: ToolContext[DualFrankaRuntime]) -> ToolResult: + """Close one Franka gripper and wait for the command to settle. + + Args: + arm: Which arm to command; the other arm is left uncommanded. + """ + ctx.check_cancelled() + return _result(ctx.robot.env.set_gripper(arm, open=False)) def _camera_alias_from_key(raw_key: Any) -> str | None: @@ -348,372 +130,534 @@ def _agent_observation_policy(meta: dict[str, Any] | None) -> dict[str, list[str } -def _image_bytes_key(index: int) -> str: - keys = [ - "_image_bytes", - "_image_cam_bytes", - "_image_nav_bytes", - "_image_wrist_bytes", - ] - return keys[index] - - -class DualFrankaPrimitives(FrankaPrimitives): - """Safe agent-facing operations over a remote dual-Franka environment.""" - - def __init__( - self, - *, - env: Any, - model: Any | None, - task_description: str, - check_cancelled: Callable[[], None], - sam3_client: Any | None = None, - vla_instruction: str | None = None, - ) -> None: - super().__init__( - env=env, - model=model, - task_description=task_description, - check_cancelled=check_cancelled, +def _run_named_vla_skill( + ctx: ToolContext[DualFrankaRuntime], + *, + skill_name: str, + boundary: str, + prompt: str, + max_chunks: int = 20, +) -> dict[str, Any]: + if ctx.robot.model is None: + raise RuntimeError(f"{skill_name} requires --vla-endpoint") + requested_prompt = str(prompt).strip() + if not requested_prompt: + raise ValueError("prompt must be non-empty") + # Task configuration owns policy conditioning; retain planner intent in logs. + effective_prompt = ctx.robot.vla_instruction or ctx.robot.task_description + + state = ctx.robot.env.get_robot_state() + start_state = state + previous_left_open = bool(state["left_arm"]["gripper_open"]) + previous_right_open = bool(state["right_arm"]["gripper_open"]) + event_z: float | None = None + event_time: float | None = None + steps_after_event = 0 + chunks_executed = 0 + steps_executed = 0 + event_step: int | None = None + boundary_reached = False + terminated = False + truncated = False + action_count = 0 + action_min: np.ndarray | None = None + action_max: np.ndarray | None = None + action_sum: np.ndarray | None = None + first_action: list[float] | None = None + first_policy_state: list[float] | None = None + first_policy_state_shape: list[int] | None = None + started_at = time.perf_counter() + + for _ in range(int(max_chunks)): + ctx.check_cancelled() + observation = dict(ctx.robot.env.get_observation()) + if first_policy_state is None: + obs_states = np.asarray(observation.get("states"), dtype=np.float32) + first_policy_state_shape = list(obs_states.shape) + first_policy_state = ( + np.round(obs_states.reshape(-1)[:20], 5).astype(float).tolist() + ) + observation["task_descriptions"] = effective_prompt + actions = np.asarray( + ctx.robot.model.predict(observation, options={"mode": "eval"}), + dtype=np.float32, ) - self._sam3_client = sam3_client - self._vla_instruction = vla_instruction or task_description - - def move_delta(self, arm: str, delta_xyz: Sequence[float]) -> dict[str, Any]: - self._check_cancelled() - return self.env.move_delta( - coerce_arm(arm), coerce_vec3(delta_xyz, name="delta_xyz") + if actions.ndim != 2 or actions.shape[1] != 20: + raise RuntimeError( + f"{skill_name} expected [chunk, 20] actions, got {actions.shape}" + ) + if not np.isfinite(actions).all(): + raise RuntimeError(f"{skill_name} received non-finite VLA actions") + chunks_executed += 1 + if first_action is None: + first_action = np.round(actions[0], 5).astype(float).tolist() + if action_count == 0: + action_min = actions.min(axis=0) + action_max = actions.max(axis=0) + action_sum = actions.sum(axis=0) + else: + action_min = np.minimum(action_min, actions.min(axis=0)) + action_max = np.maximum(action_max, actions.max(axis=0)) + action_sum = action_sum + actions.sum(axis=0) + action_count += int(actions.shape[0]) + + for action in actions: + ctx.check_cancelled() + result = ctx.robot.env.chunk_step(action[None, :]) + steps_executed += 1 + terminated = terminated or bool(result.get("terminated")) + truncated = truncated or bool(result.get("truncated")) + state = ctx.robot.env.get_robot_state() + left = state["left_arm"] + right = state["right_arm"] + left_open = bool(left["gripper_open"]) + right_open = bool(right["gripper_open"]) + + if boundary == "grasp": + if event_z is None and previous_right_open and not right_open: + event_z = float(right["tcp_pose"][2]) + event_step = steps_executed + steps_after_event = 0 + elif event_z is not None: + steps_after_event += 1 + boundary_reached = ( + not right_open + and steps_after_event >= 2 + and float(right["tcp_pose"][2]) - event_z >= 0.15 + ) + elif boundary == "handoff": + if event_time is None and not previous_right_open and right_open: + event_time = time.monotonic() + event_step = steps_executed + steps_after_event = 0 + elif event_time is not None: + steps_after_event += 1 + boundary_reached = ( + right_open + and steps_after_event >= 2 + and time.monotonic() - event_time >= 1.5 + ) + elif boundary == "place": + if event_z is None and not previous_left_open and left_open: + event_z = float(left["tcp_pose"][2]) + event_step = steps_executed + steps_after_event = 0 + elif event_z is not None: + steps_after_event += 1 + boundary_reached = ( + left_open + and steps_after_event >= 2 + and float(left["tcp_pose"][2]) - event_z >= 0.10 + ) + else: # pragma: no cover - internal programming guard + raise ValueError(f"unknown VLA skill boundary: {boundary}") + + previous_left_open = left_open + previous_right_open = right_open + if boundary_reached or terminated or truncated: + break + if boundary_reached or terminated or truncated: + break + + stop_rule: dict[str, Any] = { + "phase": boundary, + "skill_name": skill_name, + "step_count": steps_executed, + "event_step": event_step, + "success_claim": False, + } + if boundary == "grasp": + stop_rule.update( + { + "condition": "right_gripper_closed_then_lifted", + "right_close_step": event_step, + "right_close_z": event_z, + "right_current_z": float(state["right_arm"]["tcp_pose"][2]), + "lift_m": ( + float(state["right_arm"]["tcp_pose"][2]) - event_z + if event_z is not None + else None + ), + "threshold_m": 0.15, + } ) - - def rotate_delta(self, arm: str, delta_rpy: Sequence[float]) -> dict[str, Any]: - self._check_cancelled() - return self.env.rotate_delta( - coerce_arm(arm), coerce_vec3(delta_rpy, name="delta_rpy") + elif boundary == "handoff": + stop_rule.update( + { + "condition": "right_gripper_opened_then_delay", + "right_open_step": event_step, + "elapsed_after_open_s": ( + time.monotonic() - event_time if event_time is not None else None + ), + "delay_s": 1.5, + "right_current_open": bool(state["right_arm"]["gripper_open"]), + } + ) + elif boundary == "place": + stop_rule.update( + { + "condition": "left_gripper_opened_then_lifted", + "left_open_step": event_step, + "left_open_z": event_z, + "left_current_z": float(state["left_arm"]["tcp_pose"][2]), + "lift_m": ( + float(state["left_arm"]["tcp_pose"][2]) - event_z + if event_z is not None + else None + ), + "threshold_m": 0.10, + } ) - def open_gripper(self, arm: str) -> dict[str, Any]: - self._check_cancelled() - return self.env.set_gripper(coerce_arm(arm), open=True) + action_summary: dict[str, Any] = { + "count": action_count, + "shape": [action_count, 20], + "finite": True, + } + if action_count and action_min is not None and action_max is not None: + action_summary.update( + { + "min": np.round(action_min, 5).astype(float).tolist(), + "max": np.round(action_max, 5).astype(float).tolist(), + "mean": np.round(action_sum / action_count, 5).astype(float).tolist(), + "first_action": first_action, + "first_policy_state_shape": first_policy_state_shape, + "first_policy_state": first_policy_state, + } + ) - def close_gripper(self, arm: str) -> dict[str, Any]: - self._check_cancelled() - return self.env.set_gripper(coerce_arm(arm), open=False) + return { + "ok": boundary_reached and not (terminated or truncated), + "skill_name": skill_name, + "boundary": boundary, + "requested_prompt": requested_prompt, + "effective_policy_prompt": effective_prompt, + "prompt_overridden": requested_prompt != effective_prompt, + "boundary_reached": boundary_reached, + "chunks_executed": chunks_executed, + "steps_executed": steps_executed, + "terminated": terminated, + "truncated": truncated, + "stop_rule": stop_rule, + "action_summary": action_summary, + "elapsed_s": time.perf_counter() - started_at, + "vla_start_robot_state": start_state, + "robot_state": state, + } - def recover_joint_posture( - self, reason: str = "", return_to_start: bool = True - ) -> dict[str, Any]: - self._check_cancelled() - return self.env.recover_joint_posture( - reason=str(reason), return_to_start=bool(return_to_start) - ) - @readonly - def describe_dual_franka_setup(self) -> dict[str, Any]: - """Return PhysicalAgent-compatible setup guidance without moving hardware.""" - meta = self.env.meta - observation_policy = _agent_observation_policy(meta) - return { - "ok": True, - "phase": "strict", - "reset_policy": ( - "The runner reset the robot at startup. Do not call reset during " - "a task unless reset is explicitly exposed and there is a clear " - "robot-side reason." +@tool +@readonly +def describe_dual_franka_setup(*, ctx: ToolContext[DualFrankaRuntime]) -> ToolResult: + """Read the dual-Franka runtime conventions, camera aliases, VLA policy conditioning text, semantic stop rules, and available primitive names before acting. This is read-only.""" + meta = ctx.robot.env.meta + observation_policy = _agent_observation_policy(meta) + data = { + "ok": True, + "phase": "strict", + "reset_policy": ( + "The runner reset the robot at startup. Do not call reset during " + "a task unless reset is explicitly exposed and there is a clear " + "robot-side reason." + ), + "coordinate_frame": "right_base", + "camera_aliases": meta.get("observation_camera_map", {}), + "projection_views": meta.get("projection_views", {}), + "agent_observation": observation_policy, + "vla": { + "policy_instruction": ( + ctx.robot.vla_instruction or ctx.robot.task_description ), - "coordinate_frame": "right_base", - "camera_aliases": meta.get("observation_camera_map", {}), - "projection_views": meta.get("projection_views", {}), - "agent_observation": observation_policy, - "vla": { - "policy_instruction": self._vla_instruction, - "num_action_chunks": 20, - "action_dim": 20, - "num_images_in_input": 3, - "external_localization_views_in_policy_input": False, - "agent_visible_images": observation_policy["inline_cameras"], - "auxiliary_artifact_images": observation_policy["auxiliary_cameras"], - "skill_stop_rules": { - "enabled": True, - "grasp_lift_m": 0.15, - "place_lift_m": 0.10, - "handoff_release_delay_s": 1.5, - "min_steps_after_gripper_event": 2, - }, + "num_action_chunks": 20, + "action_dim": 20, + "num_images_in_input": 3, + "external_localization_views_in_policy_input": False, + "agent_visible_images": observation_policy["inline_cameras"], + "auxiliary_artifact_images": observation_policy["auxiliary_cameras"], + "skill_stop_rules": { + "enabled": True, + "grasp_lift_m": 0.15, + "place_lift_m": 0.10, + "handoff_release_delay_s": 1.5, + "min_steps_after_gripper_event": 2, }, - "sam3": { - "tool": "segment", - "status": ( - "optional; returns an error and falls back to manual camera " - "projection when no SAM3 client is configured" - ), - "usage": ( - "Use text prompt or one [row, col] positive camera point, then " - "inspect the returned mask overlay before trusting point_xyz. " - "SAM3 text grounding is phrase-sensitive; if a prompt returns a very low score, " - "retry a shorter/rephrased prompt or point prompt rather than " - "lowering min_score blindly." - ), - }, - "available_primitives": [spec["name"] for spec in TOOLS_SPEC], - "operator_guidance": ( - "Named VLA semantic boundaries are segment boundaries, not proof " - "of physical success. Verify images, gripper widths/open flags, " - "joint_health, and projection evidence after every action. " - "recover_joint_posture re-commands and preserves each gripper's " - "open/closed state; inspect its gripper_preserved result before " - "continuing." + }, + "sam3": { + "tool": "segment", + "status": ( + "optional; returns an error and falls back to manual camera " + "projection when no SAM3 client is configured" ), - } - - def _run_named_vla_skill( - self, - *, - skill_name: str, - boundary: str, - prompt: str, - max_chunks: int = 20, - ) -> dict[str, Any]: - if self.model is None: - raise RuntimeError(f"{skill_name} requires --vla-endpoint") - requested_prompt = str(prompt).strip() - if not requested_prompt: - raise ValueError("prompt must be non-empty") - # Task configuration owns policy conditioning; retain planner intent in logs. - effective_prompt = self._vla_instruction - if not 1 <= int(max_chunks) <= 20: - raise ValueError("max_chunks must be between 1 and 20") - - state = self.env.get_robot_state() - start_state = state - previous_left_open = bool(state["left_arm"]["gripper_open"]) - previous_right_open = bool(state["right_arm"]["gripper_open"]) - event_z: float | None = None - event_time: float | None = None - steps_after_event = 0 - chunks_executed = 0 - steps_executed = 0 - event_step: int | None = None - boundary_reached = False - terminated = False - truncated = False - action_count = 0 - action_min: np.ndarray | None = None - action_max: np.ndarray | None = None - action_sum: np.ndarray | None = None - first_action: list[float] | None = None - first_policy_state: list[float] | None = None - first_policy_state_shape: list[int] | None = None - started_at = time.perf_counter() - - for _ in range(int(max_chunks)): - self._check_cancelled() - observation = dict(self.env.get_observation()) - if first_policy_state is None: - obs_states = np.asarray(observation.get("states"), dtype=np.float32) - first_policy_state_shape = list(obs_states.shape) - first_policy_state = ( - np.round(obs_states.reshape(-1)[:20], 5).astype(float).tolist() - ) - observation["task_descriptions"] = effective_prompt - actions = np.asarray( - self.model.predict(observation, options={"mode": "eval"}), - dtype=np.float32, - ) - if actions.ndim != 2 or actions.shape[1] != 20: - raise RuntimeError( - f"{skill_name} expected [chunk, 20] actions, got {actions.shape}" - ) - if not np.isfinite(actions).all(): - raise RuntimeError(f"{skill_name} received non-finite VLA actions") - chunks_executed += 1 - if first_action is None: - first_action = np.round(actions[0], 5).astype(float).tolist() - if action_count == 0: - action_min = actions.min(axis=0) - action_max = actions.max(axis=0) - action_sum = actions.sum(axis=0) - else: - action_min = np.minimum(action_min, actions.min(axis=0)) - action_max = np.maximum(action_max, actions.max(axis=0)) - action_sum = action_sum + actions.sum(axis=0) - action_count += int(actions.shape[0]) - - for action in actions: - self._check_cancelled() - result = self.env.chunk_step(action[None, :]) - steps_executed += 1 - terminated = terminated or bool(result.get("terminated")) - truncated = truncated or bool(result.get("truncated")) - state = self.env.get_robot_state() - left = state["left_arm"] - right = state["right_arm"] - left_open = bool(left["gripper_open"]) - right_open = bool(right["gripper_open"]) - - if boundary == "grasp": - if event_z is None and previous_right_open and not right_open: - event_z = float(right["tcp_pose"][2]) - event_step = steps_executed - steps_after_event = 0 - elif event_z is not None: - steps_after_event += 1 - boundary_reached = ( - not right_open - and steps_after_event >= 2 - and float(right["tcp_pose"][2]) - event_z >= 0.15 - ) - elif boundary == "handoff": - if event_time is None and not previous_right_open and right_open: - event_time = time.monotonic() - event_step = steps_executed - steps_after_event = 0 - elif event_time is not None: - steps_after_event += 1 - boundary_reached = ( - right_open - and steps_after_event >= 2 - and time.monotonic() - event_time >= 1.5 - ) - elif boundary == "place": - if event_z is None and not previous_left_open and left_open: - event_z = float(left["tcp_pose"][2]) - event_step = steps_executed - steps_after_event = 0 - elif event_z is not None: - steps_after_event += 1 - boundary_reached = ( - left_open - and steps_after_event >= 2 - and float(left["tcp_pose"][2]) - event_z >= 0.10 - ) - else: # pragma: no cover - internal programming guard - raise ValueError(f"unknown VLA skill boundary: {boundary}") - - previous_left_open = left_open - previous_right_open = right_open - if boundary_reached or terminated or truncated: - break - if boundary_reached or terminated or truncated: - break - - stop_rule: dict[str, Any] = { - "phase": boundary, - "skill_name": skill_name, - "step_count": steps_executed, - "event_step": event_step, - "success_claim": False, - } - if boundary == "grasp": - stop_rule.update( - { - "condition": "right_gripper_closed_then_lifted", - "right_close_step": event_step, - "right_close_z": event_z, - "right_current_z": float(state["right_arm"]["tcp_pose"][2]), - "lift_m": ( - float(state["right_arm"]["tcp_pose"][2]) - event_z - if event_z is not None - else None - ), - "threshold_m": 0.15, - } - ) - elif boundary == "handoff": - stop_rule.update( - { - "condition": "right_gripper_opened_then_delay", - "right_open_step": event_step, - "elapsed_after_open_s": ( - time.monotonic() - event_time - if event_time is not None - else None - ), - "delay_s": 1.5, - "right_current_open": bool(state["right_arm"]["gripper_open"]), - } - ) - elif boundary == "place": - stop_rule.update( - { - "condition": "left_gripper_opened_then_lifted", - "left_open_step": event_step, - "left_open_z": event_z, - "left_current_z": float(state["left_arm"]["tcp_pose"][2]), - "lift_m": ( - float(state["left_arm"]["tcp_pose"][2]) - event_z - if event_z is not None - else None - ), - "threshold_m": 0.10, - } - ) + "usage": ( + "Use text prompt or one [row, col] positive camera point, then " + "inspect the returned mask overlay before trusting point_xyz. " + "SAM3 text grounding is phrase-sensitive; if a prompt returns a very low score, " + "retry a shorter/rephrased prompt or point prompt rather than " + "lowering min_score blindly." + ), + }, + "available_primitives": [ + item.name for item in DUAL_FRANKA_TOOLS if item.name != "finish" + ], + "operator_guidance": ( + "Named VLA semantic boundaries are segment boundaries, not proof " + "of physical success. Verify images, gripper widths/open flags, " + "joint_health, and projection evidence after every action. " + "recover_joint_posture re-commands and preserves each gripper's " + "open/closed state; inspect its gripper_preserved result before " + "continuing." + ), + } + if ctx.robot.session is not None and ctx.robot.session._mode == "exploration": + data["phase"] = "exploration" + data["reset_policy"] = ( + "No automatic reset was performed by the client. Before motion in " + "each session call request_scene_reset and wait for operator confirmation." + ) + data["scene_ready"] = ctx.robot.session._scene_ready + return _result(data) - action_summary: dict[str, Any] = { - "count": action_count, - "shape": [action_count, 20], - "finite": True, - } - if action_count and action_min is not None and action_max is not None: - action_summary.update( - { - "min": np.round(action_min, 5).astype(float).tolist(), - "max": np.round(action_max, 5).astype(float).tolist(), - "mean": np.round(action_sum / action_count, 5) - .astype(float) - .tolist(), - "first_action": first_action, - "first_policy_state_shape": first_policy_state_shape, - "first_policy_state": first_policy_state, - } - ) - return { - "ok": boundary_reached and not (terminated or truncated), - "skill_name": skill_name, - "boundary": boundary, - "requested_prompt": requested_prompt, - "effective_policy_prompt": effective_prompt, - "prompt_overridden": requested_prompt != effective_prompt, - "boundary_reached": boundary_reached, - "chunks_executed": chunks_executed, - "steps_executed": steps_executed, - "terminated": terminated, - "truncated": truncated, - "stop_rule": stop_rule, - "action_summary": action_summary, - "elapsed_s": time.perf_counter() - started_at, - "vla_start_robot_state": start_state, - "robot_state": state, - } - - def vla_right_grasp(self, prompt: str, max_chunks: int = 20) -> dict[str, Any]: - return self._run_named_vla_skill( +@tool +def vla_right_grasp( + prompt: str, + max_chunks: Annotated[ + int, Field(ge=1, le=20, json_schema_extra={"default": 20}) + ] = 20, + *, + ctx: ToolContext[DualFrankaRuntime], +) -> ToolResult: + """Run the learned right-grasp VLA segment. The active task prompt decides which object is currently allowed; this tool only defines the capability boundary: right gripper closes and the right TCP lifts. + + Args: + prompt: Planner-facing segment intent. This is recorded in the tool result; the current live clean-desk checkpoint still receives its fixed training instruction during policy inference. + """ + return _result( + _run_named_vla_skill( + ctx, skill_name="vla_right_grasp", boundary="grasp", prompt=prompt, max_chunks=max_chunks, ) + ) + - def vla_handoff(self, prompt: str, max_chunks: int = 20) -> dict[str, Any]: - return self._run_named_vla_skill( +@tool +def vla_handoff( + prompt: str, + max_chunks: Annotated[ + int, Field(ge=1, le=20, json_schema_extra={"default": 20}) + ] = 20, + *, + ctx: ToolContext[DualFrankaRuntime], +) -> ToolResult: + """Run the learned bimanual handoff VLA segment. The capability boundary is right-gripper release followed by the configured settle delay; do not rule-base pre-position either arm for it. + + Args: + prompt: Planner-facing segment intent. This is recorded in the tool result; the current live clean-desk checkpoint still receives its fixed training instruction during policy inference. + """ + return _result( + _run_named_vla_skill( + ctx, skill_name="vla_handoff", boundary="handoff", prompt=prompt, max_chunks=max_chunks, ) + ) - def vla_left_place(self, prompt: str, max_chunks: int = 20) -> dict[str, Any]: - return self._run_named_vla_skill( + +@tool +def vla_left_place( + prompt: str, + max_chunks: Annotated[ + int, Field(ge=1, le=20, json_schema_extra={"default": 20}) + ] = 20, + *, + ctx: ToolContext[DualFrankaRuntime], +) -> ToolResult: + """Run the learned left-placement VLA segment. The active task decides the destination; this tool only defines the capability boundary: left gripper opens and the left TCP lifts. + + Args: + prompt: Planner-facing segment intent. This is recorded in the tool result; the current live clean-desk checkpoint still receives its fixed training instruction during policy inference. + """ + return _result( + _run_named_vla_skill( + ctx, skill_name="vla_left_place", boundary="place", prompt=prompt, max_chunks=max_chunks, ) + ) + + +@tool +def recover_joint_posture( + reason: Annotated[str, Field(json_schema_extra={"default": ""})] = "", + return_to_start: Annotated[bool, Field(json_schema_extra={"default": True})] = True, + *, + ctx: ToolContext[DualFrankaRuntime], +) -> ToolResult: + """Reset both arms to their healthy configured joint posture while preserving each gripper's open/closed state. Closed grippers are re-commanded before/after the joint reset so held objects stay clamped, then both TCPs return near their prior poses.""" + ctx.check_cancelled() + return _result( + ctx.robot.env.recover_joint_posture( + reason=reason, return_to_start=return_to_start + ) + ) + + +def _perception_result(data: dict[str, Any]) -> ToolResult: + data = dict(data) + image = data.pop("_image_cam_bytes", None) + result = _result(data) + if image is not None: + result.images.append(image) + return result + + +@tool +@readonly +def back_project( + row: Annotated[int, Field(ge=0)], + col: Annotated[int, Field(ge=0)], + camera: Annotated[str, Field(json_schema_extra={"default": "d455"})] = "d455", + target_name: Annotated[ + str, Field(json_schema_extra={"default": "target"}) + ] = "target", + step: int | None = None, + window_radius: Annotated[int, Field(ge=0, json_schema_extra={"default": 2})] = 2, + *, + ctx: ToolContext[DualFrankaRuntime], +) -> ToolResult: + """Back-project one pixel from a registered RGBD camera view into shared right-base coordinates. Use a camera listed by view_env_state/view_camera_meta; default is the configured primary metric localization camera. + + Args: + camera: Registered projection view name, e.g. d455 or base. Valid names come from perception.projection_views and the current state's saved artifacts. + """ + return _perception_result( + perception.back_project( + row=row, + col=col, + camera=camera, + target_name=target_name, + step=step, + window_radius=window_radius, + state=ctx.state, + ) + ) + + +@tool +@readonly +def segment( + camera: Annotated[str, Field(json_schema_extra={"default": "d455"})] = "d455", + prompt: Annotated[str, Field(json_schema_extra={"default": ""})] = "", + point: Annotated[list[int], Field(min_length=2, max_length=2)] | None = None, + target_name: Annotated[ + str, Field(json_schema_extra={"default": "target"}) + ] = "target", + step: int | None = None, + min_score: Annotated[ + float, Field(ge=0, le=1, json_schema_extra={"default": 0.2}) + ] = 0.2, + min_valid_depth_pixels: Annotated[ + int, Field(ge=1, json_schema_extra={"default": 25}) + ] = 25, + *, + ctx: ToolContext[DualFrankaRuntime], +) -> ToolResult: + """Use SAM3 on a registered RGB image with either a text prompt or one positive [row, col] point, return a mask overlay for verification, and estimate the mask median point in shared right-base coordinates. + + Args: + camera: Registered projection view name, e.g. d455 or base. Valid names come from perception.projection_views and the current state's saved artifacts. + prompt: Text prompt for SAM3. Prefer short object/relation phrases; for the clean-desk box use 'white interior of the black cardboard box' or 'cardboard box'. Avoid over-specific surface words such as 'floor' when grounding is weak. Provide exactly one of prompt or point. + point: Positive SAM3 point in camera image coordinates [row, col]. Provide exactly one of prompt or point. + """ + return _perception_result( + perception.segment( + camera=camera, + prompt=prompt, + point=point, + target_name=target_name, + step=step, + min_score=min_score, + min_valid_depth_pixels=min_valid_depth_pixels, + state=ctx.state, + sam3_client=ctx.robot.sam3_client, + ) + ) + + +@tool +def request_scene_reset( + reason: str, + expected_scene_state: Annotated[str, Field(json_schema_extra={"default": ""})] = "", + *, + ctx: ToolContext[DualFrankaRuntime], +) -> ToolResult: + """Exploration-only real-robot reset gate. Ask the human operator to remove/secure held objects and restore the tabletop scene for another attempt, wait for terminal confirmation, then reset the robot posture. This does not automatically restore physical objects like a simulator. + + Args: + reason: Why the scene needs to be restored. + expected_scene_state: Short instruction for the operator describing the desired restored layout. + """ + ctx.check_cancelled() + return _result(ctx.robot.session._request_scene_reset(reason, expected_scene_state)) + + +@tool +@readonly +def request_operator_verdict( + question: Annotated[ + str, + Field( + json_schema_extra={ + "default": "Does the current real-robot scene satisfy the task?" + } + ), + ] = "Does the current real-robot scene satisfy the task?", + *, + ctx: ToolContext[DualFrankaRuntime], +) -> ToolResult: + """Exploration-only human feedback gate. Ask the operator to mark the current physical task state as success, failure, or continue before the planner finishes or starts another attempt.""" + session = ctx.robot.session + handler = ( + session._request_operator_verdict + if session._mode == "exploration" + else session._evaluation_verdict + ) + return _result(handler(question)) + + +@tool +@readonly +def finish( + status: str, summary: str, *, ctx: ToolContext[DualFrankaRuntime] +) -> ToolResult: + """Call when the task is complete or unrecoverable. Halts the agent loop. Save any artifacts (recipe, audit) BEFORE calling finish. + + Args: + status: Outcome, e.g. 'success', 'failure', or 'stuck'. + summary: Short natural-language summary of the run. + """ + session = ctx.robot.session + handler = ( + session._guarded_finish + if session._mode == "exploration" + else session._evaluation_finish + ) + return _result( + handler( + lambda **kwargs: {"_finish": True, **kwargs}, status=status, summary=summary + ) + ) def dump_state( - primitives: DualFrankaPrimitives, + runtime: DualFrankaRuntime, state: EnvState, *, command: dict[str, Any] | None, @@ -721,9 +665,9 @@ def dump_state( elapsed_s: float | None, ) -> StepRecord: """Capture per-arm robot state and synchronized camera images.""" - observation = primitives.env.get_observation() - robot_state = primitives.env.get_robot_state() - metadata = primitives.env.get_camera_meta() + observation = runtime.env.get_observation() + robot_state = runtime.env.get_robot_state() + metadata = runtime.env.get_camera_meta() with state.record_step( state=robot_state, command=command, @@ -801,10 +745,7 @@ def dump_state( return state.get(step) -@readonly -def view_env_state(step: int = -1, *, state: EnvState) -> dict[str, Any]: - """Return one recorded dual-Franka state with configured inline camera views.""" - record = state.get(step) +def build_observation(state: EnvState, record: StepRecord) -> ToolResult: output = record.to_blob() output["images"] = [] output["artifact_images"] = [] @@ -829,13 +770,43 @@ def view_env_state(step: int = -1, *, state: EnvState) -> dict[str, Any]: state.artifact_path(artifact, step=record.step_idx) ) output["artifact_images"].append(view) + images = [] for index, view in enumerate(policy["inline_cameras"]): artifact = f"{view}.png" if index >= 4 or not state.exists(artifact, step=record.step_idx): continue - output[_image_bytes_key(index)] = state.load_bytes( - artifact, step=record.step_idx - ) + images.append(state.load_bytes(artifact, step=record.step_idx)) output["images"].append(view) output["image_block_order"] = list(output["images"]) - return output + return ToolResult(data=output, images=images) + + +@tool +@readonly +def view_env_state( + step: Annotated[int, Field(json_schema_extra={"default": -1})] = -1, + *, + ctx: ToolContext[DualFrankaRuntime], +) -> ToolResult: + """Read a dual-Franka state snapshot. The D455 image is returned inline; left_wrist, base, and right_wrist are returned as artifact paths for targeted read_image inspection.""" + return build_observation(ctx.state, ctx.state.get(step)) + + +DUAL_FRANKA_TOOLS = ( + finish, + describe_dual_franka_setup, + view_env_state, + view_camera_meta, + back_project, + segment, + move_delta, + rotate_delta, + open_gripper, + close_gripper, + recover_joint_posture, + request_scene_reset, + request_operator_verdict, + vla_right_grasp, + vla_handoff, + vla_left_place, +) diff --git a/robots/franka/perception.py b/robots/franka/perception.py index 89aa2de64..9ac7b7f64 100644 --- a/robots/franka/perception.py +++ b/robots/franka/perception.py @@ -25,7 +25,6 @@ from robots.franka.runtime_config import get_calibration_path from rpent.session import EnvState -from rpent.tools.toolkit import readonly class PerceptionError(ValueError): @@ -77,7 +76,6 @@ def load_calibration_bundle( } -@readonly def view_perception_setup( *, state: EnvState | None = None, @@ -103,7 +101,6 @@ def view_perception_setup( } -@readonly def back_project( *, row: int, @@ -249,7 +246,6 @@ def _save_selected_pixel_overlay( return None -@readonly def back_project_correspondence( *, third_person_row: int | None = None, diff --git a/robots/franka/toolkit.py b/robots/franka/toolkit.py index ae48baf16..7a7211301 100644 --- a/robots/franka/toolkit.py +++ b/robots/franka/toolkit.py @@ -12,31 +12,41 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Franka toolkit integrated with RPent's centralized environment state.""" +"""Native Franka toolkit and per-session robot resources.""" from __future__ import annotations -from functools import partial +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any -from robots.franka import perception as franka_perception from robots.franka import tools as franka_tools from robots.franka.runtime_config import set_calibration_path -from rpent.dashboard.events import DashboardEventSink -from rpent.session import EnvState -from rpent.tools.toolkit import Toolkit +from rpent.dashboard.events import DashboardEventSink, StepRecordEvent +from rpent.session import EnvState, StepRecord +from rpent.tools import Toolkit, ToolResult from rpent.utils.logging import get_output_dir if TYPE_CHECKING: + from robots.franka.env_client import FrankaEnvClient from rpent.memory.manager import MemoryManager + from rpent.robots.components.pi05_vla_client import Pi05VLAClient -class FrankaToolkit(Toolkit): - """Common RPent tools plus safe single-Franka planner primitives.""" +@dataclass +class FrankaRuntime: + """Environment and optional VLA client shared by one Franka session.""" - _tools_module = franka_tools - _primitives_cls = franka_tools.FrankaPrimitives + env: FrankaEnvClient + model: Pi05VLAClient | None + task_description: str + + +class FrankaToolkit(Toolkit[FrankaRuntime]): + """Common native tools plus single-Franka motion and perception.""" + + _runtime_type = FrankaRuntime + _robot_tools = franka_tools.FRANKA_TOOLS def __init__( self, @@ -46,69 +56,59 @@ def __init__( memory: MemoryManager, state_output_dir: Path | str | None = None, ) -> None: - state = EnvState(Path(state_output_dir or get_output_dir())) - super().__init__( - dashboard_events=dashboard_events, - state=state, - memory=memory, - ) + runtime_kwargs = dict(runtime_kwargs) calibration_path = runtime_kwargs.pop("calibration_path", None) if calibration_path is not None: set_calibration_path(calibration_path) - self._primitives = self._primitives_cls( - check_cancelled=self.raise_if_cancelled, - **runtime_kwargs, + output_dir = Path(state_output_dir or get_output_dir()) + state = EnvState(output_dir) + super().__init__( + state=state, + memory=memory, + robot=self._runtime_type(**runtime_kwargs), + output_dir=output_dir, + tools=self._robot_tools, + dashboard_events=dashboard_events, ) - self._register_tools() - self._state.reset() - record = self._tools_module.dump_state( - self._primitives, - self._state, + # The env client already resets the robot when it connects. + state.reset() + record = self._dump_state( command=None, result=None, elapsed_s=None, ) - self._publish_step(record) - - def _register_tools(self) -> None: - state_handlers = { - "view_env_state": partial(franka_tools.view_env_state, state=self._state), - "view_camera_meta": partial( - franka_tools.view_camera_meta, - state=self._state, - ), - "view_perception_setup": partial( - franka_perception.view_perception_setup, - state=self._state, - ), - "back_project": partial( - franka_perception.back_project, - state=self._state, - ), - "back_project_correspondence": partial( - franka_perception.back_project_correspondence, - state=self._state, - ), - } - for spec in self._tools_module.TOOLS_SPEC: - name = spec["name"] - handler = state_handlers.get(name) or getattr(self._primitives, name) - self.add_tool(name, spec, handler) + self._dashboard_events.emit(StepRecordEvent(record=record, env_state=state)) - def get_env_state( + def _capture_observation( self, *, command: dict[str, Any], - result: dict[str, Any], + result: ToolResult, elapsed_s: float, - ) -> dict[str, Any]: - record = self._tools_module.dump_state( - self._primitives, - self._state, + ) -> tuple[dict[str, Any], list[bytes]]: + record = self._dump_state( + command=command, + result=result.to_dict(), + elapsed_s=elapsed_s, + ) + observation = self._build_observation(record) + observation.data["agent_elapsed_s"] = elapsed_s + return observation.data, observation.images + + def _dump_state( + self, + *, + command: dict[str, Any] | None, + result: dict[str, Any] | None, + elapsed_s: float | None, + ) -> StepRecord: + return franka_tools.dump_state( + self._robot, + self.state, command=command, result=result, elapsed_s=elapsed_s, ) - output = self._tools_module.view_env_state(record.step_idx, state=self._state) - output["agent_elapsed_s"] = elapsed_s - return output + + def _build_observation(self, record: StepRecord) -> ToolResult: + return franka_tools.build_observation(self.state, record) diff --git a/robots/franka/tools.py b/robots/franka/tools.py index 3b2b9ac00..8bf8dc272 100644 --- a/robots/franka/tools.py +++ b/robots/franka/tools.py @@ -12,210 +12,102 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Franka planner tools, primitives, and canonical state capture.""" +"""Native Franka tools and canonical RGB-D state capture.""" from __future__ import annotations -from collections.abc import Callable, Sequence -from typing import Any +from typing import TYPE_CHECKING, Annotated, Any, Literal import numpy as np +from pydantic import Field, FiniteFloat +from robots.franka import perception from rpent.session import EnvState, StepRecord -from rpent.tools.toolkit import readonly - -TOOLS_SPEC = [ - { - "name": "view_env_state", - "description": "Read a Franka state snapshot and its synchronized RGB images.", - "input_schema": { - "type": "object", - "properties": {"step": {"type": "integer", "default": -1}}, - }, - }, - { - "name": "view_camera_meta", - "description": "Read camera intrinsics, crop, depth, and calibration metadata.", - "input_schema": { - "type": "object", - "properties": {"step": {"type": "integer", "default": -1}}, - }, - }, - { - "name": "view_perception_setup", - "description": "Read calibrated camera geometry and projection conventions.", - "input_schema": { - "type": "object", - "properties": {"step": {"type": "integer", "default": -1}}, - }, - }, - { - "name": "back_project", - "description": "Back-project one wrist or external-camera pixel into Franka base coordinates.", - "input_schema": { - "type": "object", - "properties": { - "row": {"type": "integer", "minimum": 0}, - "col": {"type": "integer", "minimum": 0}, - "step": {"type": "integer"}, - "camera": {"type": "string", "enum": ["wrist", "third_person"]}, - "debug": {"type": "boolean", "default": False}, - }, - "required": ["row", "col"], - }, - }, - { - "name": "back_project_correspondence", - "description": "Fuse matched wrist and external-camera pixels into a Franka base point.", - "input_schema": { - "type": "object", - "properties": { - "third_person_row": {"type": "integer", "minimum": 0}, - "third_person_col": {"type": "integer", "minimum": 0}, - "wrist_row": {"type": "integer", "minimum": 0}, - "wrist_col": {"type": "integer", "minimum": 0}, - "pixels": {"type": "array", "items": {"type": "object"}}, - "step": {"type": "integer"}, - "debug": {"type": "boolean", "default": False}, - }, - }, - }, - { - "name": "move_delta", - "description": "Move the Franka TCP by a bounded base-frame xyz delta in meters.", - "input_schema": { - "type": "object", - "properties": { - "delta_xyz": { - "type": "array", - "items": {"type": "number"}, - "minItems": 3, - "maxItems": 3, - } - }, - "required": ["delta_xyz"], - }, - }, - { - "name": "rotate_delta", - "description": "Rotate the Franka TCP by a bounded base-frame rpy delta in radians.", - "input_schema": { - "type": "object", - "properties": { - "delta_rpy": { - "type": "array", - "items": {"type": "number"}, - "minItems": 3, - "maxItems": 3, - } - }, - "required": ["delta_rpy"], - }, - }, - { - "name": "open_gripper", - "description": "Open the Franka gripper and wait for the command to settle.", - "input_schema": {"type": "object", "properties": {}}, - }, - { - "name": "close_gripper", - "description": "Close the Franka gripper and wait for the command to settle.", - "input_schema": {"type": "object", "properties": {}}, - }, - { - "name": "vla_grasp", - "description": "Run bounded real-world VLA action chunks for a local grasp attempt.", - "input_schema": { - "type": "object", - "properties": { - "prompt": {"type": "string"}, - "max_chunks": {"type": "integer", "minimum": 1, "maximum": 20}, - }, - "required": ["prompt"], - }, - }, -] - - -def coerce_vec3(value: Sequence[float], *, name: str) -> np.ndarray: - """Return a finite float32 three-vector or raise a useful error.""" - array = np.asarray(value, dtype=np.float32) - if array.shape != (3,): - raise ValueError(f"{name} must contain exactly 3 values, got {array.shape}") - if not np.isfinite(array).all(): - raise ValueError(f"{name} must contain only finite values") - return array - - -class FrankaPrimitives: - """Safe agent-facing operations over a remote Franka environment.""" - - def __init__( - self, - *, - env: Any, - model: Any | None, - task_description: str, - check_cancelled: Callable[[], None], - ) -> None: - self.env = env - self.model = model - self.task_description = task_description - self._check_cancelled = check_cancelled - - def reset(self) -> dict[str, Any]: - return self.env.reset() - - def move_delta(self, delta_xyz: Sequence[float]) -> dict[str, Any]: - self._check_cancelled() - return self.env.move_delta(coerce_vec3(delta_xyz, name="delta_xyz")) - - def rotate_delta(self, delta_rpy: Sequence[float]) -> dict[str, Any]: - self._check_cancelled() - return self.env.rotate_delta(coerce_vec3(delta_rpy, name="delta_rpy")) - - def open_gripper(self) -> dict[str, Any]: - self._check_cancelled() - return self.env.set_gripper(open=True) - - def close_gripper(self) -> dict[str, Any]: - self._check_cancelled() - return self.env.set_gripper(open=False) - - def vla_grasp(self, prompt: str, max_chunks: int = 4) -> dict[str, Any]: - if self.model is None: - raise RuntimeError("vla_grasp requires --vla-endpoint") - if not prompt.strip(): - raise ValueError("prompt must be non-empty") - if not 1 <= int(max_chunks) <= 20: - raise ValueError("max_chunks must be between 1 and 20") - - chunk_results: list[dict[str, Any]] = [] - observation: dict[str, Any] | None = None - for _ in range(int(max_chunks)): - self._check_cancelled() - if observation is None: - observation = dict(self.env.get_observation()) - observation["task_descriptions"] = prompt or self.task_description - actions = self.model.predict(observation, options={"mode": "eval"}) - result = self.env.chunk_step(actions) - chunk_results.append(result) - if result.get("terminated") or result.get("truncated"): - break - # Reuse the obs chunk_step already returned instead of re-fetching it. - next_obs = result.get("observation") - observation = dict(next_obs) if isinstance(next_obs, dict) else None - - return { +from rpent.tools import ToolContext, ToolResult, readonly, tool + +if TYPE_CHECKING: + from robots.franka.toolkit import FrankaRuntime + + +def _result(data: dict[str, Any]) -> ToolResult: + data = dict(data) + error = data.pop("error", None) + return ToolResult(data=data, error=error) + + +@tool +def move_delta( + delta_xyz: Annotated[list[FiniteFloat], Field(min_length=3, max_length=3)], + *, + ctx: ToolContext[FrankaRuntime], +) -> ToolResult: + """Move the Franka TCP by a bounded base-frame xyz delta in meters.""" + ctx.check_cancelled() + return _result(ctx.robot.env.move_delta(np.asarray(delta_xyz, dtype=np.float32))) + + +@tool +def rotate_delta( + delta_rpy: Annotated[list[FiniteFloat], Field(min_length=3, max_length=3)], + *, + ctx: ToolContext[FrankaRuntime], +) -> ToolResult: + """Rotate the Franka TCP by a bounded base-frame rpy delta in radians.""" + ctx.check_cancelled() + return _result(ctx.robot.env.rotate_delta(np.asarray(delta_rpy, dtype=np.float32))) + + +@tool +def open_gripper(*, ctx: ToolContext[FrankaRuntime]) -> ToolResult: + """Open the Franka gripper and wait for the command to settle.""" + ctx.check_cancelled() + return _result(ctx.robot.env.set_gripper(open=True)) + + +@tool +def close_gripper(*, ctx: ToolContext[FrankaRuntime]) -> ToolResult: + """Close the Franka gripper and wait for the command to settle.""" + ctx.check_cancelled() + return _result(ctx.robot.env.set_gripper(open=False)) + + +@tool +def vla_grasp( + prompt: str, + max_chunks: Annotated[int, Field(ge=1, le=20)] = 4, + *, + ctx: ToolContext[FrankaRuntime], +) -> ToolResult: + """Run bounded real-world VLA action chunks for a local grasp attempt.""" + runtime = ctx.robot + if runtime.model is None: + raise RuntimeError("vla_grasp requires --vla-endpoint") + if not prompt.strip(): + raise ValueError("prompt must be non-empty") + observation = None + for chunk in range(int(max_chunks)): + ctx.check_cancelled() + if observation is None: + observation = dict(runtime.env.get_observation()) + observation["task_descriptions"] = prompt + actions = runtime.model.predict(observation, options={"mode": "eval"}) + result = runtime.env.chunk_step(actions) + if result.get("terminated") or result.get("truncated"): + break + next_obs = result.get("observation") + observation = dict(next_obs) if isinstance(next_obs, dict) else None + return _result( + { "ok": True, - "chunks_executed": len(chunk_results), - "last_chunk": chunk_results[-1] if chunk_results else None, - "robot_state": self.env.get_robot_state(), + "chunks_executed": chunk + 1, + "last_chunk": result, + "robot_state": runtime.env.get_robot_state(), } + ) def dump_state( - primitives: FrankaPrimitives, + runtime: FrankaRuntime, state: EnvState, *, command: dict[str, Any] | None, @@ -223,9 +115,9 @@ def dump_state( elapsed_s: float | None, ) -> StepRecord: """Capture robot state and synchronized camera artifacts in ``EnvState``.""" - observation = primitives.env.get_observation() - robot_state = primitives.env.get_robot_state() - metadata = primitives.env.get_camera_meta() + observation = runtime.env.get_observation() + robot_state = runtime.env.get_robot_state() + metadata = runtime.env.get_camera_meta() with state.record_step( state=robot_state, command=command, @@ -255,34 +147,131 @@ def dump_state( return state.get(step) +def build_observation(state: EnvState, record: StepRecord) -> ToolResult: + """Return recorded JSON and wrist/external PNGs in their declared order.""" + data = record.to_blob() + images = [] + for name, field in (("camera", "image_cam_path"), ("wrist", "image_wrist_path")): + if state.exists(f"{name}.png", step=record.step_idx): + data[field] = str(state.artifact_path(f"{name}.png", step=record.step_idx)) + images.append(state.load_bytes(f"{name}.png", step=record.step_idx)) + return ToolResult(data=data, images=images) + + +@tool @readonly -def view_env_state(step: int = -1, *, state: EnvState) -> dict[str, Any]: - """Return one recorded Franka state with image blocks for the planner.""" - record = state.get(step) - output = record.to_blob() - if state.exists("wrist.png", step=record.step_idx): - output["image_wrist_path"] = str( - state.artifact_path("wrist.png", step=record.step_idx) - ) - output["_image_wrist_bytes"] = state.load_bytes( - "wrist.png", step=record.step_idx - ) - if state.exists("camera.png", step=record.step_idx): - output["image_cam_path"] = str( - state.artifact_path("camera.png", step=record.step_idx) +def view_env_state( + step: Annotated[int, Field(json_schema_extra={"default": -1})] = -1, + *, + ctx: ToolContext[FrankaRuntime], +) -> ToolResult: + """Read a Franka state snapshot and its synchronized RGB images.""" + return build_observation(ctx.state, ctx.state.get(step)) + + +@tool +@readonly +def view_camera_meta( + step: Annotated[int, Field(json_schema_extra={"default": -1})] = -1, + *, + ctx: ToolContext[FrankaRuntime], +) -> ToolResult: + """Read camera intrinsics, crop, depth, and calibration metadata.""" + if not ctx.state.exists("camera_meta.json", step=step): + return ToolResult(data={"step": step}, error="camera metadata is unavailable") + return _result( + { + "step": ctx.state.get(step).step_idx, + "camera_meta": ctx.state.load("camera_meta.json", step=step), + } + ) + + +@tool +@readonly +def view_perception_setup( + step: Annotated[int, Field(json_schema_extra={"default": -1})] = -1, + *, + ctx: ToolContext[FrankaRuntime], +) -> ToolResult: + """Read calibrated camera geometry and projection conventions.""" + return _result(perception.view_perception_setup(state=ctx.state, step=step)) + + +@tool +@readonly +def back_project( + row: Annotated[int, Field(ge=0)], + col: Annotated[int, Field(ge=0)], + step: int | None = None, + camera: Literal["wrist", "third_person"] = "wrist", + debug: Annotated[bool, Field(json_schema_extra={"default": False})] = False, + *, + ctx: ToolContext[FrankaRuntime], +) -> ToolResult: + """Back-project one wrist or external-camera pixel into Franka base coordinates.""" + return _result( + perception.back_project( + row=row, + col=col, + step=step, + camera=camera, + debug=debug, + state=ctx.state, ) - output["_image_cam_bytes"] = state.load_bytes( - "camera.png", step=record.step_idx + ) + + +@tool +@readonly +def back_project_correspondence( + third_person_row: Annotated[int, Field(ge=0)] | None = None, + third_person_col: Annotated[int, Field(ge=0)] | None = None, + wrist_row: Annotated[int, Field(ge=0)] | None = None, + wrist_col: Annotated[int, Field(ge=0)] | None = None, + pixels: list[dict[str, Any]] | None = None, + step: int | None = None, + debug: Annotated[bool, Field(json_schema_extra={"default": False})] = False, + *, + ctx: ToolContext[FrankaRuntime], +) -> ToolResult: + """Fuse matched wrist and external-camera pixels into a Franka base point.""" + return _result( + perception.back_project_correspondence( + third_person_row=third_person_row, + third_person_col=third_person_col, + wrist_row=wrist_row, + wrist_col=wrist_col, + pixels=pixels, + step=step, + debug=debug, + state=ctx.state, ) - return output + ) +@tool @readonly -def view_camera_meta(step: int = -1, *, state: EnvState) -> dict[str, Any]: - """Return camera metadata captured for one state step.""" - if not state.exists("camera_meta.json", step=step): - return {"error": "camera metadata is unavailable", "step": step} - return { - "step": state.get(step).step_idx, - "camera_meta": state.load("camera_meta.json", step=step), - } +def finish(status: str, summary: str, *, ctx: ToolContext) -> ToolResult: + """Call when the task is complete or unrecoverable. Halts the agent loop. Save any artifacts (recipe, audit) BEFORE calling finish. + + Args: + status: Outcome, e.g. 'success', 'failure', or 'stuck'. + summary: Short natural-language summary of the run. + """ + return _result({"_finish": True, "status": status, "summary": summary}) + + +FRANKA_TOOLS = ( + finish, + view_env_state, + view_camera_meta, + view_perception_setup, + back_project, + back_project_correspondence, + move_delta, + rotate_delta, + open_gripper, + close_gripper, + vla_grasp, +) diff --git a/robots/libero/flash/replay.py b/robots/libero/flash/replay.py index 4efe1d99a..b957b99b8 100644 --- a/robots/libero/flash/replay.py +++ b/robots/libero/flash/replay.py @@ -32,15 +32,18 @@ import re from collections.abc import Callable from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np -from robots.libero import tools as libero_tools from robots.libero.flash.prompts import build as prompt_for from rpent.robots.components.molmo_client import MolmoClient from rpent.session import EnvState +if TYPE_CHECKING: + from robots.libero.toolkit import LiberoToolkit + + #: ``__t_s``, the tag the CLI builds per cell. _CELL = re.compile(r"^(10|goal|object|spatial)_(task|swap)_t(\d+)_s(\d+)$") @@ -71,7 +74,7 @@ def profile( - state: EnvState, + toolkit: LiberoToolkit, step: int, camera: str, col: float, @@ -88,15 +91,17 @@ def profile( """ points = [] for offset in np.linspace(-spread, spread, samples): - found = libero_tools.back_project( - row=int(round(row + offset)), - col=int(round(col)), - step=step, - camera=camera, - resolution="high", - state=state, + found = toolkit.execute_tool( + "back_project", + { + "row": int(round(row + offset)), + "col": int(round(col)), + "step": step, + "camera": camera, + "resolution": "high", + }, ) - world = found.get("world_xyz") if isinstance(found, dict) else None + world = found.data.get("world_xyz") if not found.is_error else None if isinstance(world, list) and len(world) >= 3 and all(np.isfinite(world[:3])): points.append([float(v) for v in world[:3]]) if not points: @@ -119,20 +124,22 @@ def profile( } -def locate(molmo: MolmoClient, state: EnvState, step: int, camera: str, query: str): - image = _image_bytes(state, step, camera) +def locate( + molmo: MolmoClient, toolkit: LiberoToolkit, step: int, camera: str, query: str +): + image = _image_bytes(toolkit.state, step, camera) if image is None: return None found = molmo.ground(image, query) if not found.found: return None col, row = found.point_xy - return profile(state, step, camera, col, row) + return profile(toolkit, step, camera, col, row) -def held_body(molmo: MolmoClient, state: EnvState, step: int, query: str): +def held_body(molmo: MolmoClient, toolkit: LiberoToolkit, step: int, query: str): """What is in the gripper, sampled on a grid because a line may miss it.""" - image = _image_bytes(state, step, "wrist") + image = _image_bytes(toolkit.state, step, "wrist") if image is None: return None found = molmo.ground(image, query) @@ -142,15 +149,17 @@ def held_body(molmo: MolmoClient, state: EnvState, step: int, query: str): points = [] for dc in (-40, 0, 40): for dr in (-40, 0, 40): - got = libero_tools.back_project( - row=int(round(row + dr)), - col=int(round(col + dc)), - step=step, - camera="wrist", - resolution="high", - state=state, + got = toolkit.execute_tool( + "back_project", + { + "row": int(round(row + dr)), + "col": int(round(col + dc)), + "step": step, + "camera": "wrist", + "resolution": "high", + }, ) - world = got.get("world_xyz") if isinstance(got, dict) else None + world = got.data.get("world_xyz") if not got.is_error else None if ( isinstance(world, list) and len(world) >= 3 @@ -198,14 +207,14 @@ def pick_succeeded(raw) -> bool: return action_result(raw).get("success") is True -def execute(toolkit: Any, name: str, arguments: dict[str, Any]) -> dict[str, Any]: +def execute( + toolkit: LiberoToolkit, name: str, arguments: dict[str, Any] +) -> dict[str, Any]: """Execute one toolkit action and turn its error result into an exception.""" - result = toolkit.execute_tool(name, arguments).result - if not isinstance(result, dict): - raise RuntimeError(f"{name} returned an invalid result: {result!r}") - if error := result.get("error"): - raise RuntimeError(f"{name} failed: {error}") - return result + result = toolkit.execute_tool(name, arguments) + if result.is_error: + raise RuntimeError(f"{name} failed: {result.error}") + return result.data def plans(root: Path) -> Path: @@ -231,7 +240,7 @@ def load(root: Path, plan_name: str) -> dict: def replay( - toolkit: Any, + toolkit: LiberoToolkit, molmo: MolmoClient, program: dict, note: Callable[[str], None] = lambda _: None, @@ -294,7 +303,7 @@ def finished() -> bool: ) else: got = locate( - molmo, state, opening, "agentview", prompt_for("survey", phrase) + molmo, toolkit, opening, "agentview", prompt_for("survey", phrase) ) xy = got["xy"] if got else None if xy is None or max(abs(xy[0]), abs(xy[1])) > REACH: @@ -340,7 +349,7 @@ def finished() -> bool: ) except Exception: continue - close = locate(molmo, state, look(), "wrist", prompt_for("refine", phrase)) + close = locate(molmo, toolkit, look(), "wrist", prompt_for("refine", phrase)) if close is None: continue gap = float(np.linalg.norm(close["xy"] - coarse)) @@ -424,7 +433,7 @@ def finished() -> bool: elif name == "set_gripper": execute(toolkit, name, arguments) held_step = look() - body = held_body(molmo, state, held_step, prompt_for("held", held_phrase)) + body = held_body(molmo, toolkit, held_step, prompt_for("held", held_phrase)) eef = np.asarray(state.get(held_step).state["robot0_eef_pos"][:2]) candidate = body["xy"] - eef if body is not None else None if candidate is not None and np.linalg.norm(candidate) <= MAX_HELD: @@ -449,7 +458,7 @@ def finished() -> bool: def run_flash( - toolkit: Any, cell_tag: str, note: Callable[[str], None] = lambda _: None + toolkit: LiberoToolkit, cell_tag: str, note: Callable[[str], None] = lambda _: None ) -> dict: """Replay the program for one cell, named the way the CLI names its cells. @@ -477,7 +486,7 @@ def run_flash( "both plan and anchors files are required" ) - molmo = toolkit.primitives.molmo_client + molmo = toolkit.molmo_client if molmo is None: raise RuntimeError( "no grounder: Flash plans need a Molmo server named by --molmo-endpoint" diff --git a/robots/libero/robot_spec.py b/robots/libero/robot_spec.py index 81c706336..9e85647ef 100644 --- a/robots/libero/robot_spec.py +++ b/robots/libero/robot_spec.py @@ -156,6 +156,7 @@ def get_toolkit( ) return LiberoToolkit( runtime_kwargs=runtime_kwargs, + output_dir=config.output_dir, dashboard_events=dashboard_events, memory=memory, mode=mode, diff --git a/robots/libero/toolkit.py b/robots/libero/toolkit.py index 030fa8686..9b653a344 100644 --- a/robots/libero/toolkit.py +++ b/robots/libero/toolkit.py @@ -12,37 +12,97 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""LIBERO toolkit: common tools + LIBERO primitives. - -Inherits the common file/IO tools from :class:`Toolkit` and registers the -LIBERO primitives (``move_to``, ``pi0_pick``, ``release``, ...) on top. -""" +"""LIBERO tool composition, observations, exploration, and recording lifecycle.""" from __future__ import annotations +import json +from dataclasses import replace from functools import partial from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np from robots.libero import tools as libero_tools -from rpent.dashboard.events import DashboardEventSink -from rpent.session import EnvState -from rpent.tools.toolkit import Toolkit, readonly -from rpent.utils.logging import get_logger, get_output_dir +from rpent.dashboard.events import DashboardEventSink, StepRecordEvent +from rpent.memory import MemoryManager +from rpent.session import EnvState, StepRecord +from rpent.tools import Toolkit, ToolResult +from rpent.utils.logging import get_logger if TYPE_CHECKING: - from rpent.memory.manager import MemoryManager + from robots.libero.env_client import LiberoEnvClient + from rpent.robots.components.molmo_client import MolmoClient + from rpent.robots.components.pi05_vla_client import Pi05VLAClient + from rpent.robots.components.sam3_client import Sam3Client logger = get_logger("libero_toolkit") -class LiberoToolkit(Toolkit): - """Toolkit for the LIBERO robot.""" +class LiberoRuntime: + """Environment clients, cached observations, and session progress.""" + + def __init__( + self, + env: LiberoEnvClient, + model: Pi05VLAClient, + sam3_client: Sam3Client, + molmo_client: MolmoClient | None = None, + flywheel_config: dict[str, Any] | None = None, + ): + self.env = env + self.model = model + self._sam3_client = sam3_client + self.molmo_client = molmo_client + self._flywheel_config = flywheel_config + self._flywheel = None + self._last_obs = None + self._last_obs_eef_pos = None + self._last_obs_gripper = None + self.executed_steps = 0 + self.mode: Literal["evaluation", "exploration"] = "evaluation" + self.solved = False + self.attempt = 1 + self.attempts_per_session = 0 + + def set_obs(self, obs): + self._last_obs = obs + states_arr = np.asarray(obs["states"]) + self._last_obs_eef_pos = np.asarray(states_arr[:3], dtype=np.float32) + # robosuite 2f85: qpos[6] in [~0, ~0.04], qpos[7] in [~-0.04, ~0.]. + # Use |qpos[6]| + |qpos[7]| ≈ finger separation proxy. + # When open ≈ 0.08; when closed ≈ 0. + gp = np.asarray(states_arr[6:8], dtype=np.float32) + self._last_obs_gripper = float(abs(gp[0]) + abs(gp[1])) + + def reset(self) -> None: + obs, _ = self.env.reset() + self.set_obs(obs) + if self._flywheel_config is not None: + from robots.libero.flywheel import create_episode_writer + + self._flywheel = create_episode_writer(self._flywheel_config, obs) + + def execute_primitive(self, name: str, handler: Any, **kwargs: Any) -> ToolResult: + self._flywheel.begin_primitive(name) + try: + return handler(**kwargs) + finally: + self._flywheel.end_primitive() + + def finalize_flywheel(self) -> Path | None: + return self._flywheel.finalize() if self._flywheel is not None else None + + +class LiberoToolkit(Toolkit[LiberoRuntime]): + """Native tools and resources for one LIBERO planner session.""" def __init__( self, *, runtime_kwargs: dict[str, Any], + output_dir: Path | str, dashboard_events: DashboardEventSink, memory: MemoryManager, mode: str = "evaluation", @@ -51,190 +111,399 @@ def __init__( ) -> None: if mode not in {"evaluation", "exploration"}: raise ValueError(f"unsupported LIBERO toolkit mode: {mode!r}") - self._state_output_dir = Path(state_output_dir or get_output_dir()) - state = EnvState(self._state_output_dir) + runtime = LiberoRuntime(**runtime_kwargs) + runtime.mode = mode + runtime.attempts_per_session = max(0, int(attempts_per_session)) + state = EnvState(state_output_dir or output_dir) + tools = libero_tools.LIBERO_TOOLS + if mode == "evaluation": + tools = tuple(tool for tool in tools if tool.name != "reset") + if runtime._flywheel_config is not None: + tools = tuple( + replace( + item, + handler=partial(runtime.execute_primitive, item.name, item.handler), + ) + if not item.readonly and item.name != "finish" + else item + for item in tools + ) super().__init__( - dashboard_events=dashboard_events, state=state, memory=memory, + robot=runtime, + output_dir=output_dir, + tools=tools, + dashboard_events=dashboard_events, ) - self._mode = mode - self._solved: bool = False - self._attempt: int = 1 - # Bound the resettable attempts owned by this planner session. - self._attempts_per_session: int = max(0, int(attempts_per_session)) - self._session_attempt: int = 1 - self.init_primitives(runtime_kwargs=runtime_kwargs) - self._register_libero_tools() - - # ------------------------------------------------------------------ - # Registration - # ------------------------------------------------------------------ - def _register_libero_tools(self) -> None: - # These read-only handlers need the run's EnvState bound in. Every - # other spec binds to a primitive-driver method and captures state by - # default unless that method is explicitly marked @readonly. - 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 spec in libero_tools.TOOLS_SPEC: - name = spec["name"] - if name == "reset" and self._mode != "exploration": - continue - 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 - handler = partial(self._execute_primitive, name, handler) - self.add_tool(name, spec, handler) - if self._mode == "exploration": - reset_spec = next( - spec for spec in libero_tools.TOOLS_SPEC if spec["name"] == "reset" - ) - self.add_tool("reset", reset_spec, self._reset_episode) - finish_spec, finish_handler = self._tools["finish"] - self.add_tool( - "finish", finish_spec, partial(self._guarded_finish, finish_handler) - ) - - def _execute_primitive(self, name: str, handler: Any, **kwargs: Any) -> Any: - self._primitives.begin_primitive(name) + runtime.reset() + record = dump_state(runtime, state, log=None) try: - return handler(**kwargs) - finally: - self._primitives.end_primitive() - - @readonly - def _guarded_finish(self, inner: Any, **kwargs: Any) -> dict[str, Any]: - """Refuse to end an unsolved session while attempts remain.""" - budget = self._attempts_per_session - if budget and not self.solved() and self._session_attempt < budget: - remaining = budget - self._session_attempt - return { - "error": "finish refused", - "reason": ( - f"This session has {remaining} of its {budget} attempts left " - "and the task is not solved. Archive this attempt, call " - "`reset`, and try another approach." - ), - } - return inner(**kwargs) - - def _reset_episode(self, reason: str) -> dict[str, Any]: - """Restart the episode while preserving the full exploration trace.""" - budget = self._attempts_per_session - if budget and self._session_attempt >= budget: - return { - "error": "reset refused", - "reason": ( - f"This session's attempt budget is spent ({budget} attempts). " - "Archive the attempt, update the handoff notes, and call " - "`finish` so the next session can continue." - ), - } - self._attempt += 1 - self._session_attempt += 1 - result = self._primitives.reset_episode(reason=reason) - result["attempt"] = self._attempt - result["notice"] = ( - f"Episode restarted; this is attempt {self._attempt}. The original " - "layout was restored. Re-run perception before acting." - ) - return result + self._dashboard_events.emit( + StepRecordEvent(record=record, env_state=self._state) + ) + except Exception: + logger.exception("Dashboard failed to publish step %s", record.step_idx) + self._action_frame_cursor = 0 - def get_env_state( - self, - *, - command: dict[str, Any], - result: dict[str, Any], - elapsed_s: float, - ) -> dict[str, Any]: + def _capture_observation( + self, *, command: dict[str, Any], result: ToolResult, elapsed_s: float + ) -> tuple[dict[str, Any], list[bytes]]: + """Save the full action log, then assemble the current observation response.""" frame_start = self._action_frame_cursor - self._action_frame_cursor = self._primitives.recorded_frame_count() - record = libero_tools.dump_state( - self._primitives, + self._action_frame_cursor = len(self._frames) + logged_result = result.to_dict() + record = dump_state( + self._robot, self._state, - log={"command": command, "result": result, "elapsed_s": elapsed_s}, + log={ + "command": command, + "result": logged_result, + "elapsed_s": elapsed_s, + }, ) - self._solved |= record.terminated + self._robot.solved |= record.terminated if self._dashboard_events.enabled: try: - frames = self._primitives.frame_slice(frame_start) + frames = self._frames[frame_start:] if frames: - candidate = f"action_{command['action']}.mp4" self._state.save( - candidate, + f"action_{command['action']}.mp4", frames, step=record.step_idx, fps=20, ) - except Exception as e: + except Exception as exc: logger.warning( - "failed to save action clip for step %s: %s", - record.step_idx, - e, + "failed to save action clip for step %s: %s", record.step_idx, exc ) - 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 + record = self._state.get(record.step_idx) + data, images = build_observation(self._state, record) + if result.is_error: + # Report the error once in this response; retain it in the saved log + # so later view_env_state calls can still inspect the failed action. + data["log"]["result"] = { + key: value for key, value in logged_result.items() if key != "error" + } + data["agent_elapsed_s"] = elapsed_s + return data, images @property - def primitives(self) -> "libero_tools.LiberoPrimitives": - """Return the action primitives this toolkit drives.""" - return self._primitives - - def init_primitives( - self, - *, - runtime_kwargs: dict[str, Any], - ) -> None: - """Wipe stale run artifacts, build the LiberoPrimitives, dump step 0.""" - self._state.reset() + def molmo_client(self) -> MolmoClient | None: + """Return the optional point grounder used by task-card replay.""" + return self._robot.molmo_client - primitives = libero_tools.LiberoPrimitives( - check_cancelled=self.raise_if_cancelled, - **runtime_kwargs, - ) - primitives.reset() - primitives.start_recording() - 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 solved(self) -> bool: + return self._robot.solved def close(self) -> None: """Finalize collected data and save the episode video independently.""" try: - episode = self._primitives.finalize_flywheel() + episode = self._robot.finalize_flywheel() if episode is not None: logger.info("flywheel episode finalized: %s", episode) - except Exception as e: - logger.warning("failed to finalize flywheel episode: %s", e) + except Exception as exc: + logger.warning("failed to finalize flywheel episode: %s", exc) try: - 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. - logger.warning(f"failed to save episode video: {e}") - - def solved(self) -> bool: - """Return whether this run has completed the task.""" - return self._solved + if self._frames: + self._state.save("episode.mp4", self._frames, step=None, fps=20) + except Exception as exc: + logger.warning("failed to save episode video: %s", exc) 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( - self._state, recipe_tag, output_dir=get_output_dir() + """Export the successful attempt from the recorded LIBERO trace.""" + return write_recipe_from_states( + self._state, recipe_tag, output_dir=self._task_output_dir + ) + + +def dump_state( + runtime: LiberoRuntime, + env_state: EnvState, + log: dict | None = None, +) -> StepRecord: + """Save one Libero observation through its owned state record.""" + raw = runtime.env.raw_obs() + 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"]], + "robot0_gripper_qpos": [float(x) for x in raw["robot0_gripper_qpos"]], + "object_names": sorted( + k[:-4] + for k in raw + if k.endswith("_pos") and "robot0" not in k and "to_robot" not in k + ), + } + log = log or {} + with env_state.record_step( + state=state, + terminated=runtime.env.terminated, + truncated=runtime.env.truncated, + command=log.get("command"), + result=log.get("result"), + elapsed_s=log.get("elapsed_s"), + extras={"task_language": runtime.env.get_task_language()}, + ) as step_idx: + _save_observation_artifacts(runtime, env_state, step_idx, raw) + return env_state.get(step_idx) + + +def build_observation( + state: EnvState, record: StepRecord +) -> tuple[dict[str, Any], list[bytes]]: + """Assemble a recorded observation and its ordered PNG images.""" + nn = record.step_idx + extras = record.extras + out: dict = { + "step": nn, + "terminated": record.terminated, + "truncated": record.truncated, + "state": record.state, + "artifacts": sorted(record.artifacts), + } + out["task_language"] = extras.get("task_language") + out["log"] = { + "command": record.command, + "result": record.result, + "elapsed_s": record.elapsed_s, + } + images: list[bytes] = [] + for names in ( + ("agentview_policy.png",), + ("agentview_high.png", "agentview.png"), + ("wrist_high.png", "wrist.png"), + ): + name = next((name for name in names if name in record.artifacts), None) + if name: + try: + images.append(state.load_bytes(name, step=nn)) + except FileNotFoundError: + pass + return out, images + + +def _save_observation_artifacts( + runtime: LiberoRuntime, + state: EnvState, + step: int, + raw: dict[str, Any], +) -> None: + """Save the policy view, then each camera's calibrated and high-res views.""" + state.save("agentview_policy.png", runtime._last_obs["main_images"], step=step) + _save_agentview_artifacts(runtime.env, state, step, raw) + _save_wrist_artifacts(runtime.env, state, step, raw) + + for camera, prefix in ( + ("agentview", "agentview"), + ("robot0_eye_in_hand", "wrist"), + ): + try: + rgb, depth = runtime.env.render_camera( + camera_name=camera, height=1024, width=1024, depth=True + ) + camera_meta = runtime.env.get_camera_meta(camera, 1024, 1024) + if camera_meta is None: + raise RuntimeError(f"{camera} camera metadata missing") + state.save(f"{prefix}_high.png", np.asarray(rgb)[::-1], step=step) + depth_metric = _metric_depth(depth, camera_meta)[::-1] + world = _world_from_depth(depth_metric, camera_meta).astype(np.float16) + state.save(f"{prefix}_world_high.npz", world, step=step) + except Exception as exc: + logger.warning("%s high-res dump failed: %s", prefix, exc) + + +def _save_agentview_artifacts( + env: LiberoEnvClient, state: EnvState, step: int, raw: dict[str, Any] +) -> None: + camera_meta = ( + env.get_camera_meta(camera_name="agentview", height=256, width=256) or {} + ) + if camera_meta: + metadata = dict(camera_meta) + metadata["projection"] = ( + "Prefer the back_project(row, col, step=NN) MCP tool; it " + "uses the 1024x1024 high-resolution world map by default. " + "Pass resolution='low' only when row/col came from the " + "256x256 calibration-frame image." ) + metadata["note"] = ( + "The agentview_depth.npz observation is aligned with agentview.png. " + "agentview_policy.png uses the Pi0 orientation and must not supply " + "pixels for back-projection." + ) + state.save("agentview_metadata.json", metadata, step=step) + + try: + image = raw.get("agentview_image") + if image is not None: + image = np.asarray(image, dtype=np.uint8) + state.save("agentview.png", image[::-1], step=step) + except Exception as exc: + logger.warning("image_cam dump failed: %s", exc) + + try: + depth = raw.get("agentview_depth") + if depth is not None: + depth_metric = _metric_depth(depth, camera_meta)[::-1] + state.save( + "agentview_depth.npz", depth_metric.astype(np.float32), step=step + ) + world = _world_from_depth(depth_metric, camera_meta).astype(np.float32) + state.save("agentview_world.npz", world, step=step) + except Exception as exc: + logger.warning("depth dump failed: %s", exc) + + +def _save_wrist_artifacts( + env: LiberoEnvClient, state: EnvState, step: int, raw: dict[str, Any] +) -> None: + try: + image = raw.get("robot0_eye_in_hand_image") + if image is None: + logger.warning("wrist image missing from raw_obs") + else: + image = np.asarray(image, dtype=np.uint8) + state.save("wrist.png", image[::-1], step=step) + except Exception as exc: + logger.warning("wrist image dump failed: %s", exc) + + try: + depth = raw.get("robot0_eye_in_hand_depth") + if depth is None: + logger.warning("wrist depth missing from raw_obs") + return + depth = np.asarray(depth, dtype=np.float32) + height, width = depth.shape[:2] + camera_meta = env.get_camera_meta( + camera_name="robot0_eye_in_hand", height=int(height), width=int(width) + ) + if camera_meta is None: + logger.warning("wrist camera meta missing; skipping wrist depth/world") + return + depth_metric = _metric_depth(depth, camera_meta)[::-1] + state.save("wrist_depth.npz", depth_metric.astype(np.float32), step=step) + world = _world_from_depth(depth_metric, camera_meta).astype(np.float32) + state.save("wrist_world.npz", world, step=step) + metadata = dict(camera_meta) + metadata["note"] = ( + "MOVING camera: extrinsic_cam2world is for THIS step " + "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." + ) + state.save("wrist_metadata.json", metadata, step=step) + except Exception as exc: + logger.warning("wrist depth/world dump failed: %s", exc) + + +def _metric_depth(depth: Any, camera_meta: dict) -> np.ndarray: + """Convert the raw depth buffer to meters before flipping to the RGB frame.""" + d = np.asarray(depth, dtype=np.float32) + if d.ndim == 3: + d = d[..., 0] + near = camera_meta.get("depth_near") + far = camera_meta.get("depth_far") + if near is not None and far is not None: + d = near / (1.0 - d * (1.0 - near / far)) + return d + + +def _world_from_depth(depth_metric: np.ndarray, camera_meta: dict) -> np.ndarray: + """Back-project calibrated image pixels to world XYZ in meters.""" + k_matrix = np.array(camera_meta["intrinsic_K"], dtype=np.float64) + extrinsic = np.array(camera_meta["extrinsic_cam2world"], dtype=np.float64) + fx, fy = k_matrix[0, 0], k_matrix[1, 1] + cx, cy = k_matrix[0, 2], k_matrix[1, 2] + height, width = depth_metric.shape + rr, cc = np.mgrid[0:height, 0:width] + z = depth_metric.astype(np.float64) + camera_points = np.stack( + [(cc - cx) * z / fx, (rr - cy) * z / fy, z, np.ones_like(z)], + axis=-1, + ) + return (camera_points @ extrinsic.T)[..., :3] + + +def _is_primitive_action(name: object) -> bool: + return name in { + "reset", + "pi0_pick", + "pi0_doubled", + "move_to", + "rotate_wrist", + "rotate_pitch", + "move_pose", + "release", + "set_gripper", + } + + +def write_recipe_from_states( + state: EnvState, recipe_tag: str, *, output_dir: Path | str +) -> str: + """Find a command sequence that gets ``terminated=True``. + + Export non-error LIBERO primitive commands and successful segment calls. + """ + records = state.records() + last_reset = max( + ( + record.step_idx + for record in records + if ( + (record.command or {}).get("action") == "reset" + and not (isinstance(record.result, dict) and record.result.get("error")) + ) + ), + default=-1, + ) + command_events = [] + for record in records: + if record.step_idx <= last_reset: + continue + 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)) + + # Never publish a failed trajectory as a recipe. The environment trace is + # authoritative; an agent's self-reported finish status is not. + solved = any( + record.terminated for record in records if record.step_idx > last_reset + ) + if not solved: + return "" + command_events.sort(key=lambda event: event[0]) + recipe_name = f"{recipe_tag}_recipe.jsonl" + recipe_path = Path(output_dir) / recipe_name + recipe_path.parent.mkdir(parents=True, exist_ok=True) + recipe_path.write_text( + "".join(json.dumps(command) + "\n" for _, command in command_events) + ) + return recipe_name diff --git a/robots/libero/tools.py b/robots/libero/tools.py index df6a01e7c..cdfe410af 100644 --- a/robots/libero/tools.py +++ b/robots/libero/tools.py @@ -17,247 +17,293 @@ from __future__ import annotations import json -from collections.abc import Callable -from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Annotated, Literal import numpy as np +from pydantic import Field + +from rpent.tools import ( + ToolContext, + ToolResult, + readonly, + tool, +) + +if TYPE_CHECKING: + from robots.libero.toolkit import LiberoRuntime + + +def _step_env(ctx: ToolContext[LiberoRuntime], action) -> None: + """Submit and count one physical action through the recording path.""" + ctx.check_cancelled() + runtime = ctx.robot + obs, reward, terminated, truncated, _info = runtime.env.step(action) + if runtime._flywheel is not None: + runtime._flywheel.add_transition(action, obs, reward, terminated, truncated) + runtime.executed_steps += 1 + runtime.set_obs(obs) + ctx.record_frame(obs["main_images"]) + + +def _vlm_chunk(ctx: ToolContext[LiberoRuntime], instruction: str): + """One model forward + ``chunk_size`` env steps. Overrides prompt.""" + runtime = ctx.robot + ctx.check_cancelled() + original_task = runtime._last_obs.get("task_descriptions") + try: + runtime._last_obs["task_descriptions"] = instruction + runtime._last_obs.setdefault("extra_view_images", None) -from robots.libero.env_client import LiberoEnvClient -from rpent.robots.components.molmo_client import MolmoClient -from rpent.robots.components.pi05_vla_client import Pi05VLAClient -from rpent.robots.components.sam3_client import Sam3Client -from rpent.session import EnvState, StepRecord -from rpent.tools.toolkit import readonly -from rpent.utils.logging import get_logger + actions = runtime.model.predict(runtime._last_obs, options={"mode": "eval"}) + ctx.check_cancelled() -logger = get_logger("libero") + vla_id = ( + runtime._flywheel.add_proposal(instruction, actions) + if runtime._flywheel is not None + else -1 + ) + chunk_obs, rewards, terminated, truncated, _info = runtime.env.chunk_step( + actions, return_all_frames=True + ) + for index, obs in enumerate(chunk_obs): + ctx.record_frame(obs["main_images"]) + if runtime._flywheel is not None: + runtime._flywheel.add_transition( + actions[index], + obs, + rewards[index], + terminated[index], + truncated[index], + vla_id=vla_id, + proposal_index=index, + ) + runtime.executed_steps += int(np.asarray(terminated).size) + runtime.set_obs(chunk_obs[-1]) + return runtime._last_obs + finally: + if original_task is not None: + runtime._last_obs["task_descriptions"] = original_task -def _normalize_xyz(xyz): - """Coerce an LLM-supplied xyz into a length-3 list[float].""" - if not isinstance(xyz, (list, tuple)) or len(xyz) != 3: - raise ValueError( - 'xyz must be a JSON array of three numbers, e.g. "xyz":[-0.05,0,0.3]' +@tool +def finish(status: str, summary: str, *, ctx: ToolContext[LiberoRuntime]) -> ToolResult: + """Call when the task is complete or unrecoverable. Halts the agent loop. Save any artifacts (recipe, audit) BEFORE calling finish. + + Args: + status: Outcome, e.g. 'success', 'failure', or 'stuck'. + summary: Short natural-language summary of the run. + """ + runtime = ctx.robot + budget = runtime.attempts_per_session + if ( + runtime.mode == "exploration" + and budget + and not runtime.solved + and runtime.attempt < budget + ): + remaining = budget - runtime.attempt + return ToolResult( + error=f"This session has {remaining} of its {budget} attempts left and the task is not solved. Archive this attempt, call `reset`, and try another approach." ) - return [float(v) for v in xyz] + return ToolResult(data={"_finish": True, "status": status, "summary": summary}) -class LiberoPrimitives: - """Wraps a single-env LIBERO-shaped env + VLA policy with primitive- - level methods. +@tool +def reset(reason: str, *, ctx: ToolContext[LiberoRuntime]) -> ToolResult: + """EXPLORE MODE ONLY. Abandon the current episode and restore the same initial scene. Archive the failed attempt first and state which strategy lever will change in the next attempt. - ``pi0_pick`` and ``pi0_doubled`` override ``obs['task_descriptions']`` - with a sub-instruction. ``move_to`` and friends are scripted (no VLM - call) and drive the underlying OSC controller directly. + Args: + reason: Why this episode is unrecoverable and what will change. """ - - def __init__( - self, - env: LiberoEnvClient, - model: Pi05VLAClient, - sam3_client: Sam3Client, - check_cancelled: Callable[[], None], - molmo_client: MolmoClient | None = None, - flywheel_config: dict[str, Any] | None = None, - ): - self.env = env - self.model = model - self._sam3_client = sam3_client - #: Only a Flash Mode replay reads this; other runs never start Molmo. - self.molmo_client = molmo_client - self._check_cancelled = check_cancelled - self._last_obs = None - self._last_obs_eef_pos = None - 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(). - self._recording = False - self._frames = [] - self._flywheel_config = flywheel_config - self._flywheel = None - - def start_recording(self): - self._recording = True - self._frames = [] - - def record_frame(self, obs): - """Append one agentview frame extracted from ``obs`` to the buffer.""" - self._frames.append(np.ascontiguousarray(np.asarray(obs["main_images"]))) - - def recorded_frame_count(self) -> int: - return len(self._frames) - - def stop_recording(self) -> list[np.ndarray]: - frames = list(self._frames) - self._recording = False - self._frames = [] - return frames - - def frame_slice(self, start: int) -> list[np.ndarray]: - return list(self._frames[int(start) :]) - - def set_obs(self, obs): - self._last_obs = obs - states_arr = np.asarray(obs["states"]) - self._last_obs_eef_pos = np.asarray(states_arr[:3], dtype=np.float32) - self._last_obs_eef_z = float(self._last_obs_eef_pos[2]) - # robosuite 2f85: qpos[6] in [~0, ~0.04], qpos[7] in [~-0.04, ~0]. - # Use |qpos[6]| + |qpos[7]| ≈ finger separation proxy. - # When open ≈ 0.08; when closed ≈ 0. - gp = np.asarray(states_arr[6:8], dtype=np.float32) - self._last_obs_gripper = float(abs(gp[0]) + abs(gp[1])) - - def reset(self): - obs, info = self.env.reset() - self.set_obs(obs) - if self._flywheel_config is not None: - from robots.libero.flywheel import create_episode_writer - - self._flywheel = create_episode_writer(self._flywheel_config, obs) - return self._last_obs, info - - def begin_primitive(self, name: str) -> None: - if self._flywheel is not None: - self._flywheel.begin_primitive(name) - - def end_primitive(self) -> None: - if self._flywheel is not None: - self._flywheel.end_primitive() - - def finalize_flywheel(self) -> Path | None: - return self._flywheel.finalize() if self._flywheel is not None else None - - def _step_env(self, action) -> None: - """Execute one env action between cancellation checkpoints.""" - self._check_cancelled() - obs, reward, terminated, truncated, _info = self.env.step(action) - if self._flywheel is not None: - self._flywheel.add_transition(action, obs, reward, terminated, truncated) - self.set_obs(obs) - if self._recording: - self.record_frame(obs) - - def reset_episode(self, reason: str = "") -> dict: - """Restart an episode and return an explore-tool-compatible result.""" - self.reset() - return { + runtime = ctx.robot + budget = runtime.attempts_per_session + if budget and runtime.attempt >= budget: + return ToolResult( + error=f"This session's attempt budget is spent ({budget} attempts). Archive the attempt, update the handoff notes, and call `finish` so the next session can continue." + ) + ctx.check_cancelled() + runtime.attempt += 1 + runtime.reset() + return ToolResult( + data={ "action": "reset", "reason": reason, - "libero_terminated": self.env.terminated or self.env.truncated, + "libero_terminated": runtime.env.terminated or runtime.env.truncated, + "attempt": runtime.attempt, + "notice": f"Episode restarted; this is attempt {runtime.attempt}. The original " + "layout was restored. Re-run perception before acting.", } + ) - def _vlm_chunk(self, instruction: str): - """One model forward + ``chunk_size`` env steps. Overrides prompt.""" - self._check_cancelled() - original_task = self._last_obs.get("task_descriptions") - try: - self._last_obs["task_descriptions"] = instruction - self._last_obs.setdefault("extra_view_images", None) - actions = self.model.predict(self._last_obs, options={"mode": "eval"}) - self._check_cancelled() +@tool +@readonly +def view_env_state( + step: Annotated[int, Field(json_schema_extra={"default": -1})] = -1, + *, + ctx: ToolContext[LiberoRuntime], +) -> ToolResult: + """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 near the gripper, occlusions, and container/cabinet interiors. - vla_id = ( - self._flywheel.add_proposal(instruction, actions) - if self._flywheel is not None - else -1 - ) + Args: + step: Step number; 0 = initial, -1 = latest. + """ + # Toolkit imports these tool definitions while assembling the session. + from robots.libero.toolkit import build_observation - if not self._recording and self._flywheel is None: - chunk_obs, _r, _t, _tr, _i = self.env.chunk_step(actions) - obs = chunk_obs[-1] if self.env.return_all_frames else chunk_obs - else: - chunk_obs, rewards, terminated, truncated, _info = self.env.chunk_step( - actions, return_all_frames=True - ) - for index, obs in enumerate(chunk_obs): - if self._recording: - self.record_frame(obs) - if self._flywheel is not None: - self._flywheel.add_transition( - actions[index], - obs, - rewards[index], - terminated[index], - truncated[index], - vla_id=vla_id, - proposal_index=index, - ) - obs = chunk_obs[-1] - self.set_obs(obs) - return self._last_obs - finally: - if original_task is not None: - self._last_obs["task_descriptions"] = original_task - - def pi0_pick( - self, - prompt: str, - *, - max_chunks: int = 24, - lift_thresh: float = 0.05, - gripper_closed_thresh: float = 0.06, - gripper_open_thresh: float = 0.0, - descent_thresh: float = 0.10, - ) -> dict: - """Closed-loop Pi0.5 pick driven by ``prompt`` as the VLA instruction. - - Success requires the EEF to descend by ``descent_thresh``, then rise by - ``lift_thresh``, with gripper opening in - [``gripper_open_thresh``, ``gripper_closed_thresh``). Terminates early - on LIBERO ``terminated`` (official success) or ``max_chunks``. - """ - instr = prompt - start_z = self._last_obs_eef_z - peak_z = start_z - min_z = start_z - # Track ascent AFTER min_z has been observed — descent then re-ascent - # is the actual "lift" signal, distinct from raw |peak - min| which - # also fires at the BOTTOM of the descent. - post_min_peak_z = start_z - min_grip = self._last_obs_gripper - last_grip = min_grip - descent_done = False - success = False - chunks_used = 0 - - for c in range(max_chunks): - self._vlm_chunk(instr) - chunks_used = c + 1 - z = self._last_obs_eef_z - grip = self._last_obs_gripper - peak_z = max(peak_z, z) - if z < min_z: - min_z = z - post_min_peak_z = z # reset after a new deeper min - else: - post_min_peak_z = max(post_min_peak_z, z) - if (start_z - min_z) >= descent_thresh: - descent_done = True - min_grip = min(min_grip, grip) - last_grip = grip - ascended = (post_min_peak_z - min_z) >= lift_thresh - closed = gripper_open_thresh <= grip < gripper_closed_thresh - if descent_done and ascended and closed: - success = True - break - if self.env.terminated or self.env.truncated: - success = self.env.terminated - break + try: + record = ctx.state.get(step) + except Exception as exc: + return ToolResult(error=f"state step not available: {exc}") + data, images = build_observation(ctx.state, record) + return ToolResult(data=data, images=images) + + +@tool +def move_to( + xyz: Annotated[list[float], Field(min_length=3, max_length=3)], + gripper: float = -1.0, + tol: float = 0.012, + step_clip: float = 0.025, + max_steps: int = 80, + action_scale: float = 0.05, + target_yaw: float | None = None, + yaw_step_clip: float = 0.1, + *, + ctx: ToolContext[LiberoRuntime], +) -> ToolResult: + """Scripted EEF servo to a world-frame XYZ target via the OSC controller. Holds orientation (use rotate_wrist / rotate_pitch / move_pose to reorient). gripper: -1 = open, +1 = close. NEVER command a single move_to with |Δxy| > 0.30 — OSC flips IK and the run corrupts; split long traversal into 2-3 mid waypoints at carry z. + + Args: + xyz: World-frame target [x, y, z] in meters + gripper: Gripper command: -1 open, +1 close (default -1) + tol: Position tolerance, m (default 0.012) + step_clip: Per-step Δxyz cap before action_scale, m (default 0.025) + max_steps: Step budget (default 80) + action_scale: OSC action scale (default 0.05) + target_yaw: Optional world-frame yaw target in radians + yaw_step_clip: Per-step yaw clip, rad (default 0.10) + """ + # Sends 7-D delta actions; the env's underlying OSC_POSE controller + # interprets ``action[:3] ∈ [-1, 1]`` as a per-step desired delta scaled + # by ``action_scale`` (so ``action=1.0`` -> ~5 cm per env step). + runtime = ctx.robot + started = runtime.executed_steps + target = np.asarray(xyz, dtype=np.float32) + for _ in range(max_steps): + cur = runtime._last_obs_eef_pos + diff = target - cur + dist = float(np.linalg.norm(diff)) + if dist < tol: + break + step_dxyz = np.clip(diff, -step_clip, step_clip) + action = np.zeros(7, dtype=np.float32) + action[:3] = step_dxyz / action_scale # -> roughly [-0.5, 0.5] + action[:3] = np.clip(action[:3], -1.0, 1.0) + if target_yaw is not None: + # add wrist yaw control via action[5] (z-axis axis-angle). + # NOTE: extract world yaw via atan2(R[1,0], R[0,0]), NOT + # as_euler('zyx')[0] — the latter returns -world_yaw for + # gripper-down configs (R[2,2]≈-1) and silently flips the + # commanded rotation direction. See feedback_rotate_wrist_yaw_sign. + from scipy.spatial.transform import Rotation as _R + + q = runtime.env.raw_obs()["robot0_eef_quat"] + _R_mat = _R.from_quat([q[0], q[1], q[2], q[3]]).as_matrix() + cur_yaw = float(np.arctan2(_R_mat[1, 0], _R_mat[0, 0])) + err = (float(target_yaw) - cur_yaw + np.pi) % (2 * np.pi) - np.pi + step_dyaw = float(np.clip(err, -yaw_step_clip, yaw_step_clip)) + action[5] = float(np.clip(step_dyaw / 0.10, -1.0, 1.0)) + action[6] = gripper + _step_env(ctx, action) + if runtime.env.terminated or runtime.env.truncated: + break + final = runtime._last_obs_eef_pos + return ToolResult( + data={ + "name": "move_to", + "target_xyz": [float(x) for x in target], + "final_eef_pos": [round(float(x), 4) for x in final], + "final_dist_m": round(float(np.linalg.norm(target - final)), 4), + "steps_used": runtime.executed_steps - started, + "max_steps": max_steps, + "terminated": runtime.env.terminated, + "truncated": runtime.env.truncated, + } + ) - return { + +@tool +def pi0_pick( + prompt: str, + max_chunks: int = 24, + lift_thresh: float = 0.05, + gripper_closed_thresh: float = 0.06, + gripper_open_thresh: float = 0.0, + descent_thresh: float = 0.10, + *, + ctx: ToolContext[LiberoRuntime], +) -> ToolResult: + """Pi0.5 closed-loop pick. Use it for the grasp; YOU then do every move_to and release. Use modest max_chunks and verify the grasp from EEF lift, gripper closure, and available images. + + Args: + prompt: Pi0 prompt (e.g. 'pick up the akita black bowl'). + max_chunks: Action-chunk budget (default 24) + lift_thresh: EEF post-descent ascent threshold for success, m (default 0.05) + gripper_closed_thresh: Finger-separation closed threshold (default 0.06) + gripper_open_thresh: Minimum finger separation accepted as a held object (default 0.0) + descent_thresh: Required descent before lift detection, m (default 0.10) + """ + runtime = ctx.robot + start_z = float(runtime._last_obs_eef_pos[2]) + peak_z = start_z + min_z = start_z + # Track ascent AFTER min_z has been observed — descent then re-ascent + # is the actual "lift" signal, distinct from raw |peak - min| which + # also fires at the BOTTOM of the descent. + post_min_peak_z = start_z + min_grip = runtime._last_obs_gripper + last_grip = min_grip + descent_done = False + success = False + chunks_used = 0 + + for c in range(max_chunks): + _vlm_chunk(ctx, prompt) + chunks_used = c + 1 + z = float(runtime._last_obs_eef_pos[2]) + grip = runtime._last_obs_gripper + peak_z = max(peak_z, z) + if z < min_z: + min_z = z + post_min_peak_z = z # reset after a new deeper min + else: + post_min_peak_z = max(post_min_peak_z, z) + if (start_z - min_z) >= descent_thresh: + descent_done = True + min_grip = min(min_grip, grip) + last_grip = grip + ascended = (post_min_peak_z - min_z) >= lift_thresh + closed = gripper_open_thresh <= grip < gripper_closed_thresh + if descent_done and ascended and closed: + success = True + break + if runtime.env.terminated or runtime.env.truncated: + success = runtime.env.terminated + break + + return ToolResult( + data={ "name": "pick", - "instruction": instr, + "instruction": prompt, "success": success, "chunks_used": chunks_used, "max_chunks": max_chunks, "peak_lift_m": post_min_peak_z - min_z, # actual post-descent ascent "min_gripper_opening": min_grip, "final_gripper_opening": last_grip, - "terminated": self.env.terminated, - "truncated": self.env.truncated, + "terminated": runtime.env.terminated, + "truncated": runtime.env.truncated, "diagnostics": { "start_eef_z": round(start_z, 4), "peak_eef_z": round(peak_z, 4), @@ -272,40 +318,41 @@ def pi0_pick( "descent_thresh": descent_thresh, }, } + ) - def pi0_doubled( - self, - prompt: str, - *, - max_chunks: int = 20, - ) -> dict: - """Closed-loop Pi0.5 contact skill. - - Intended for non-pick contact interactions such as turning knobs, - toggling stoves, or short pushes. Success is the official LIBERO - termination predicate, not a private object-pose oracle. - """ - instr = prompt - task_success = False - chunks_used = 0 - - for c in range(max_chunks): - self._vlm_chunk(instr) - chunks_used = c + 1 - if self.env.terminated or self.env.truncated: - task_success = self.env.terminated - break - return { +@tool +def pi0_doubled( + prompt: str, max_chunks: int = 20, *, ctx: ToolContext[LiberoRuntime] +) -> ToolResult: + """Pi0.5 closed-loop contact skill for non-pick interactions (e.g. stove/knob/button/short push). Returned success/task_success only mirrors official termination; for intermediate contact skills, success=false does not necessarily mean the contact interaction failed. Inspect image/state evidence. Do not use it as a general pick/place shortcut. + + Args: + prompt: Contact-skill prompt, e.g. 'turn on the stove'. + max_chunks: Action-chunk budget (default 20) + """ + runtime = ctx.robot + task_success = False + chunks_used = 0 + + for c in range(max_chunks): + _vlm_chunk(ctx, prompt) + chunks_used = c + 1 + if runtime.env.terminated or runtime.env.truncated: + task_success = runtime.env.terminated + break + + return ToolResult( + data={ "name": "pi0_doubled", - "instruction": instr, + "instruction": prompt, "success": task_success, "task_success": task_success, "contact_skill_executed": chunks_used > 0, "chunks_used": chunks_used, "max_chunks": max_chunks, - "terminated": self.env.terminated, - "truncated": self.env.truncated, + "terminated": runtime.env.terminated, + "truncated": runtime.env.truncated, "diagnostics": { "mode": "contact_skill_success_by_termination", "success_meaning": ( @@ -314,139 +361,137 @@ def pi0_doubled( ), }, } + ) - def move_to( - self, - xyz, - *, - max_steps: int = 80, - gripper: float = -1.0, - step_clip: float = 0.025, - tol: float = 0.012, - action_scale: float = 0.05, - target_yaw: float | None = None, - yaw_step_clip: float = 0.10, - ) -> dict: - """Scripted EEF servo to a world-frame target xyz. - - Sends 7-D delta actions; the env's underlying OSC_POSE controller - interprets ``action[:3] ∈ [-1, 1]`` as a per-step desired delta scaled - by ``action_scale`` (so ``action=1.0`` -> ~5 cm per env step). - ``gripper``: +1.0 keeps it closed (holding object), -1.0 opens. - """ - target = np.asarray(_normalize_xyz(xyz), dtype=np.float32) - traj = [] - for step in range(max_steps): - cur = self._last_obs_eef_pos - diff = target - cur - dist = float(np.linalg.norm(diff)) - traj.append( - { - "step": step, - "eef_pos": [round(float(x), 4) for x in cur], - "dist_to_target_m": round(dist, 4), - } - ) - if dist < tol: - break - step_dxyz = np.clip(diff, -step_clip, step_clip) - action = np.zeros(7, dtype=np.float32) - action[:3] = step_dxyz / action_scale # -> roughly [-0.5, 0.5] - action[:3] = np.clip(action[:3], -1.0, 1.0) - if target_yaw is not None: - # add wrist yaw control via action[5] (z-axis axis-angle). - # NOTE: extract world yaw via atan2(R[1,0], R[0,0]), NOT - # as_euler('zyx')[0] — the latter returns -world_yaw for - # gripper-down configs (R[2,2]≈-1) and silently flips the - # commanded rotation direction. See feedback_rotate_wrist_yaw_sign. - from scipy.spatial.transform import Rotation as _R - - q = self.env.raw_obs()["robot0_eef_quat"] - _R_mat = _R.from_quat([q[0], q[1], q[2], q[3]]).as_matrix() - cur_yaw = float(np.arctan2(_R_mat[1, 0], _R_mat[0, 0])) - err = (float(target_yaw) - cur_yaw + np.pi) % (2 * np.pi) - np.pi - step_dyaw = float(np.clip(err, -yaw_step_clip, yaw_step_clip)) - action[5] = float(np.clip(step_dyaw / 0.10, -1.0, 1.0)) - action[6] = gripper - self._step_env(action) - if self.env.terminated or self.env.truncated: - break - final = self._last_obs_eef_pos - return { - "name": "move_to", - "target_xyz": [float(x) for x in target], - "final_eef_pos": [round(float(x), 4) for x in final], - "final_dist_m": round(float(np.linalg.norm(target - final)), 4), - "steps_used": len(traj), - "max_steps": max_steps, - "terminated": self.env.terminated, - "truncated": self.env.truncated, + +@tool +def release(max_steps: int = 20, *, ctx: ToolContext[LiberoRuntime]) -> ToolResult: + """Open the gripper for up to max_steps env steps while holding EEF in place. Triggers libero termination if the matching On/In predicate is met. + + Args: + max_steps: Step budget (default 20) + """ + runtime = ctx.robot + started = runtime.executed_steps + assert max_steps > 0, f"max_steps must be > 0, got {max_steps}" + start_grip = runtime._last_obs_gripper + peak_grip = start_grip + for _ in range(max_steps): + action = np.zeros(7, dtype=np.float32) + action[6] = -1.0 # open + _step_env(ctx, action) + peak_grip = max(peak_grip, runtime._last_obs_gripper) + if runtime.env.terminated or runtime.env.truncated: + break + return ToolResult( + data={ + "name": "release", + "steps_used": runtime.executed_steps - started, + "start_gripper_opening": round(start_grip, 4), + "peak_gripper_opening": round(peak_grip, 4), + "final_gripper_opening": round(runtime._last_obs_gripper, 4), + "terminated": runtime.env.terminated, + "truncated": runtime.env.truncated, } + ) - def rotate_wrist( - self, - *, - target_yaw: float | None = None, - delta_yaw: float | None = None, - gripper: float = 1.0, - max_steps: int = 40, - tol: float = 0.02, - step_clip: float = 0.10, - ) -> dict: - """Rotate wrist around world z-axis. Provide EITHER target_yaw (absolute) - or delta_yaw (relative, applied as a single rotation goal). - - Uses ``action[5]`` (axis-angle z component) to drive wrist yaw via the - OSC controller. Holds xyz pose constant during rotation. - - Yaw is the world-frame z-rotation, recovered as - ``atan2(R[1,0], R[0,0])`` where R is the eef rotation matrix in the - world frame. (Note: ``as_euler('zyx')[0]`` returns the *negative* - of this value for gripper-down configurations because the Z-Y-X - decomposition picks the chart with γ ≈ π, flipping α. Bug fixed - 2026-05-19 — previous implementation rotated the wrist in the - opposite direction of the commanded yaw.) - """ - from scipy.spatial.transform import Rotation as _R - - def _yaw_of(quat_xyzw): - # robot0_eef_quat in libero+robosuite is xyzw (scipy convention). - q = quat_xyzw - rot = _R.from_quat([q[0], q[1], q[2], q[3]]) - R = rot.as_matrix() - # World-frame yaw: angle of the eef x-axis projected onto the - # world xy plane. Robust to gripper-down (R[2,2]≈-1) which is - # where the euler 'zyx' chart flips sign. - return float(np.arctan2(R[1, 0], R[0, 0])) - - raw = self.env.raw_obs() - cur_quat = raw["robot0_eef_quat"] - start_yaw = _yaw_of(cur_quat) - if target_yaw is None and delta_yaw is None: - return {"name": "rotate_wrist", "error": "need target_yaw or delta_yaw"} - if target_yaw is None: - target_yaw = start_yaw + float(delta_yaw) - - traj = [] - for step in range(max_steps): - raw = self.env.raw_obs() - cur_yaw = _yaw_of(raw["robot0_eef_quat"]) - err = float(target_yaw - cur_yaw) - # wrap to [-pi, pi] - err = (err + np.pi) % (2 * np.pi) - np.pi - traj.append({"step": step, "yaw": round(cur_yaw, 4), "err": round(err, 4)}) - if abs(err) < tol: - break - step_dyaw = float(np.clip(err, -step_clip, step_clip)) - action = np.zeros(7, dtype=np.float32) - action[5] = step_dyaw / 0.10 # scale to ~[-1,1] action range - action[5] = float(np.clip(action[5], -1.0, 1.0)) - action[6] = float(gripper) - self._step_env(action) - if self.env.terminated or self.env.truncated: - break - final_yaw = _yaw_of(self.env.raw_obs()["robot0_eef_quat"]) - return { + +@tool +def set_gripper( + gripper: float = -1.0, steps: int = 5, *, ctx: ToolContext[LiberoRuntime] +) -> ToolResult: + """Hold the current EEF pose and drive the gripper command for `steps` env steps. Use to firm up a grip mid-carry. + + Args: + gripper: Gripper command: -1 open, +1 close (default -1) + steps: Number of env steps (default 5) + """ + runtime = ctx.robot + started = runtime.executed_steps + for _ in range(steps): + action = np.zeros(7, dtype=np.float32) + action[6] = gripper + _step_env(ctx, action) + if runtime.env.terminated or runtime.env.truncated: + break + return ToolResult( + data={ + "name": "set_gripper", + "gripper": gripper, + "steps": runtime.executed_steps - started, + "terminated": runtime.env.terminated, + "truncated": runtime.env.truncated, + } + ) + + +@tool +def rotate_wrist( + target_yaw: float | None = None, + delta_yaw: float | None = None, + gripper: float = 1.0, + max_steps: int = 40, + tol: float = 0.02, + step_clip: float = 0.1, + *, + ctx: ToolContext[LiberoRuntime], +) -> ToolResult: + """Rotate the wrist around the world Z-axis. Provide either target_yaw (absolute) or delta_yaw (relative). Holds xyz fixed. + + Args: + target_yaw: Absolute world-frame yaw target, rad + delta_yaw: Relative yaw delta, rad + gripper: Gripper command held during rotation (default +1) + max_steps: Step budget (default 40) + tol: Yaw tolerance, rad (default 0.02) + step_clip: Per-step yaw clip, rad (default 0.10) + """ + # Uses ``action[5]`` (axis-angle z component) to drive wrist yaw via the + # OSC controller. Holds xyz pose constant during rotation. + from scipy.spatial.transform import Rotation as _R + + runtime = ctx.robot + started = runtime.executed_steps + + def _yaw_of(quat_xyzw): + # robot0_eef_quat in libero+robosuite is xyzw (scipy convention). + q = quat_xyzw + rot = _R.from_quat([q[0], q[1], q[2], q[3]]) + R = rot.as_matrix() + # World-frame yaw: angle of the eef x-axis projected onto the + # world xy plane. Robust to gripper-down (R[2,2]≈-1) which is + # where the euler 'zyx' chart flips sign. + return float(np.arctan2(R[1, 0], R[0, 0])) + + raw = runtime.env.raw_obs() + cur_quat = raw["robot0_eef_quat"] + start_yaw = _yaw_of(cur_quat) + if target_yaw is None and delta_yaw is None: + return ToolResult( + data={"name": "rotate_wrist"}, error="need target_yaw or delta_yaw" + ) + if target_yaw is None: + target_yaw = start_yaw + float(delta_yaw) + + for _ in range(max_steps): + raw = runtime.env.raw_obs() + cur_yaw = _yaw_of(raw["robot0_eef_quat"]) + err = float(target_yaw - cur_yaw) + # wrap to [-pi, pi] + err = (err + np.pi) % (2 * np.pi) - np.pi + if abs(err) < tol: + break + step_dyaw = float(np.clip(err, -step_clip, step_clip)) + action = np.zeros(7, dtype=np.float32) + action[5] = step_dyaw / 0.10 # scale to ~[-1,1] action range + action[5] = float(np.clip(action[5], -1.0, 1.0)) + action[6] = gripper + _step_env(ctx, action) + if runtime.env.terminated or runtime.env.truncated: + break + final_yaw = _yaw_of(runtime.env.raw_obs()["robot0_eef_quat"]) + return ToolResult( + data={ "name": "rotate_wrist", "start_yaw": round(start_yaw, 4), "target_yaw": round(float(target_yaw), 4), @@ -454,81 +499,86 @@ def _yaw_of(quat_xyzw): "final_err": round( float((target_yaw - final_yaw + np.pi) % (2 * np.pi) - np.pi), 4 ), - "steps_used": len(traj), - "terminated": self.env.terminated, - "truncated": self.env.truncated, + "steps_used": runtime.executed_steps - started, + "terminated": runtime.env.terminated, + "truncated": runtime.env.truncated, } + ) - def rotate_pitch( - self, - *, - target_pitch: float | None = None, - delta_pitch: float | None = None, - gripper: float = 1.0, - max_steps: int = 40, - tol: float = 0.02, - step_clip: float = 0.10, - ) -> dict: - """Tilt the gripper around the world X-axis ("pitch"). - - Pitch is defined as the angle between the eef z-axis and the - world -z direction, measured in the world yz-plane: - - pitch = atan2(R[1, 2], -R[2, 2]) - - - pitch = 0 -> gripper z-axis aligned with world -z (default - "gripper down" rest pose). - - pitch = +pi/2 -> gripper z-axis points in world +y (gripper - "looking forward" along world +y). - - pitch = -pi/2 -> gripper z-axis points in world -y. - - Driven by ``action[3]`` (axis-angle X component) of the OSC_POSE - controller. Sign verified empirically (probe_pitch.py 2026-05-19): - action[3]=+1.0 tilts eef z toward world +y, matching this pitch - definition with no sign flip. - - Holds xyz, yaw, and gripper constant during rotation. Use BEFORE - threading the gripper into a narrow opening whose front face - normal is along world ±y (e.g. microwave cavity in libero_10 t9). - - Provide EITHER ``target_pitch`` (absolute) or ``delta_pitch`` - (relative). Both in radians. - """ - from scipy.spatial.transform import Rotation as _R - - def _pitch_of(quat_xyzw): - q = quat_xyzw - R = _R.from_quat([q[0], q[1], q[2], q[3]]).as_matrix() - return float(np.arctan2(R[1, 2], -R[2, 2])) - - raw = self.env.raw_obs() - start_pitch = _pitch_of(raw["robot0_eef_quat"]) - if target_pitch is None and delta_pitch is None: - return {"name": "rotate_pitch", "error": "need target_pitch or delta_pitch"} - if target_pitch is None: - target_pitch = start_pitch + float(delta_pitch) - - traj = [] - for step in range(max_steps): - raw = self.env.raw_obs() - cur_pitch = _pitch_of(raw["robot0_eef_quat"]) - err = float(target_pitch - cur_pitch) - err = (err + np.pi) % (2 * np.pi) - np.pi - traj.append( - {"step": step, "pitch": round(cur_pitch, 4), "err": round(err, 4)} - ) - if abs(err) < tol: - break - step_dpitch = float(np.clip(err, -step_clip, step_clip)) - action = np.zeros(7, dtype=np.float32) - action[3] = step_dpitch / 0.10 - action[3] = float(np.clip(action[3], -1.0, 1.0)) - action[6] = float(gripper) - self._step_env(action) - if self.env.terminated or self.env.truncated: - break - final_pitch = _pitch_of(self.env.raw_obs()["robot0_eef_quat"]) - return { + +@tool +def rotate_pitch( + target_pitch: float | None = None, + delta_pitch: float | None = None, + gripper: float = 1.0, + max_steps: int = 40, + tol: float = 0.02, + step_clip: float = 0.1, + *, + ctx: ToolContext[LiberoRuntime], +) -> ToolResult: + """Tilt the gripper around the world X-axis. Provide either target_pitch (absolute) or delta_pitch (relative). Holds xyz and yaw fixed. Use before threading the gripper into a narrow opening whose front face normal is along world ±y (e.g. microwave cavity). + + Args: + target_pitch: Absolute world-frame pitch target, rad + delta_pitch: Relative pitch delta, rad + gripper: Gripper command held during rotation (default +1) + max_steps: Step budget (default 40) + tol: Pitch tolerance, rad (default 0.02) + step_clip: Per-step pitch clip, rad (default 0.10) + """ + # Pitch is defined as the angle between the eef z-axis and the + # world -z direction, measured in the world yz-plane: + # + # pitch = atan2(R[1, 2], -R[2, 2]) + # + # - pitch = 0 -> gripper z-axis aligned with world -z (default + # "gripper down" rest pose). + # - pitch = +pi/2 -> gripper z-axis points in world +y (gripper + # "looking forward" along world +y). + # - pitch = -pi/2 -> gripper z-axis points in world -y. + # + # Driven by ``action[3]`` (axis-angle X component) of the OSC_POSE + # controller. Sign verified empirically (probe_pitch.py 2026-05-19): + # action[3]=+1.0 tilts eef z toward world +y, matching this pitch + # definition with no sign flip. + from scipy.spatial.transform import Rotation as _R + + runtime = ctx.robot + started = runtime.executed_steps + + def _pitch_of(quat_xyzw): + q = quat_xyzw + R = _R.from_quat([q[0], q[1], q[2], q[3]]).as_matrix() + return float(np.arctan2(R[1, 2], -R[2, 2])) + + raw = runtime.env.raw_obs() + start_pitch = _pitch_of(raw["robot0_eef_quat"]) + if target_pitch is None and delta_pitch is None: + return ToolResult( + data={"name": "rotate_pitch"}, error="need target_pitch or delta_pitch" + ) + if target_pitch is None: + target_pitch = start_pitch + float(delta_pitch) + + for _ in range(max_steps): + raw = runtime.env.raw_obs() + cur_pitch = _pitch_of(raw["robot0_eef_quat"]) + err = float(target_pitch - cur_pitch) + err = (err + np.pi) % (2 * np.pi) - np.pi + if abs(err) < tol: + break + step_dpitch = float(np.clip(err, -step_clip, step_clip)) + action = np.zeros(7, dtype=np.float32) + action[3] = step_dpitch / 0.10 + action[3] = float(np.clip(action[3], -1.0, 1.0)) + action[6] = gripper + _step_env(ctx, action) + if runtime.env.terminated or runtime.env.truncated: + break + final_pitch = _pitch_of(runtime.env.raw_obs()["robot0_eef_quat"]) + return ToolResult( + data={ "name": "rotate_pitch", "start_pitch": round(start_pitch, 4), "target_pitch": round(float(target_pitch), 4), @@ -536,1361 +586,384 @@ def _pitch_of(quat_xyzw): "final_err": round( float((target_pitch - final_pitch + np.pi) % (2 * np.pi) - np.pi), 4 ), - "steps_used": len(traj), - "terminated": self.env.terminated, - "truncated": self.env.truncated, - } - - def move_pose( - self, - xyz, - *, - target_pitch: float | None = None, - target_yaw: float | None = None, - gripper: float = -1.0, - step_clip: float = 0.02, - pitch_step: float = 0.08, - yaw_step: float = 0.08, - tol: float = 0.012, - ori_tol: float = 0.05, - action_scale: float = 0.05, - max_steps: int = 150, - ) -> dict: - """Servo position AND orientation (pitch + yaw) SIMULTANEOUSLY. - - Unlike ``move_to`` (holds orientation) + ``rotate_pitch`` (holds - xyz), this co-varies xyz and wrist tilt every env.step. Co-variation - lets the OSC controller thread cabinet-front-low poses where a - decoupled position servo (fixed gripper-down orientation) drives - the wrist into a singularity and stalls — mimicking pi0's curved - reach-in. - """ - from scipy.spatial.transform import Rotation as _R - - def _pitch_of(q): - R = _R.from_quat([q[0], q[1], q[2], q[3]]).as_matrix() - return float(np.arctan2(R[1, 2], -R[2, 2])) - - def _yaw_of(q): - R = _R.from_quat([q[0], q[1], q[2], q[3]]).as_matrix() - return float(np.arctan2(R[1, 0], R[0, 0])) - - target = np.asarray(_normalize_xyz(xyz), dtype=np.float32) - traj = [] - step = 0 - for step in range(max_steps): - cur = self._last_obs_eef_pos - q = self.env.raw_obs()["robot0_eef_quat"] - diff = target - cur - dist = float(np.linalg.norm(diff)) - p_err = ( - 0.0 - if target_pitch is None - else float((target_pitch - _pitch_of(q) + np.pi) % (2 * np.pi) - np.pi) - ) - y_err = ( - 0.0 - if target_yaw is None - else float((target_yaw - _yaw_of(q) + np.pi) % (2 * np.pi) - np.pi) - ) - traj.append( - { - "step": step, - "eef": [round(float(x), 4) for x in cur], - "dist": round(dist, 4), - "p_err": round(p_err, 3), - } - ) - if dist < tol and abs(p_err) < ori_tol and abs(y_err) < ori_tol: - break - action = np.zeros(7, dtype=np.float32) - sd = np.clip(diff, -step_clip, step_clip) - action[:3] = np.clip(sd / action_scale, -1.0, 1.0) - action[3] = float( - np.clip(np.clip(p_err, -pitch_step, pitch_step) / 0.10, -1.0, 1.0) - ) - action[5] = float( - np.clip(np.clip(y_err, -yaw_step, yaw_step) / 0.10, -1.0, 1.0) - ) - action[6] = float(gripper) - self._step_env(action) - if self.env.terminated or self.env.truncated: - break - final = self._last_obs_eef_pos - fq = self.env.raw_obs()["robot0_eef_quat"] - return { - "name": "move_pose", - "final_eef_pos": [round(float(x), 4) for x in final], - "final_dist_m": round(float(np.linalg.norm(target - final)), 4), - "final_pitch": round(_pitch_of(fq), 4), - "steps_used": step + 1, - "terminated": self.env.terminated, - "truncated": self.env.truncated, - } - - def release( - self, - *, - max_steps: int = 20, - ) -> dict: - """Open gripper for ``max_steps`` env steps while keeping eef in place. - - Returns once libero terminates (success) or step budget exhausted. - """ - assert max_steps > 0, f"max_steps must be > 0, got {max_steps}" - start_grip = self._last_obs_gripper - peak_grip = start_grip - for step in range(max_steps): - action = np.zeros(7, dtype=np.float32) - action[6] = -1.0 # open - self._step_env(action) - peak_grip = max(peak_grip, self._last_obs_gripper) - if self.env.terminated or self.env.truncated: - break - return { - "name": "release", - "steps_used": step + 1, - "start_gripper_opening": round(start_grip, 4), - "peak_gripper_opening": round(peak_grip, 4), - "final_gripper_opening": round(self._last_obs_gripper, 4), - "terminated": self.env.terminated, - "truncated": self.env.truncated, - } - - def set_gripper( - self, - *, - gripper: float = -1.0, - steps: int = 5, - ) -> dict: - """Hold the current EEF pose and drive ``gripper`` for ``steps`` env steps.""" - g = float(gripper) - n = int(steps) - for _ in range(n): - action = np.zeros(7, dtype=np.float32) - action[6] = g - self._step_env(action) - if self.env.terminated or self.env.truncated: - break - return { - "name": "set_gripper", - "gripper": g, - "steps": n, - "terminated": self.env.terminated, - "truncated": self.env.truncated, - } - - # ---- introspection helpers (for LLM-in-the-loop) ---- - - @readonly - def segment( - self, - prompt: str = "", - camera: str = "agentview", - 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. - - This tool deliberately does not render camera views or create wrist/high-res - artifacts. Errors are structured so the agent can continue with image - inspection and ``back_project``. - """ - 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() - has_prompt = bool(prompt) - has_point = point is not None - if has_prompt == has_point: - return {"error": "segment needs exactly one of prompt or point"} - try: - image_name, world_name, artifact_pairs = _select_segment_artifacts( - state, record, camera - ) - except ValueError as e: - return {"error": str(e)} - if image_name is None or world_name is None: - return { - "error": "complete segment artifacts not found", - "step": nn, - "camera": camera, - "checked_artifacts": [ - name - for image, world in artifact_pairs - for name in (image, world) - if name - ], - } - - try: - data = self._sam3_client.segment( - state.load_bytes(image_name, step=nn), - text_prompt=prompt if has_prompt else None, - point=point, - min_score=min_score, - ) - except ValueError as e: - return { - "error": str(e), - "step": nn, - "camera": camera, - "image_artifact": image_name, - } - except Exception as e: - return { - "error": f"segmentation service call failed: {e}", - "step": nn, - "camera": camera, - "image_artifact": image_name, - "fallback": "Use manual visual localization and back_project.", - } - - 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): - try: - world_map = state.load(world_name, step=nn) - except Exception as exc: - world_result = { - "world_xyz": None, - "world_error": f"world map artifact not available: {exc}", - "expected_world_artifact": world_name, - } - else: - 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, - "world_error": data.reason or "segmentation did not find a mask", - } - - segment_blob = { - "found": data.found, - "mode": "text" if has_prompt else "point", - "camera": camera, - "source_step": nn, - "segment_index": segment_index, - "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, - "mask_shape": list(data.mask_shape) if data.mask_shape else None, - } - if has_prompt: - segment_blob["prompt"] = prompt - else: - segment_blob["point"] = point - if not data.found: - segment_blob["error"] = data.reason or "SAM3 found no mask" - segment_blob.update(world_result) - saved_segment = state.save( - segment_name, - segment_blob, - step=nn, - ) - - result = { - "found": data.found, - "step": nn, - "camera": camera, - "image_artifact": image_name, - "score": segment_blob["score"], - "box": segment_blob["box"], - "world_xyz": segment_blob["world_xyz"], - "world_error": segment_blob.get("world_error"), + "steps_used": runtime.executed_steps - started, + "terminated": runtime.env.terminated, + "truncated": runtime.env.truncated, } - if saved_segment is None: - result["error"] = f"failed to persist segment artifact {segment_name}" - result["code"] = "segment_artifact_save_failed" - result["attempted_segment_artifact"] = segment_name - if "error" in segment_blob: - result["segmentation_error"] = segment_blob["error"] - result["fallback"] = "Use manual visual localization and back_project." - else: - result["segment_artifact"] = saved_segment - if saved_segment is not None and "error" in segment_blob: - result["error"] = segment_blob["error"] - result["fallback"] = "Use manual visual localization and back_project." - if saved_overlay is not None: - result["overlay_artifact"] = saved_overlay - result["_image_bytes"] = state.load_bytes(saved_overlay, step=nn) - return result - - -def _is_primitive_action(name: object) -> bool: - """Whether ``name`` is a state-advancing LIBERO primitive. - - A primitive is any non-read-only method on :class:`LiberoPrimitives`; - read-only tools and non-strings read as ``False``. - """ - if not isinstance(name, str): - return False - method = getattr(LiberoPrimitives, name, None) - return method is not None and not bool(getattr(method, "_readonly", False)) - + ) -def write_recipe_from_states( - state: EnvState, recipe_tag: str, *, output_dir: Path | str -) -> str: - """Find a command sequence that gets ``terminated=True``. - Export non-error LIBERO primitive commands and successful segment calls. +@tool +def move_pose( + xyz: Annotated[list[float], Field(min_length=3, max_length=3)], + target_pitch: float | None = None, + target_yaw: float | None = None, + gripper: float = -1.0, + step_clip: float = 0.02, + pitch_step: float = 0.08, + yaw_step: float = 0.08, + tol: float = 0.012, + ori_tol: float = 0.05, + action_scale: float = 0.05, + max_steps: int = 150, + *, + ctx: ToolContext[LiberoRuntime], +) -> ToolResult: + """Servo position AND orientation (pitch + yaw) SIMULTANEOUSLY. Unlike move_to (holds orientation) + rotate_pitch (holds xyz), this co-varies xyz and wrist tilt every env.step. Use to thread cabinet-front / low-shelf poses where a decoupled position servo drives the wrist into an IK singularity and stalls. + + Args: + xyz: World-frame target [x, y, z] in meters + target_pitch: Absolute pitch target, rad + target_yaw: Absolute yaw target, rad + gripper: Gripper command held during the move (default -1) + step_clip: Per-step Δxyz cap, m (default 0.02) + pitch_step: Per-step pitch clip, rad (default 0.08) + yaw_step: Per-step yaw clip, rad (default 0.08) + tol: Position tolerance, m (default 0.012) + ori_tol: Orientation tolerance, rad (default 0.05) + action_scale: OSC action scale (default 0.05) + max_steps: Step budget (default 150) """ - records = state.records() - last_reset = max( - ( - record.step_idx - for record in records - if ( - (record.command or {}).get("action") == "reset" - and not (isinstance(record.result, dict) and record.result.get("error")) - ) - ), - default=-1, - ) - command_events = [] - for record in records: - if record.step_idx <= last_reset: - continue - 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)) - - # Never publish a failed trajectory as a recipe. The environment trace is - # authoritative; an agent's self-reported finish status is not. - solved = any( - record.terminated for record in records if record.step_idx > last_reset - ) - if not solved: - return "" - command_events.sort(key=lambda event: event[0]) - recipe_name = f"{recipe_tag}_recipe.jsonl" - recipe_path = Path(output_dir) / recipe_name - recipe_path.parent.mkdir(parents=True, exist_ok=True) - recipe_path.write_text( - "".join(json.dumps(command) + "\n" for _, command in command_events) - ) - return recipe_name - - -def _metric_depth(depth: Any, camera_meta: dict) -> np.ndarray: - d = np.asarray(depth, dtype=np.float32) - if d.ndim == 3: - d = d[..., 0] - near = camera_meta.get("depth_near") - far = camera_meta.get("depth_far") - if near is not None and far is not None: - d = near / (1.0 - d * (1.0 - near / far)) - return d - - -def _world_from_depth(depth_metric: np.ndarray, camera_meta: dict) -> np.ndarray: - k_matrix = np.array(camera_meta["intrinsic_K"], dtype=np.float64) - extrinsic = np.array(camera_meta["extrinsic_cam2world"], dtype=np.float64) - fx, fy = k_matrix[0, 0], k_matrix[1, 1] - cx, cy = k_matrix[0, 2], k_matrix[1, 2] - height, width = depth_metric.shape - rr, cc = np.mgrid[0:height, 0:width] - z = depth_metric.astype(np.float64) - camera_points = np.stack( - [(cc - cx) * z / fx, (rr - cy) * z / fy, z, np.ones_like(z)], - axis=-1, - ) - return (camera_points @ extrinsic.T)[..., :3] - - -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() - 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"]], - "robot0_gripper_qpos": [float(x) for x in raw["robot0_gripper_qpos"]], - "object_names": sorted( - k[:-4] - for k in raw - if k.endswith("_pos") and "robot0" not in k and "to_robot" not in k - ), - } - log = log or {} - with env_state.record_step( - state=state, - terminated=primitives.env.terminated, - truncated=primitives.env.truncated, - command=log.get("command"), - result=log.get("result"), - elapsed_s=log.get("elapsed_s"), - extras={ - "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 --- - agentview_meta = ( - primitives.env.get_camera_meta( - camera_name="agentview", - height=256, - width=256, + from scipy.spatial.transform import Rotation as _R + + runtime = ctx.robot + started = runtime.executed_steps + + def _pitch_of(q): + R = _R.from_quat([q[0], q[1], q[2], q[3]]).as_matrix() + return float(np.arctan2(R[1, 2], -R[2, 2])) + + def _yaw_of(q): + R = _R.from_quat([q[0], q[1], q[2], q[3]]).as_matrix() + return float(np.arctan2(R[1, 0], R[0, 0])) + + target = np.asarray(xyz, dtype=np.float32) + for _ in range(max_steps): + cur = runtime._last_obs_eef_pos + q = runtime.env.raw_obs()["robot0_eef_quat"] + diff = target - cur + dist = float(np.linalg.norm(diff)) + p_err = ( + 0.0 + if target_pitch is None + else float((target_pitch - _pitch_of(q) + np.pi) % (2 * np.pi) - np.pi) ) - or {} - ) - if agentview_meta: - cam_meta_out = dict(agentview_meta) - cam_meta_out["projection"] = ( - "Prefer the back_project(row, col, step=NN) MCP tool; it " - "uses the 1024x1024 high-resolution world map by default. " - "Pass resolution='low' only when row/col came from the " - "256x256 calibration-frame image." + y_err = ( + 0.0 + if target_yaw is None + else float((target_yaw - _yaw_of(q) + np.pi) % (2 * np.pi) - np.pi) ) - cam_meta_out["note"] = ( - "The agentview_depth.npz observation is aligned with agentview.png. " - "agentview_policy.png uses the Pi0 orientation and must not supply " - "pixels for back-projection." + if dist < tol and abs(p_err) < ori_tol and abs(y_err) < ori_tol: + break + action = np.zeros(7, dtype=np.float32) + sd = np.clip(diff, -step_clip, step_clip) + action[:3] = np.clip(sd / action_scale, -1.0, 1.0) + action[3] = float( + np.clip(np.clip(p_err, -pitch_step, pitch_step) / 0.10, -1.0, 1.0) ) - env_state.save( - "agentview_metadata.json", - cam_meta_out, - step=step_idx, + action[5] = float( + np.clip(np.clip(y_err, -yaw_step, yaw_step) / 0.10, -1.0, 1.0) ) + action[6] = gripper + _step_env(ctx, action) + if runtime.env.terminated or runtime.env.truncated: + break + final = runtime._last_obs_eef_pos + fq = runtime.env.raw_obs()["robot0_eef_quat"] + return ToolResult( + data={ + "name": "move_pose", + "final_eef_pos": [round(float(x), 4) for x in final], + "final_dist_m": round(float(np.linalg.norm(target - final)), 4), + "final_pitch": round(_pitch_of(fq), 4), + "steps_used": runtime.executed_steps - started, + "terminated": runtime.env.terminated, + "truncated": runtime.env.truncated, + } + ) - # --- per-step RGB in the depth/K frame (vertical-flip of the raw buffer) --- - # 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) - env_state.save( - "agentview.png", - ci[::-1], - step=step_idx, - ) - except Exception as e: - logger.warning("image_cam dump failed: %s", e) - - # --- per-step metric depth (agentview), native orientation, in meters --- - try: - d = raw.get("agentview_depth") - if d is not None: - # Vertical flip to align with the camera matrices: robosuite's - # camera_utils projection M = K_exp @ inv(extrinsic) expects the - # 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 - # Pixels in agentview.png align with agentview_depth.npz and the - # per-step agentview metadata saved with the record artifacts. - d = _metric_depth(d, agentview_meta)[::-1] - env_state.save( - "agentview_depth.npz", - d.astype(np.float32), - step=step_idx, - ) - world = _world_from_depth(d, agentview_meta).astype(np.float32) - env_state.save( - "agentview_world.npz", - world, - step=step_idx, - ) - except Exception as e: - logger.warning("depth dump failed: %s", e) - - # --- per-step wrist camera (robot0_eye_in_hand), calibration frame --- - try: - wimg = raw.get("robot0_eye_in_hand_image") - if wimg is None: - logger.warning("wrist image missing from raw_obs") - else: - wimg = np.asarray(wimg) - if wimg.dtype != np.uint8: - wimg = wimg.astype(np.uint8) - env_state.save( - "wrist.png", - wimg[::-1], - step=step_idx, - ) - except Exception as e: - logger.warning("wrist image dump failed: %s", e) - - try: - wdpt = raw.get("robot0_eye_in_hand_depth") - if wdpt is None: - logger.warning("wrist depth missing from raw_obs") - else: - wdpt_arr = np.asarray(wdpt, dtype=np.float32) - height, width = wdpt_arr.shape[:2] - wmeta = primitives.env.get_camera_meta( - camera_name="robot0_eye_in_hand", - height=int(height), - width=int(width), - ) - if wmeta is None: - logger.warning("wrist camera meta missing; skipping wrist depth/world") - else: - wdpt_metric = _metric_depth(wdpt_arr, wmeta)[::-1] - env_state.save( - "wrist_depth.npz", - wdpt_metric.astype(np.float32), - step=step_idx, - ) - world_w = _world_from_depth(wdpt_metric, wmeta).astype(np.float32) - env_state.save( - "wrist_world.npz", - world_w, - step=step_idx, - ) - wmeta_out = dict(wmeta) - wmeta_out["note"] = ( - "MOVING camera: extrinsic_cam2world is for THIS step " - "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, - ) - except Exception as e: - logger.warning("wrist depth/world dump failed: %s", e) +@tool +@readonly +def view_camera_meta( + camera: Literal["agentview", "wrist"] = "agentview", + step: Annotated[int, Field(json_schema_extra={"default": -1})] = -1, + *, + ctx: ToolContext[LiberoRuntime], +) -> ToolResult: + """Read per-step camera calibration metadata from recorded artifacts. + Args: + camera: Camera metadata to read (default agentview). + step: Metadata step to use; -1 = latest. + """ + state = ctx.state try: - rgb_hi, depth_hi = primitives.env.render_camera( - camera_name="agentview", - height=1024, - width=1024, - depth=True, - ) - meta_hi = primitives.env.get_camera_meta("agentview", 1024, 1024) - if meta_hi is None: - raise RuntimeError("agentview camera metadata missing") - 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) - env_state.save( - "agentview_world_high.npz", - world_hi, - step=step_idx, - ) + 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: - logger.warning("agentview high-res dump failed: %s", e) + return ToolResult(error=f"{camera} camera metadata not found: {e}") - try: - rgb_wrist_hi, depth_wrist_hi = primitives.env.render_camera( - camera_name="robot0_eye_in_hand", - height=1024, - width=1024, - depth=True, - ) - meta_wrist_hi = primitives.env.get_camera_meta("robot0_eye_in_hand", 1024, 1024) - if meta_wrist_hi is None: - raise RuntimeError("robot0_eye_in_hand camera metadata missing") - 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) - env_state.save( - "wrist_world_high.npz", - world_wrist_hi, - step=step_idx, - ) - except Exception as e: - logger.warning("wrist high-res dump failed: %s", e) - - -# --------------------------------------------------------------------------- -# Tool schema declarations (Anthropic-shaped canonical schema) -# --------------------------------------------------------------------------- - -TOOLS_SPEC = [ - { - "name": "reset", - "description": ( - "EXPLORE MODE ONLY. Abandon the current episode and restore the " - "same initial scene. Archive the failed attempt first and state " - "which strategy lever will change in the next attempt." - ), - "input_schema": { - "type": "object", - "properties": { - "reason": { - "type": "string", - "description": ( - "Why this episode is unrecoverable and what will change." - ), - } - }, - "required": ["reason"], - }, - }, - { - "name": "view_env_state", - "description": ( - "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 " - "near the gripper, occlusions, and container/cabinet interiors." - ), - "input_schema": { - "type": "object", - "properties": { - "step": { - "type": "integer", - "default": -1, - "description": "Step number; 0 = initial, -1 = latest.", - }, - }, - }, - }, - { - "name": "move_to", - "description": ( - "Scripted EEF servo to a world-frame XYZ target via the OSC " - "controller. Holds orientation (use rotate_wrist / rotate_pitch " - "/ move_pose to reorient). gripper: -1 = open, +1 = close. NEVER " - "command a single move_to with |Δxy| > 0.30 — OSC flips IK and " - "the run corrupts; split long traversal into 2-3 mid waypoints " - "at carry z." - ), - "input_schema": { - "type": "object", - "properties": { - "xyz": { - "type": "array", - "description": "World-frame target [x, y, z] in meters", - "items": {"type": "number"}, - "minItems": 3, - "maxItems": 3, - }, - "gripper": { - "type": "number", - "description": "Gripper command: -1 open, +1 close (default -1)", - }, - "tol": { - "type": "number", - "description": "Position tolerance, m (default 0.012)", - }, - "step_clip": { - "type": "number", - "description": "Per-step Δxyz cap before action_scale, m (default 0.025)", - }, - "max_steps": { - "type": "integer", - "description": "Step budget (default 80)", - }, - "action_scale": { - "type": "number", - "description": "OSC action scale (default 0.05)", - }, - "target_yaw": { - "type": ["number", "null"], - "description": "Optional world-frame yaw target in radians", - }, - "yaw_step_clip": { - "type": "number", - "description": "Per-step yaw clip, rad (default 0.10)", - }, - }, - "required": ["xyz"], - }, - }, - { - "name": "pi0_pick", - "description": ( - "Pi0.5 closed-loop pick. Use it for the grasp; YOU then do " - "every move_to and release. Use modest max_chunks and verify " - "the grasp from EEF lift, gripper closure, and available images." - ), - "input_schema": { - "type": "object", - "properties": { - "prompt": { - "type": "string", - "description": "Pi0 prompt (e.g. 'pick up the akita black bowl').", - }, - "max_chunks": { - "type": "integer", - "description": "Action-chunk budget (default 24)", - }, - "lift_thresh": { - "type": "number", - "description": "EEF post-descent ascent threshold for success, m (default 0.05)", - }, - "gripper_closed_thresh": { - "type": "number", - "description": "Finger-separation closed threshold (default 0.06)", - }, - "gripper_open_thresh": { - "type": "number", - "description": "Minimum finger separation accepted as a held object (default 0.0)", - }, - "descent_thresh": { - "type": "number", - "description": "Required descent before lift detection, m (default 0.10)", - }, - }, - "required": ["prompt"], - }, - }, - { - "name": "pi0_doubled", - "description": ( - "Pi0.5 closed-loop contact skill for non-pick interactions " - "(e.g. stove/knob/button/short push). Returned success/task_success " - "only mirrors official termination; for intermediate contact " - "skills, success=false does not necessarily mean the contact " - "interaction failed. Inspect image/state evidence. Do not use it " - "as a general pick/place shortcut." - ), - "input_schema": { - "type": "object", - "properties": { - "prompt": { - "type": "string", - "description": "Contact-skill prompt, e.g. 'turn on the stove'.", - }, - "max_chunks": { - "type": "integer", - "description": "Action-chunk budget (default 20)", - }, - }, - "required": ["prompt"], - }, - }, - { - "name": "release", - "description": ( - "Open the gripper for up to max_steps env steps while holding " - "EEF in place. Triggers libero termination if the matching " - "On/In predicate is met." - ), - "input_schema": { - "type": "object", - "properties": { - "max_steps": { - "type": "integer", - "description": "Step budget (default 20)", - }, - }, - }, - }, - { - "name": "set_gripper", - "description": ( - "Hold the current EEF pose and drive the gripper command for " - "`steps` env steps. Use to firm up a grip mid-carry." - ), - "input_schema": { - "type": "object", - "properties": { - "gripper": { - "type": "number", - "description": "Gripper command: -1 open, +1 close (default -1)", - }, - "steps": { - "type": "integer", - "description": "Number of env steps (default 5)", - }, - }, - }, - }, - { - "name": "rotate_wrist", - "description": ( - "Rotate the wrist around the world Z-axis. Provide either " - "target_yaw (absolute) or delta_yaw (relative). Holds xyz fixed." - ), - "input_schema": { - "type": "object", - "properties": { - "target_yaw": { - "type": ["number", "null"], - "description": "Absolute world-frame yaw target, rad", - }, - "delta_yaw": { - "type": ["number", "null"], - "description": "Relative yaw delta, rad", - }, - "gripper": { - "type": "number", - "description": "Gripper command held during rotation (default +1)", - }, - "max_steps": { - "type": "integer", - "description": "Step budget (default 40)", - }, - "tol": { - "type": "number", - "description": "Yaw tolerance, rad (default 0.02)", - }, - "step_clip": { - "type": "number", - "description": "Per-step yaw clip, rad (default 0.10)", - }, - }, - }, - }, - { - "name": "rotate_pitch", - "description": ( - "Tilt the gripper around the world X-axis. Provide either " - "target_pitch (absolute) or delta_pitch (relative). Holds xyz " - "and yaw fixed. Use before threading the gripper into a narrow " - "opening whose front face normal is along world ±y (e.g. " - "microwave cavity)." - ), - "input_schema": { - "type": "object", - "properties": { - "target_pitch": { - "type": ["number", "null"], - "description": "Absolute world-frame pitch target, rad", - }, - "delta_pitch": { - "type": ["number", "null"], - "description": "Relative pitch delta, rad", - }, - "gripper": { - "type": "number", - "description": "Gripper command held during rotation (default +1)", - }, - "max_steps": { - "type": "integer", - "description": "Step budget (default 40)", - }, - "tol": { - "type": "number", - "description": "Pitch tolerance, rad (default 0.02)", - }, - "step_clip": { - "type": "number", - "description": "Per-step pitch clip, rad (default 0.10)", - }, - }, - }, - }, - { - "name": "move_pose", - "description": ( - "Servo position AND orientation (pitch + yaw) SIMULTANEOUSLY. " - "Unlike move_to (holds orientation) + rotate_pitch (holds xyz), " - "this co-varies xyz and wrist tilt every env.step. Use to thread " - "cabinet-front / low-shelf poses where a decoupled position " - "servo drives the wrist into an IK singularity and stalls." - ), - "input_schema": { - "type": "object", - "properties": { - "xyz": { - "type": "array", - "description": "World-frame target [x, y, z] in meters", - "items": {"type": "number"}, - "minItems": 3, - "maxItems": 3, - }, - "target_pitch": { - "type": ["number", "null"], - "description": "Absolute pitch target, rad", - }, - "target_yaw": { - "type": ["number", "null"], - "description": "Absolute yaw target, rad", - }, - "gripper": { - "type": "number", - "description": "Gripper command held during the move (default -1)", - }, - "step_clip": { - "type": "number", - "description": "Per-step Δxyz cap, m (default 0.02)", - }, - "pitch_step": { - "type": "number", - "description": "Per-step pitch clip, rad (default 0.08)", - }, - "yaw_step": { - "type": "number", - "description": "Per-step yaw clip, rad (default 0.08)", - }, - "tol": { - "type": "number", - "description": "Position tolerance, m (default 0.012)", - }, - "ori_tol": { - "type": "number", - "description": "Orientation tolerance, rad (default 0.05)", - }, - "action_scale": { - "type": "number", - "description": "OSC action scale (default 0.05)", - }, - "max_steps": { - "type": "integer", - "description": "Step budget (default 150)", - }, - }, - "required": ["xyz"], - }, - }, - { - "name": "view_camera_meta", - "description": ( - "Read per-step camera calibration metadata from recorded artifacts." - ), - "input_schema": { - "type": "object", - "properties": { - "camera": { - "type": "string", - "enum": ["agentview", "wrist"], - "description": "Camera metadata to read (default agentview).", - }, - "step": { - "type": "integer", - "default": -1, - "description": "Metadata step to use; -1 = latest.", - }, - }, - }, - }, - { - "name": "segment", - "description": ( - "SAM3 visual segmentation over an existing run artifact. It never " - "renders a new camera view. Provide exactly one text prompt or " - "single positive point. A successful top-ranked mask is projected " - "through the matching world map to produce world_xyz." - ), - "input_schema": { - "type": "object", - "properties": { - "prompt": { - "type": "string", - "description": "Object/text prompt to segment.", - }, - "camera": { - "type": "string", - "enum": ["agentview", "wrist"], - "description": "Artifact camera to use (default agentview).", - }, - "step": { - "type": "integer", - "default": -1, - "description": "Step to segment; -1 = latest.", - }, - "point": { - "type": ["array", "null"], - "description": ( - "Optional single positive point as [row, col]. " - "Mutually exclusive with prompt." - ), - "items": {"type": "integer"}, - "minItems": 2, - "maxItems": 2, - }, - "min_score": { - "type": "number", - "description": "Minimum accepted mask score (default 0.2).", - }, - }, - }, - }, - { - "name": "back_project", - "description": ( - "Back-project a pixel (row, col) to a world XYZ point using the " - "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 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 " - "256 image. The pixel coordinates must come " - "from the same camera and resolution requested here. Use " - "camera='agentview' for global tabletop layout and object " - "locations; use camera='wrist' for close-range details near the " - "gripper, occlusions, and container/cabinet interiors. " - "Sample several pixels on the object and median their xy for " - "robustness.\n\n" - "REGION MODE: pass row_range=[r0,r1] and col_range=[c0,c1] instead " - "of row/col to get the midpoint of world xy over that pixel window, " - "with an optional world-z band (z_min, z_max). Use it for the " - "center of a container cavity or flat region, where a single-pixel " - "or mask-median estimate is biased toward an edge/rim." - ), - "input_schema": { - "type": "object", - "properties": { - "row": { - "type": ["integer", "null"], - "description": "Pixel row (0=top) in the selected resolution image.", - }, - "col": { - "type": ["integer", "null"], - "description": "Pixel column (0=left) in the selected resolution image.", - }, - "step": { - "type": "integer", - "default": -1, - "description": "Depth/world-map step; 0 = initial, -1 = latest.", - }, - "camera": { - "type": "string", - "enum": ["agentview", "wrist"], - "description": "Camera to back-project from (default agentview).", - }, - "resolution": { - "type": "string", - "enum": ["high", "low"], - "description": ( - "Coordinate system for row/col (default high). " - "Use low only when row/col came from the " - "embedded/standard 256 image." - ), - }, - "row_range": { - "type": ["array", "null"], - "items": {"type": "integer"}, - "description": "Region mode: [r0, r1] pixel row window. Requires col_range.", - }, - "col_range": { - "type": ["array", "null"], - "items": {"type": "integer"}, - "description": "Region mode: [c0, c1] pixel col window. Requires row_range.", - }, - "z_min": { - "type": ["number", "null"], - "description": "Region mode: keep only pixels with world z >= z_min.", - }, - "z_max": { - "type": ["number", "null"], - "description": "Region mode: keep only pixels with world z <= z_max.", - }, - }, - }, - }, -] + if camera == "agentview": + return ToolResult(data={"camera": "agentview", "camera_meta": meta}) + return ToolResult( + data={"camera": "wrist", "step": record.step_idx, "camera_meta": meta} + ) +@tool @readonly -def view_env_state(step: int = -1, *, state: EnvState) -> dict: +def segment( + prompt: str = "", + camera: Literal["agentview", "wrist"] = "agentview", + step: Annotated[int, Field(json_schema_extra={"default": -1})] = -1, + point: Annotated[list[int] | None, Field(min_length=2, max_length=2)] = None, + min_score: float = 0.2, + *, + ctx: ToolContext[LiberoRuntime], +) -> ToolResult: + """SAM3 visual segmentation over an existing run artifact. Provide exactly one text prompt or single positive point. A successful top-ranked mask is projected through the matching world map to produce world_xyz. + + Args: + prompt: Object/text prompt to segment. + camera: Artifact camera to use (default agentview). + step: Step to segment; -1 = latest. + point: Optional single positive point as [row, col]. Mutually exclusive with prompt. + min_score: Minimum accepted mask score (default 0.2). + """ + runtime = ctx.robot + state = ctx.state try: record = state.get(step) except Exception as exc: - return {"error": f"state step not available: {exc}"} - + return ToolResult(error=f"state step not available: {exc}") nn = record.step_idx - extras = record.extras - out: dict = { - "step": nn, - "terminated": record.terminated, - "truncated": record.truncated, - "state": record.state, - "artifacts": sorted(record.artifacts), - } - out["task_language"] = extras.get("task_language") - out["log"] = { - "command": record.command, - "result": record.result, - "elapsed_s": record.elapsed_s, - } - 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")), - ): - 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( - 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 = [ + + prompt = prompt.strip() + has_prompt = bool(prompt) + has_point = point is not None + if has_prompt == has_point: + return ToolResult(error="segment needs exactly one of prompt or point") + artifact_pairs = [ (f"{camera}_high.png", f"{camera}_world_high.npz"), (f"{camera}.png", f"{camera}_world.npz"), ] - for image_name, world_name in pairs: + for image_name, world_name in artifact_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) + and state.exists(image_name, step=nn) + and state.exists(world_name, step=nn) ): - return image_name, world_name, pairs - return None, None, pairs - - -def _next_segment_index(record: StepRecord) -> int: - idx = 0 - while f"segment_{idx:02d}.json" in record.artifacts: - idx += 1 - return idx + break + else: + return ToolResult( + data={ + "step": nn, + "camera": camera, + "checked_artifacts": [ + name for image, world in artifact_pairs for name in (image, world) + ], + }, + error="complete segment artifacts not found", + ) + try: + data = runtime._sam3_client.segment( + state.load_bytes(image_name, step=nn), + text_prompt=prompt if has_prompt else None, + point=point, + min_score=min_score, + ) + except ValueError as e: + return ToolResult( + data={"step": nn, "camera": camera, "image_artifact": image_name}, + error=str(e), + ) + except Exception as e: + return ToolResult( + data={ + "step": nn, + "camera": camera, + "image_artifact": image_name, + "fallback": "Use manual visual localization and back_project.", + }, + error=f"segmentation service call failed: {e}", + ) -def _mask_to_world( - mask: np.ndarray, world_map: np.ndarray, min_valid: int = 10 -) -> dict: - if world_map.ndim != 3 or world_map.shape[2] < 3: - return { + segment_index = 0 + while f"segment_{segment_index:02d}.json" in record.artifacts: + segment_index += 1 + 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: + try: + world_map = state.load(world_name, step=nn) + except Exception as exc: + world_result = { + "world_xyz": None, + "world_error": f"world map artifact not available: {exc}", + "expected_world_artifact": world_name, + } + else: + world_result = _mask_to_world(mask, world_map) + world_result["world_artifact"] = world_name + image = state.load(image_name, step=nn) + if image.ndim == 3 and image.shape[:2] == mask.shape: + overlay = image.copy() + red = np.zeros(overlay.shape[-1], dtype=np.float32) + red[0] = 255 + overlay[mask] = ( + 0.55 * overlay[mask].astype(np.float32) + 0.45 * red + ).astype(np.uint8) + saved_overlay = state.save(overlay_name, overlay, step=nn) + else: + world_result = { "world_xyz": None, - "world_error": f"invalid world map shape: {tuple(world_map.shape)}", - "n_pixels": int(mask.sum()), - "n_valid": 0, - "mask_resized_to_world_shape": False, + "world_error": data.reason or "segmentation did not find a mask", } - if mask.shape != world_map.shape[:2]: - return { - "world_xyz": None, - "world_error": ( - f"mask/world shape mismatch: mask={tuple(mask.shape)}, " - f"world={tuple(world_map.shape[:2])}" - ), - "n_pixels": int(mask.sum()), - "n_valid": 0, - "mask_resized_to_world_shape": False, - } - - ys, xs = np.where(mask) - if ys.size == 0: - return {"world_xyz": None, "world_error": "empty mask"} + segment_blob = { + "found": data.found, + "mode": "text" if has_prompt else "point", + "camera": camera, + "source_step": nn, + "segment_index": segment_index, + "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, + "mask_shape": list(data.mask_shape) if data.mask_shape else None, + } + if has_prompt: + segment_blob["prompt"] = prompt + else: + segment_blob["point"] = point + if not data.found: + segment_blob["error"] = data.reason or "SAM3 found no mask" + segment_blob.update(world_result) + saved_segment = state.save( + segment_name, + segment_blob, + step=nn, + ) - pts = world_map[ys, xs].astype(np.float64) - valid = np.isfinite(pts).all(axis=1) & (np.abs(pts).sum(axis=1) > 1e-6) - pts = pts[valid] result = { - "centroid_pixel": [ - int(round(float(np.median(xs)))), - int(round(float(np.median(ys)))), - ], - "n_pixels": int(mask.sum()), - "n_valid": int(pts.shape[0]), - "mask_resized_to_world_shape": False, + "found": data.found, + "step": nn, + "camera": camera, + "image_artifact": image_name, + "score": segment_blob["score"], + "box": segment_blob["box"], + "world_xyz": segment_blob["world_xyz"], } - if pts.shape[0] < min_valid: - result.update( - { - "world_xyz": None, - "world_error": f"too few valid depth pixels ({int(pts.shape[0])})", - } - ) - return result - - result["world_xyz"] = [ - round(float(np.median(pts[:, 0])), 4), - round(float(np.median(pts[:, 1])), 4), - round(float(np.median(pts[:, 2])), 4), - ] - return result - - -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 - - -@readonly -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')"} - - try: - 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": record.step_idx, "camera_meta": meta} - - + if data.found: + result["world_error"] = segment_blob.get("world_error") + error = None + if saved_segment is None: + error = f"failed to persist segment artifact {segment_name}" + result["attempted_segment_artifact"] = segment_name + if "error" in segment_blob: + error += "\n" + json.dumps({"segmentation_error": segment_blob["error"]}) + else: + result["segment_artifact"] = saved_segment + result["segment_path"] = str(state.artifact_path(saved_segment, step=nn)) + error = segment_blob.get("error") + if error is not None: + result["fallback"] = "Use manual visual localization and back_project." + images: list[bytes] = [] + if saved_overlay is not None: + result["overlay_artifact"] = saved_overlay + result["overlay_path"] = str(state.artifact_path(saved_overlay, step=nn)) + images.append(state.load_bytes(saved_overlay, step=nn)) + return ToolResult(data=result, error=error, images=images) + + +@tool @readonly def back_project( row: int | None = None, col: int | None = None, - step: int = -1, - camera: str = "agentview", - resolution: str = "high", - row_range: list | None = None, - col_range: list | None = None, + step: Annotated[int, Field(json_schema_extra={"default": -1})] = -1, + camera: Literal["agentview", "wrist"] = "agentview", + resolution: Literal["high", "low"] = "high", + row_range: list[int] | None = None, + col_range: list[int] | 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"): - return {"error": f"bad camera '{camera}' (use 'agentview' or 'wrist')"} - if resolution not in ("high", "low"): - return {"error": f"bad resolution '{resolution}' (use 'high' or 'low')"} - + ctx: ToolContext[LiberoRuntime], +) -> ToolResult: + """Back-project a pixel (row, col) to a world XYZ point using the selected camera's precomputed world map. Row 0 = top of image, col 0 = left. Returns world_xyz in meters. + + USE THIS to find where an object is in the world — look at 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 256 image. The pixel coordinates must come from the same camera and resolution requested here. Use camera='agentview' for global tabletop layout and object locations; use camera='wrist' for close-range details near the gripper, occlusions, and container/cabinet interiors. Sample several pixels on the object and median their xy for robustness. + + REGION MODE: pass row_range=[r0,r1] and col_range=[c0,c1] instead of row/col to get the midpoint of world xy over that pixel window, with an optional world-z band (z_min, z_max). Use it for the center of a container cavity or flat region, where a single-pixel or mask-median estimate is biased toward an edge/rim. + + Args: + row: Pixel row (0=top) in the selected resolution image. + col: Pixel column (0=left) in the selected resolution image. + step: Depth/world-map step; 0 = initial, -1 = latest. + camera: Camera to back-project from (default agentview). + resolution: Coordinate system for row/col (default high). Use low only when row/col came from the embedded/standard 256 image. + row_range: Region mode: [r0, r1] pixel row window. Requires col_range. + col_range: Region mode: [c0, c1] pixel col window. Requires row_range. + z_min: Region mode: keep only pixels with world z >= z_min. + z_max: Region mode: keep only pixels with world z <= z_max. + """ + state = ctx.state region_mode = row_range is not None or col_range is not None if not region_mode and (row is None or col is None): - return { - "error": ( - "provide either (row, col) for a single pixel, or " - "row_range=[r0,r1] and col_range=[c0,c1] for a region center" - ) - } + return ToolResult( + error="provide either (row, col) for a single pixel, or row_range=[r0,r1] and col_range=[c0,c1] for a region center" + ) try: record = state.get(step) except Exception as e: - return {"error": f"state step not available: {e}"} + return ToolResult(error=f"state step not available: {e}") nn = record.step_idx hi_artifact = f"{camera}_world_high.npz" low_artifact = f"{camera}_world.npz" source_artifact = hi_artifact if resolution == "high" else low_artifact if source_artifact not in record.artifacts: - return { - "error": ( - f"{camera} {resolution}-resolution world map not recorded for step {nn}" - ) - } + return ToolResult( + error=f"{camera} {resolution}-resolution world map not recorded for step {nn}" + ) try: - world_map = state.load(str(source_artifact), step=nn) + world_map = state.load(source_artifact, step=nn) except Exception as e: - return { - "error": ( - f"{camera} {resolution}-resolution artifact not found " - f"for step {nn}: {e}" - ) - } + return ToolResult( + error=f"{camera} {resolution}-resolution artifact not found for step {nn}: {e}" + ) height, width = world_map.shape[:2] if region_mode: if row_range is None or col_range is None: - return { - "error": "region mode needs BOTH row_range=[r0,r1] and col_range=[c0,c1]" - } + return ToolResult( + error="region mode needs BOTH row_range=[r0,r1] and col_range=[c0,c1]" + ) try: - r0, r1 = int(row_range[0]), int(row_range[1]) - c0, c1 = int(col_range[0]), int(col_range[1]) - except Exception: - return {"error": "row_range/col_range must each be [min, max] integers"} + r0, r1 = row_range[0], row_range[1] + c0, c1 = col_range[0], col_range[1] + except IndexError: + return ToolResult( + error="row_range/col_range must each be [min, max] integers" + ) r0, r1 = sorted((max(0, r0), min(height, r1))) c0, c1 = sorted((max(0, c0), min(width, c1))) if r1 <= r0 or c1 <= c0: - return { - "error": ( - f"empty region after clamping to image {height}x{width}: " - f"rows [{r0},{r1}] cols [{c0},{c1}]" - ) - } + return ToolResult( + error=f"empty region after clamping to image {height}x{width}: rows [{r0},{r1}] cols [{c0},{c1}]" + ) window = ( world_map[r0:r1, c0:c1].reshape(-1, world_map.shape[2]).astype(np.float64) ) @@ -1904,44 +977,41 @@ def back_project( if z_max is not None: pts = pts[pts[:, 2] <= float(z_max)] if pts.shape[0] < 8: - return { - "error": ( - f"too few valid pixels in region after z-filter " - f"({int(pts.shape[0])}); widen the window or the z band" - ), - "n_valid_before_zfilter": n_total, - } + return ToolResult( + data={"n_valid_before_zfilter": n_total}, + error=f"too few valid pixels in region after z-filter ({int(pts.shape[0])}); widen the window or the z band", + ) xs, ys, zs = pts[:, 0], pts[:, 1], pts[:, 2] center = [ round(float((xs.min() + xs.max()) / 2.0), 4), round(float((ys.min() + ys.max()) / 2.0), 4), round(float(np.median(zs)), 4), ] - return { - "camera": camera, - "resolution": resolution, - "mode": "region", - "row_range": [r0, r1], - "col_range": [c0, c1], - "z_band": [z_min, z_max], - "center_xyz": center, - "median_xyz": [ - round(float(np.median(xs)), 4), - round(float(np.median(ys)), 4), - round(float(np.median(zs)), 4), - ], - "n_valid": int(pts.shape[0]), - "step": nn, - "image_size": [height, width], - "source_artifact": source_artifact, - } + return ToolResult( + data={ + "camera": camera, + "resolution": resolution, + "mode": "region", + "row_range": [r0, r1], + "col_range": [c0, c1], + "z_band": [z_min, z_max], + "center_xyz": center, + "median_xyz": [ + round(float(np.median(xs)), 4), + round(float(np.median(ys)), 4), + round(float(np.median(zs)), 4), + ], + "n_valid": int(pts.shape[0]), + "step": nn, + "image_size": [height, width], + "source_artifact": source_artifact, + } + ) if row < 0 or row >= height or col < 0 or col >= width: - return { - "error": ( - f"pixel ({row},{col}) out of bounds; {camera} image is {height}x{width}" - ) - } + return ToolResult( + error=f"pixel ({row},{col}) out of bounds; {camera} image is {height}x{width}" + ) depth_m = None if source_artifact == low_artifact: @@ -1953,21 +1023,18 @@ def back_project( if depth.ndim == 3: depth = depth[..., 0] except Exception as e: - return {"error": f"{camera} depth not found for step {nn}: {e}"} + return ToolResult(error=f"{camera} depth not found for step {nn}: {e}") depth_m = float(depth[row, col]) if not np.isfinite(depth_m) or depth_m <= 0 or depth_m > 10: - return { - "error": ( - f"invalid {camera} depth {depth_m:.3f}m at pixel " - f"({row},{col}); pick a different pixel" - ) - } + return ToolResult( + error=f"invalid {camera} depth {depth_m:.3f}m at pixel ({row},{col}); pick a different pixel" + ) world_xyz_raw = world_map[row, col] if ( not np.isfinite(world_xyz_raw).all() or float(np.abs(world_xyz_raw[:3]).sum()) <= 1e-6 ): - return {"error": f"invalid {camera} world xyz at pixel ({row},{col})"} + return ToolResult(error=f"invalid {camera} world xyz at pixel ({row},{col})") world_xyz = [round(float(v), 4) for v in world_xyz_raw[:3]] out = { @@ -1981,4 +1048,79 @@ def back_project( } if depth_m is not None: out["depth_m"] = round(depth_m, 4) - return out + return ToolResult(data=out) + + +def _mask_to_world( + mask: np.ndarray, world_map: np.ndarray, min_valid: int = 10 +) -> dict: + if world_map.ndim != 3 or world_map.shape[2] < 3: + return { + "world_xyz": None, + "world_error": f"invalid world map shape: {tuple(world_map.shape)}", + "n_pixels": int(mask.sum()), + "n_valid": 0, + "mask_resized_to_world_shape": False, + } + + if mask.shape != world_map.shape[:2]: + return { + "world_xyz": None, + "world_error": ( + f"mask/world shape mismatch: mask={tuple(mask.shape)}, " + f"world={tuple(world_map.shape[:2])}" + ), + "n_pixels": int(mask.sum()), + "n_valid": 0, + "mask_resized_to_world_shape": False, + } + + ys, xs = np.where(mask) + if ys.size == 0: + return {"world_xyz": None, "world_error": "empty mask"} + + pts = world_map[ys, xs].astype(np.float64) + valid = np.isfinite(pts).all(axis=1) & (np.abs(pts).sum(axis=1) > 1e-6) + pts = pts[valid] + result = { + "centroid_pixel": [ + int(round(float(np.median(xs)))), + int(round(float(np.median(ys)))), + ], + "n_pixels": int(mask.sum()), + "n_valid": int(pts.shape[0]), + "mask_resized_to_world_shape": False, + } + if pts.shape[0] < min_valid: + result.update( + { + "world_xyz": None, + "world_error": f"too few valid depth pixels ({int(pts.shape[0])})", + } + ) + return result + + result["world_xyz"] = [ + round(float(np.median(pts[:, 0])), 4), + round(float(np.median(pts[:, 1])), 4), + round(float(np.median(pts[:, 2])), 4), + ] + return result + + +LIBERO_TOOLS = ( + finish, + reset, + view_env_state, + move_to, + pi0_pick, + pi0_doubled, + release, + set_gripper, + rotate_wrist, + rotate_pitch, + move_pose, + view_camera_meta, + segment, + back_project, +) diff --git a/robots/robocasa/primitives.py b/robots/robocasa/primitives.py deleted file mode 100644 index 02e728387..000000000 --- a/robots/robocasa/primitives.py +++ /dev/null @@ -1,553 +0,0 @@ -# Copyright 2026 The RPent Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""RoboCasaPrimitives — action primitives, perception, and VLA execution for RoboCasa.""" - -import os - -import numpy as np - -from robots.robocasa.rldx_skill import RLDXSkill - -OSC_POS_SCALE = 0.05 # action 1.0 -> 0.05 m target delta -OSC_ROT_SCALE = 0.5 # action 1.0 -> 0.5 rad - - -class RoboCasaPrimitives: - def __init__(self, env_client, workdir, hi_res, vla_client, check_cancelled=None): - self.env = env_client - self.workdir = workdir - self.hi_res = hi_res - # Cancellation checkpoint callback (e.g. Toolkit.raise_if_cancelled), - # invoked before long-running env.step loops so an interrupted tool - # operation stops promptly instead of draining the whole episode. - self._check_cancelled = check_cancelled - # World maps (per-pixel world xyz) are dumped EVERY step — they are how the agent - # localizes every object, so they are never optional. Disk is bounded by keeping only - # the last `_keep_heavy` steps' heavy npy on disk (pruned in dump_state) + each - # launcher's per-cell workdir cleanup. - self._keep_heavy = int(os.environ.get("RLDX_KEEP_HEAVY_NPY", "25")) - # reset (restart the episode) is ONLY legitimate in EXPLORE mode (reset-based recipe - # search). It is FORBIDDEN in multi-seed / matched-scene evaluation (a give-up-and- - # restart). Gated by RLDX_ALLOW_RESET (default 0 = off); explore runs opt in. - self._allow_reset = bool(int(os.environ.get("RLDX_ALLOW_RESET", "0"))) - os.makedirs(workdir, exist_ok=True) - self.env.reset() - self._pos_jac = None # 3x3 action(arm xyz) -> world dpos - self._fwd_offset = None # world_forward_heading = base_yaw + offset - self._cam_meta_cache = {} - self._rldx = RLDXSkill( - self.env, vla_client=vla_client, check_cancelled=check_cancelled - ) - # "mid-call" desync guard: True whenever a NON-VLA primitive (move/navigate/ - # manual grasp) or a reset has stepped the env since the last rldx_skill call. - # The next rldx_skill then reseeds its per-sim-step frame history (else the VLA - # sees a pre-manual history stitched onto a post-manual current frame -> OOD). - # Consecutive VLA calls with NO manual step in between keep history continuity. - self._vla_desync = True - # Recording (off by default — zero overhead when not recording) - self._recording = False - self._frames = [] - - # ---- recording methods ---- - def start_recording(self): - self._recording = True - self._frames = [] - - def record_frame(self, img=None): - """Snapshot the current agentview to the frame buffer, if recording.""" - if img is None: - img = self.env.render_camera( - camera_name="agentview", - height=256, - width=256, - depth=False, - ) - img = np.asarray(img, dtype=np.uint8) - img = np.ascontiguousarray(img) - self._frames.append(img) - - def recorded_frame_count(self) -> int: - return len(self._frames) - - def frame_slice(self, start: int) -> list[np.ndarray]: - return list(self._frames[int(start) :]) - - def stop_recording(self) -> list[np.ndarray]: - frames = list(self._frames) - self._recording = False - self._frames = [] - return frames - - # ---- action helpers ---- - def _zero(self, base_mode=-1.0): - a = np.zeros(12) - a[11] = base_mode - return a - - def _hold_gripper_val(self, g): - # +1 close/hold, -1 open - return float(np.clip(g, -1, 1)) - - def _resolve_grip(self, gripper, target_q): - """Return the a[6] gripper command for a motion step. - gripper="hold"/None (DEFAULT for moves) -> SERVO the fingers back to `target_q`, - the width they had when the motion began. This is the carry-safe hold: the - gripper action is a CLOSE-VELOCITY command, so a sustained +1 keeps driving the - fingers shut and SQUEEZES a small object OUT (verified: bread qpos 0.0376 -> 0.0005 - during a +1 carry). Servoing to the grasped width holds the object without crushing - it and without letting it drift open. A numeric gripper (+1 close / -1 open) is an - EXPLICIT override and passes through unchanged.""" - if isinstance(gripper, str) or gripper is None: - cur = float(self.env.gripper_qpos[0]) - return float( - np.clip(60.0 * (cur - target_q), -1.0, 1.0) - ) # +a[6] closes (qpos↓) - return self._hold_gripper_val(gripper) - - def _step_arm(self, dpos=(0, 0, 0), drot=(0, 0, 0), gripper=-1.0, n=1): - a = self._zero(base_mode=-1.0) - a[0:3] = np.clip(np.asarray(dpos) / OSC_POS_SCALE, -1, 1) - a[3:6] = np.clip(np.asarray(drot) / OSC_ROT_SCALE, -1, 1) - a[6] = self._hold_gripper_val(gripper) - for _ in range(n): - if self._check_cancelled is not None: - self._check_cancelled() - self.env.step(a) - if self._recording: - self.record_frame() - return self.env.eef_pos - - # ---- Phase 2: online jacobian + move_to ---- - def _calibrate_pos_jacobian(self, gripper=-1.0): - """Probe 3 unit arm-xyz actions, measure world dpos -> 3x3 jacobian J s.t. - world_dpos ~= J @ action_xyz. move_to inverts J to map desired world delta.""" - cols = [] - for axis in range(3): - p0 = self.env.eef_pos.copy() - a = self._zero() - a[axis] = 0.4 - a[6] = gripper - for _ in range(3): - if self._check_cancelled is not None: - self._check_cancelled() - self.env.step(a) - if self._recording: - self.record_frame() - d = (self.env.eef_pos - p0) / (0.4 * 3) # world dpos per unit action - cols.append(d) - # settle back is not needed (closed-loop re-reads); keep going - self._pos_jac = np.stack(cols, axis=1) # 3x3: world_dpos = J @ a_xyz - return self._pos_jac - - def move_to(self, xyz, gripper="hold", step_clip=0.02, max_steps=200, tol=0.012): - """Closed-loop OSC servo of the eef to a WORLD xyz target. - gripper="hold" (DEFAULT) maintains the CURRENT finger width — carry a grasped - object WITHOUT crushing it (a sustained +1 squeezes small objects out) or letting - it drop. Pass +1 to actively CLOSE, -1 to actively OPEN.""" - self._vla_desync = True - target = np.asarray(xyz, dtype=np.float64) - target_q = float(self.env.gripper_qpos[0]) # finger width to hold - if self._pos_jac is None: - self._calibrate_pos_jacobian(gripper=self._resolve_grip(gripper, target_q)) - Jinv = np.linalg.pinv(self._pos_jac) - for i in range(max_steps): - cur = self.env.eef_pos - err = target - cur - dist = float(np.linalg.norm(err)) - if dist < tol: - return { - "ok": True, - "steps": i, - "final_dist": dist, - "eef": cur.tolist(), - "gripper_qpos": round(float(self.env.gripper_qpos[0]), 4), - } - step_world = err if dist <= step_clip else err / dist * step_clip - a_xyz = np.clip(Jinv @ step_world, -1, 1) - a = self._zero() - a[0:3] = a_xyz - a[6] = self._resolve_grip(gripper, target_q) - if self._check_cancelled is not None: - self._check_cancelled() - self.env.step(a) - if self._recording: - self.record_frame() - cur = self.env.eef_pos - return { - "ok": False, - "steps": max_steps, - "final_dist": float(np.linalg.norm(target - cur)), - "eef": cur.tolist(), - "gripper_qpos": round(float(self.env.gripper_qpos[0]), 4), - } - - def move_delta(self, dxyz, gripper="hold", step_clip=0.02, max_steps=80): - self._vla_desync = True - return self.move_to( - self.env.eef_pos + np.asarray(dxyz), gripper, step_clip, max_steps - ) - - def rotate_pitch(self, target_pitch=0.6, gripper=1, n=12): - """Tilt the wrist forward (axis-angle about control-x). Reuses arm drot.""" - self._vla_desync = True - per = float(np.clip(target_pitch, -1.5, 1.5)) / n - for _ in range(n): - self._step_arm(drot=(per, 0, 0), gripper=gripper, n=1) - return {"ok": True, "eef": self.env.eef_pos.tolist()} - - def set_gripper(self, gripper=1, steps=10): - self._vla_desync = True - g = self._hold_gripper_val(gripper) - a = self._zero() - a[6] = g - for _ in range(steps): - if self._check_cancelled is not None: - self._check_cancelled() - self.env.step(a) - if self._recording: - self.record_frame() - return {"ok": True, "gripper_qpos": self.env.gripper_qpos.tolist()} - - def release(self, steps=10): - return self.set_gripper(-1.0, steps=steps) - - # ---- Phase 4: scripted grasp ---- - def scripted_grasp(self, xyz, approach_z=0.10, grasp_z_offset=0.0, step_clip=0.02): - """Open -> hover above target -> descend -> close -> lift. Coarse; replace - with RLDX closed-loop grasp for hard objects.""" - t = np.asarray(xyz, dtype=np.float64) - self.set_gripper(-1.0, steps=4) - r = self.move_to(t + [0, 0, approach_z], gripper=-1.0, step_clip=step_clip) - if not r["ok"]: - return {**r, "stage": "approach"} - r = self.move_to( - t + [0, 0, grasp_z_offset], gripper=-1.0, step_clip=0.012, tol=0.01 - ) - if not r["ok"]: - return {**r, "stage": "descent"} - self.set_gripper(+1.0, steps=14) - r = self.move_to(t + [0, 0, approach_z + 0.05], gripper="hold", step_clip=0.015) - if not r["ok"]: - return {**r, "stage": "lift"} - return { - "ok": True, - "gripper_qpos": self.env.gripper_qpos.tolist(), - "eef": self.env.eef_pos.tolist(), - } - - # ---- Phase 3: navigation ---- - def _base_pose(self): - o = self.env.current_raw_obs - return ( - np.asarray(o["robot0_base_pos"], dtype=np.float64), - np.asarray(o["robot0_base_quat"], dtype=np.float64), - ) - - def move_base(self, forward=0, lateral=0, turn=0, steps=10, gripper="hold"): - """Raw base velocity command (robot-local: +fwd, +lateral, +turn yaw). - gripper="hold" (DEFAULT) maintains the finger width while driving (carry-safe).""" - self._vla_desync = True - target_q = float(self.env.gripper_qpos[0]) - a = self._zero(base_mode=1.0) - a[7:10] = [ - np.clip(forward, -1, 1), - np.clip(lateral, -1, 1), - np.clip(turn, -1, 1), - ] - bp0, _ = self._base_pose() - for _ in range(steps): - a[6] = self._resolve_grip(gripper, target_q) - if self._check_cancelled is not None: - self._check_cancelled() - self.env.step(a) - if self._recording: - self.record_frame() - bp1, _ = self._base_pose() - return { - "ok": True, - "base_moved": (bp1 - bp0).tolist(), - "base_pos": bp1.tolist(), - } - - def _yaw(self): - from scipy.spatial.transform import Rotation as R - - return float( - R.from_quat( - np.asarray(self.env.current_raw_obs["robot0_base_quat"]) - ).as_euler("xyz")[2] - ) - - def _calibrate_forward(self, gripper=1.0): - """Drive forward briefly, measure the WORLD direction the base actually goes, - so navigate_to can steer regardless of the base->world frame offset.""" - p0, _ = self._base_pose() - y0 = self._yaw() - a = self._zero(base_mode=1.0) - a[6] = self._hold_gripper_val(gripper) - a[7] = 1.0 - for _ in range(6): - if self._check_cancelled is not None: - self._check_cancelled() - self.env.step(a) - if self._recording: - self.record_frame() - p1, _ = self._base_pose() - disp = (p1 - p0)[:2] - if np.linalg.norm(disp) > 0.005: - self._fwd_offset = np.arctan2(disp[1], disp[0]) - y0 - else: - self._fwd_offset = 0.0 - return self._fwd_offset - - def navigate_to(self, xy, tol=0.20, max_steps=300, gripper="hold"): - """Drive the mobile base toward a WORLD (x,y) target. Online-calibrates the - base forward->world heading, then turns to face + drives forward at full - speed (closed-loop). Holds the arm (base_mode>0). Base ~2.3mm/step. - gripper="hold" (DEFAULT) maintains the finger width while driving (carry-safe).""" - self._vla_desync = True - target = np.asarray(xy[:2], dtype=np.float64) - target_q = float(self.env.gripper_qpos[0]) - if self._fwd_offset is None: - self._calibrate_forward(self._resolve_grip(gripper, target_q)) - start = self._base_pose()[0][:2].copy() - for i in range(max_steps): - bp, _ = self._base_pose() - to = target - bp[:2] - dist = float(np.linalg.norm(to)) - if dist < tol: - self._pos_jac = None # base moved -> recalibrate arm - moved = float(np.linalg.norm(bp[:2] - start)) - return { - "ok": True, - "steps": i, - "final_dist": dist, - "moved": moved, - "start_pos": start.tolist(), - "base_pos": bp.tolist(), - } - world_dir = np.arctan2(to[1], to[0]) - cur_forward = self._yaw() + self._fwd_offset - dyaw = (world_dir - cur_forward + np.pi) % (2 * np.pi) - np.pi - a = self._zero(base_mode=1.0) - a[6] = self._resolve_grip(gripper, target_q) - if abs(dyaw) > 0.30: # turn to face the target - a[9] = float(np.sign(dyaw)) - else: # drive forward + small steer - a[7] = 1.0 - a[9] = float(np.clip(dyaw * 1.5, -0.4, 0.4)) - if self._check_cancelled is not None: - self._check_cancelled() - self.env.step(a) - if self._recording: - self.record_frame() - bp, _ = self._base_pose() - self._pos_jac = None - moved = float(np.linalg.norm(bp[:2] - start)) - # stuck = ran out of steps having barely moved (rammed a fixture, no path-planning) - return { - "ok": False, - "steps": max_steps, - "final_dist": float(np.linalg.norm(target - bp[:2])), - "moved": moved, - "stuck": moved < 0.12, - "start_pos": start.tolist(), - "base_pos": bp.tolist(), - } - - def dump_success_criteria(self): - """Return this task's EXACT success condition text (the env's _check_success + - the helper/fixture predicates it calls), so the agent knows WHAT counts as done. - This is the success LOGIC (conditions on named objects/fixtures + thresholds) — - it does NOT reveal GT object COORDINATES (the agent still localizes every object - from perception). The caller writes it to the run's EnvState.""" - return self.env.get_success_criteria_text() - - def task_progress(self): - """NUMERIC progress signal toward this task's success — generic for ALL tasks. - The state json only exposes the final `success` bool; an agent driving a task - whose success is a COUNTER/THRESHOLD (washed_time>=25, success_time>=5) or a set - of sub-predicates (gripper_obj_far, is_dishwasher_closed, contact_check) is then - flying blind — it can't tell water-on from water-off, or 3/25 from 24/25. - We surface the env's OWN intermediate quantities WITHOUT hand-coding any task: - 1) self. int/float/bool COUNTERS/FLAGS referenced in _check_success - (washed_time, success_time, _turned_on, ...), read live. - 2) the INTERMEDIATE LOCALS the real _check_success computes this step - (each sub-predicate bool / distance / range), captured by tracing one - read-only call of _check_success. - This is the success CRITERION's current value — NOT ground-truth object coords - (the agent still localizes objects from perception).""" - try: - return self.env.get_task_progress() - except Exception: - return {} - - # ---- Phase 1: perception (state dict + rendered arrays) ---- - def current_state_dict(self) -> dict: - """Return the proprio + task + success state dict (NO object coords). - - Rendered RGB/depth/world arrays are produced by the caller via - ``_render_observation_artifacts`` (kept out of primitives so the primitives - object stays free of any output-dir / file-IO concern). - """ - o = self.env.current_raw_obs - return { - "task_language": o.get("language", self.env.get_task_language()) or "", - "success": self.env.check_success(), - # NUMERIC progress toward success (counters/sub-predicates the env's own - # _check_success computes) so the agent has a feedback loop, not just a bool. - "task_progress": self.task_progress(), - "robocasa_terminated": self.env.terminated, - "state": { - "robot0_eef_pos": self.env.eef_pos.tolist(), - "robot0_eef_quat": self.env.eef_quat.tolist(), - "robot0_gripper_qpos": self.env.gripper_qpos.tolist(), - "robot0_base_pos": np.asarray(o["robot0_base_pos"]).tolist(), - "robot0_base_quat": np.asarray(o["robot0_base_quat"]).tolist(), - }, - } - - # ---- VLA execution ---- - def run_rldx_skill( - self, - base_clip, - max_chunks, - use_prompt, - prompt, - force_reset, - n_action_steps, - settle_patience, - settle_eps, - ): - """Execute RLDX with the environment's live, full task language.""" - del use_prompt # Accepted for compatibility with historical task recipes. - configured_max_chunks = os.environ.get("RLDX_MAX_CHUNKS") - if configured_max_chunks is not None: - max_chunks = int(configured_max_chunks) - configured_action_steps = os.environ.get("RLDX_ACTION_STEPS_PER_CHUNK") - if configured_action_steps is not None: - n_action_steps = int(configured_action_steps) - configured_settle_patience = os.environ.get("RLDX_SETTLE_PATIENCE") - if configured_settle_patience is not None: - settle_patience = int(configured_settle_patience) - for name, value in ( - ("max_chunks", max_chunks), - ("n_action_steps", n_action_steps), - ("settle_patience", settle_patience), - ): - if value < 1: - return {"error": f"{name} must be positive; VLA was not executed"} - task_lang = ( - self.env.current_raw_obs.get("language") or self.env.get_task_language() - ) - if not task_lang: - return { - "error": "RoboCasa task language is unavailable; VLA was not executed", - "effective_prompt": "", - "prompt_overridden": False, - } - - prompt_overridden = prompt != task_lang - # Auto-reseed history if a non-VLA primitive ran since the last VLA call - # (read _vla_desync BEFORE clearing it) - fr = bool(force_reset) or self._vla_desync - self._vla_desync = False - result = self._rldx.run( - task_lang, - max_chunks, - n_action_steps, - base_clip=base_clip, - settle_patience=settle_patience, - settle_eps=settle_eps, - force_reset=fr, - recording=self._recording, - record_frame=self.record_frame, - ) - result["effective_prompt"] = task_lang - result["effective_max_chunks"] = max_chunks - result["effective_n_action_steps"] = n_action_steps - result["effective_settle_patience"] = settle_patience - result["prompt_overridden"] = prompt_overridden - if prompt_overridden: - result["requested_prompt"] = prompt - return result - - # ---- reset ---- - def reset(self): - # reset is ONLY for EXPLORE mode (reset-based multi-attempt recipe search). - # In no-reset / matched-scene evaluation it is a give-up-and-restart and is - # FORBIDDEN — the policy must solve the scene in one shot like the fullshot eval. - # Gated by RLDX_ALLOW_RESET (default 0 = disabled); explore launchers opt in. - if not self._allow_reset: - return { - "error": "reset is DISABLED in this run (no-reset/matched evaluation). " - "Solve the scene in one shot; do not restart the episode." - } - # EXPLORE MODE: restart the episode for a fresh attempt. New layout/ - # object placement sampled; arm/base calibration is invalidated. - self._vla_desync = True - self.env.reset() - self._pos_jac = None - self._fwd_offset = None - self._rldx.reset_session() - return {"ok": True, "reset": True, "eef": self.env.eef_pos.tolist()} - - # ---- VLA wrappers (public API for execute) ---- - def rldx_skill( - self, - base_clip=None, - max_chunks=70, - use_prompt=None, - prompt="", - force_reset=False, - n_action_steps=8, - settle_patience=int(os.environ.get("RLDX_SETTLE_PATIENCE", 999)), - settle_eps=0.012, - ): - # rldx_skill = full base motion (fullshot); base NOT clamped. - return self.run_rldx_skill( - base_clip, - max_chunks, - use_prompt, - prompt, - force_reset, - n_action_steps, - settle_patience, - settle_eps, - ) - - def rldx_arm( - self, - base_clip=0.1, - max_chunks=70, - use_prompt=None, - prompt="", - force_reset=False, - n_action_steps=8, - settle_patience=int(os.environ.get("RLDX_SETTLE_PATIENCE", 999)), - settle_eps=0.012, - ): - # rldx_arm = base CLAMPED to a small range so the VLA can micro-align - # for the grasp but can't drive away. - return self.run_rldx_skill( - base_clip, - max_chunks, - use_prompt, - prompt, - force_reset, - n_action_steps, - settle_patience, - settle_eps, - ) diff --git a/robots/robocasa/rldx_skill.py b/robots/robocasa/rldx_skill.py index abefbc356..e041ca525 100644 --- a/robots/robocasa/rldx_skill.py +++ b/robots/robocasa/rldx_skill.py @@ -25,6 +25,7 @@ import os from collections import deque +from collections.abc import Callable import imageio.v2 as imageio import numpy as np @@ -36,14 +37,9 @@ class RLDXSkill: - def __init__( - self, env_client: RoboCasaEnvClient, vla_client=None, check_cancelled=None - ): + def __init__(self, env_client: RoboCasaEnvClient, vla_client): self.env = env_client # RoboCasaEnvClient - self._vla_client = vla_client # VLA RPC client (when set, _load() uses it instead of loading the model directly) - self._check_cancelled = ( - check_cancelled # optional cancellation checkpoint callback - ) + self._vla_client = vla_client self._vdi = None # video delta indices, e.g. [-6,-4,-2,0] self._hist = None # deque of raw frame dicts self._unmap = None # lazy: eval's PandaOmronKeyConverter.unmap_action @@ -223,8 +219,9 @@ def run( settle_eps=0.012, settle_patience=2, force_reset=False, - recording=False, - record_frame=None, + *, + step_env: Callable[[np.ndarray], None], + check_cancelled: Callable[[], None], ): """Drive RLDX-1 closed-loop until it FINISHES, not a fixed tiny budget. The VLA has no terminate signal (like the eval, which runs to env-success), so we stop @@ -235,6 +232,7 @@ def run( base_clip=None -> full base motion; base_clip=v -> clamp base_motion to [-v,v] (whole-body policy: never zero it). Returns status + grasp signals so the LLM decides: continue (call again) vs done vs genuinely-failed.""" + check_cancelled() self._load() # Start a fresh video buffer for this run() call (no-op if RLDX_VIDEO_DIR unset). self._frames = [] if self._video_dir is not None else None @@ -273,8 +271,7 @@ def run( obs = self._build_obs(prompt) options = {"reset_memory": [fresh]} fresh = False - if self._check_cancelled is not None: - self._check_cancelled() + check_cancelled() actions = self._vla_client.predict(obs, options) # gym Dict -> native flat 12-d [eef_pos(3),eef_rot(3),gripper(1),base(4),mode(1)] horizon = actions["action.gripper_close"].shape[1] @@ -301,11 +298,7 @@ def run( base_motion, np.asarray(actions["action.control_mode"])[0, step], ) - if self._check_cancelled is not None: - self._check_cancelled() - self.env.step(a) - if recording: - record_frame() + step_env(a) applied += 1 self._record_frame( prompt @@ -376,15 +369,14 @@ def run( return result def reset_session(self): - if self._vla_client is not None: - try: - self._vla_client.reset_session() - except Exception: - logger.warning( - "VLA reset_session RPC failed; RLDX memory/RTC state may " - "not be reset for the next task", - exc_info=True, - ) + try: + self._vla_client.reset_session() + except Exception: + logger.warning( + "VLA reset_session RPC failed; RLDX memory/RTC state may " + "not be reset for the next task", + exc_info=True, + ) self._last_prompt = None # post-reset: next call is a fresh task if self._hist is not None: self._hist.clear() diff --git a/robots/robocasa/robot_spec.py b/robots/robocasa/robot_spec.py index 78dde5d97..ccb4500b0 100644 --- a/robots/robocasa/robot_spec.py +++ b/robots/robocasa/robot_spec.py @@ -168,6 +168,7 @@ def get_toolkit( ) return RoboCasaToolkit( runtime_kwargs=runtime_kwargs, + output_dir=config.output_dir, dashboard_events=dashboard_events, memory=memory, ) @@ -365,7 +366,7 @@ def _init_runtime( } connectors = { "env": lambda rpc: { - "env_client": RoboCasaEnvClient( + "env": RoboCasaEnvClient( rpc, expected_meta={ "task_name": args.task_name, @@ -375,10 +376,9 @@ def _init_runtime( "camera_w": 256, }, ), - "workdir": str(output_dir), "hi_res": args.hi_res or None, }, - "vla": lambda rpc: {"vla_client": RoboCasaVLAClient(rpc)}, + "vla": lambda rpc: {"model": RoboCasaVLAClient(rpc)}, } timeouts = {"env": 120.0, "vla": 300.0} selected = set(starters) if components is None else components diff --git a/robots/robocasa/toolkit.py b/robots/robocasa/toolkit.py index 0a509ba9b..a67581e3d 100644 --- a/robots/robocasa/toolkit.py +++ b/robots/robocasa/toolkit.py @@ -12,153 +12,345 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""RoboCasa toolkit: common tools + RoboCasa primitives. - -Inherits the common file/IO tools from :class:`Toolkit` and registers the -RoboCasa primitives (``move_to``, ``rldx_skill``, ``release``, ...) on top. -""" +"""RoboCasa runtime, tool composition, and observation persistence.""" from __future__ import annotations -from functools import partial +import os +from pathlib import Path from typing import TYPE_CHECKING, Any +import numpy as np + from robots.robocasa import tools as robocasa_tools -from rpent.dashboard.events import DashboardEventSink -from rpent.session import EnvState -from rpent.tools.toolkit import Toolkit -from rpent.utils.logging import get_logger, get_output_dir +from robots.robocasa.rldx_skill import RLDXSkill +from rpent.dashboard.events import DashboardEventSink, StepRecordEvent +from rpent.session import EnvState, StepRecord +from rpent.tools import Toolkit, ToolResult +from rpent.utils.logging import get_logger if TYPE_CHECKING: - from rpent.memory.manager import MemoryManager + from robots.robocasa.env_client import RoboCasaEnvClient + from robots.robocasa.vla_client import RoboCasaVLAClient + from rpent.memory import MemoryManager logger = get_logger("robocasa_toolkit") -class RoboCasaToolkit(Toolkit): - """Toolkit for the RoboCasa robot.""" +class RoboCasaRuntime: + """Clients, VLA history, controller caches, and fixed run configuration.""" + + def __init__( + self, + env: RoboCasaEnvClient, + model: RoboCasaVLAClient, + hi_res: int | None = None, + ): + self.env = env + self.hi_res = hi_res + self._keep_heavy = int(os.environ.get("RLDX_KEEP_HEAVY_NPY", "25")) + self._allow_reset = bool(int(os.environ.get("RLDX_ALLOW_RESET", "0"))) + self._pos_jac = None + self._fwd_offset = None + self._cam_meta_cache = {} + self._rldx = RLDXSkill(env, vla_client=model) + # Manual steps invalidate VLA history; consecutive VLA calls retain it. + self._vla_desync = True + + def reset(self) -> None: + self._vla_desync = True + self.env.reset() + self._pos_jac = None + self._fwd_offset = None + self._rldx.reset_session() + + +class RoboCasaToolkit(Toolkit[RoboCasaRuntime]): + """Native tools and observations for one RoboCasa planner session.""" def __init__( self, *, runtime_kwargs: dict[str, Any], + output_dir: Path | str, dashboard_events: DashboardEventSink, memory: MemoryManager, ) -> None: - """Create a RoboCasa toolkit, wiring the primitives and tools.""" - state = EnvState(get_output_dir()) + runtime = RoboCasaRuntime(**runtime_kwargs) + state = EnvState(output_dir) super().__init__( - dashboard_events=dashboard_events, state=state, memory=memory, + robot=runtime, + output_dir=output_dir, + tools=robocasa_tools.ROBOCASA_TOOLS, + dashboard_events=dashboard_events, ) - self.init_primitives(runtime_kwargs=runtime_kwargs) - self._register_robocasa_tools() - - # ---- registration: one explicit add_tool per RoboCasa tool ---- - def _register_robocasa_tools(self) -> None: - # Stateless perception tools: bind a state= kwarg via partial. - state_handlers = { - "view_env_state": partial(robocasa_tools.view_env_state, state=self._state), - "back_project_batch": partial( - robocasa_tools.back_project_batch, state=self._state - ), - "query_world_map": partial( - robocasa_tools.query_world_map, state=self._state - ), - } - for spec in robocasa_tools.TOOLS_SPEC: - name = spec["name"] - if name in state_handlers: - handler = state_handlers[name] - elif name == "finish": - handler = robocasa_tools.finish - 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( + runtime.env.reset() + record = dump_state(runtime, state, log=None) + try: + state.save( + "success_criteria.md", + runtime.env.get_success_criteria_text(), + step=None, + ) + except Exception as exc: + logger.warning("failed to save success_criteria.md: %s", exc) + try: + self._dashboard_events.emit(StepRecordEvent(record=record, env_state=state)) + except Exception: + logger.exception("Dashboard failed to publish step %s", record.step_idx) + self._action_frame_cursor = 0 + + def _capture_observation( self, *, command: dict[str, Any], - result: dict[str, Any], + result: ToolResult, elapsed_s: float, - ) -> dict[str, Any]: + ) -> tuple[dict[str, Any], list[bytes]]: frame_start = self._action_frame_cursor - self._action_frame_cursor = self._primitives.recorded_frame_count() - record = robocasa_tools.dump_state( - self._primitives, + self._action_frame_cursor = len(self._frames) + logged_result = result.to_dict() + record = dump_state( + self._robot, self._state, - log={"command": command, "result": result, "elapsed_s": elapsed_s}, + log={"command": command, "result": logged_result, "elapsed_s": elapsed_s}, ) if self._dashboard_events.enabled: try: - frames = self._primitives.frame_slice(frame_start) + frames = self._frames[frame_start:] if frames: - candidate = f"action_{command['action']}.mp4" self._state.save( - candidate, + f"action_{command['action']}.mp4", frames, step=record.step_idx, fps=20, ) - except Exception as e: + except Exception as exc: logger.warning( - "failed to save action clip for step %s: %s", - record.step_idx, - e, + "failed to save action clip for step %s: %s", record.step_idx, exc ) - out = robocasa_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 + record = self._state.get(record.step_idx) + data, images = build_observation(self._state, record) + if result.is_error: + data["log"]["result"] = { + key: value for key, value in logged_result.items() if key != "error" + } + data["agent_elapsed_s"] = elapsed_s + return data, images - def init_primitives( - self, - *, - runtime_kwargs: dict[str, Any], - ) -> None: - """Wipe stale run artifacts, build the primitives, dump step 0.""" - self._state.reset() + def solved(self) -> bool: + """Read success from the final environment record, independent of finish.""" + record = self._state.latest_record() + return bool(record is not None and record.extras.get("success", False)) + + def close(self) -> None: + """Save this robot's accumulated episode frames.""" + try: + if self._frames: + self._state.save("episode.mp4", self._frames, step=None, fps=20) + except Exception as exc: + logger.warning("failed to save episode video: %s", exc) - from robots.robocasa.primitives import RoboCasaPrimitives + def write_recipe(self, recipe_tag: str) -> str: + """Export non-error RoboCasa commands from the recorded trace.""" + return write_recipe_from_states(self._state, recipe_tag) - primitives = RoboCasaPrimitives( - check_cancelled=self.raise_if_cancelled, - **runtime_kwargs, + +# Heavy npy artifacts pruned after the ``_keep_heavy`` window elapses (the +# agent localizes from the latest frame; old world/depth maps are dead weight +# that once filled the 100GB root and deadlocked everything). +_HEAVY_ARTIFACTS = ( + "agentview_depth.npz", + "agentview_world.npz", + "wrist_depth.npz", + "wrist_world.npz", + "agentview_world_high.npz", + "navview_world.npz", +) + +# Artifact base names exposed to the agent for each camera (high-res first). +_CAMERA_IMAGE_ARTIFACTS = { + "agentview": ("agentview_high.png", "agentview.png"), + "navview": ("navview.png",), + "wrist": ("wrist_high.png", "wrist.png"), +} + + +def dump_state( + runtime: RoboCasaRuntime, + env_state: EnvState, + log: dict | None = None, +) -> StepRecord: + """Record one RoboCasa observation through its owned state record. + + Appends a new :class:`StepRecord` (proprio + task + success + vla_desync) + via :meth:`EnvState.record_step`, saves the rendered RGB / depth / world + artifacts for that step, and prunes heavy npy artifacts that fell out of + the ``_keep_heavy`` window. + """ + raw = runtime.env.current_raw_obs + extras = { + "task_language": raw.get("language", runtime.env.get_task_language()) or "", + "success": runtime.env.check_success(), + "task_progress": runtime.env.get_task_progress(), + "vla_desync": runtime._vla_desync, + } + state = { + "robot0_eef_pos": runtime.env.eef_pos.tolist(), + "robot0_eef_quat": runtime.env.eef_quat.tolist(), + "robot0_gripper_qpos": runtime.env.gripper_qpos.tolist(), + "robot0_base_pos": np.asarray(raw["robot0_base_pos"]).tolist(), + "robot0_base_quat": np.asarray(raw["robot0_base_quat"]).tolist(), + } + log = log or {} + with env_state.record_step( + state=state, + terminated=runtime.env.terminated, + truncated=False, + command=log.get("command"), + result=log.get("result"), + elapsed_s=log.get("elapsed_s"), + extras=extras, + ) as step_idx: + _save_observation_artifacts(runtime, env_state, step_idx) + env_state.prune_artifacts( + _HEAVY_ARTIFACTS, step=step_idx, keep_last=runtime._keep_heavy ) - primitives.reset() - primitives.start_recording() - self._action_frame_cursor = primitives.recorded_frame_count() - record = robocasa_tools.dump_state(primitives, self._state, log=None) - try: - self._state.save( - "success_criteria.md", - primitives.dump_success_criteria(), - step=None, + return env_state.get(step_idx) + + +def _save_observation_artifacts( + runtime: RoboCasaRuntime, + env_state: EnvState, + step_idx: int, +) -> None: + """Render and save all per-step observation artifacts for ``step_idx``.""" + env = runtime.env + hi_res = runtime.hi_res + + # ---- agentview + wrist: rgb, depth, world map, camera meta ---- + for cam, image_name, depth_name, world_name, meta_name in ( + ( + "agentview", + "agentview.png", + "agentview_depth.npz", + "agentview_world.npz", + "agentview_metadata.json", + ), + ( + "wrist", + "wrist.png", + "wrist_depth.npz", + "wrist_world.npz", + "wrist_metadata.json", + ), + ): + rgb, depth = env.render_camera(cam, depth=True) + env_state.save(image_name, rgb, step=step_idx) + env_state.save(depth_name, depth.astype(np.float32), step=step_idx) + env_state.save(world_name, env.world_map(cam).astype(np.float32), step=step_idx) + if cam not in runtime._cam_meta_cache or cam == "wrist": + runtime._cam_meta_cache[cam] = env.get_camera_meta(cam) + env_state.save(meta_name, runtime._cam_meta_cache[cam], step=step_idx) + + # ---- hi-res agentview (SAM grounding / fine localize) ---- + if hi_res: + hrgb, _ = env.render_camera("agentview", hi_res, hi_res, depth=True) + env_state.save("agentview_high.png", hrgb, step=step_idx) + env_state.save( + "agentview_world_high.npz", + env.world_map("agentview", hi_res, hi_res).astype(np.float16), + step=step_idx, + ) + + # ---- navview: base-mounted forward-down floor camera (follows the base) ---- + nrgb, _ = env.render_camera("navview", depth=True) + nworld = env.world_map("navview").astype(np.float32) + env_state.save("navview.png", nrgb, step=step_idx) + env_state.save("navview_world.npz", nworld, step=step_idx) + floor = (nworld[:, :, 2] < 0.12) & (nworld[:, :, 2] > -0.2) + overlay = nrgb.copy() + overlay[floor] = [0, 255, 0] + env_state.save("navview_floor.png", overlay, step=step_idx) + + +def build_observation( + state: EnvState, record: StepRecord +) -> tuple[dict[str, Any], list[bytes]]: + """Return recorded data and ordered agentview, navigation, and wrist PNGs.""" + nn = record.step_idx + extras = record.extras + out: dict = { + "step": nn, + "task_progress": extras.get("task_progress", {}), + "task_language": extras.get("task_language", ""), + "state": record.state, + "robocasa_terminated": record.terminated, + "vla_desync": extras.get("vla_desync", False), + "success": extras.get("success", False), + "log": { + "command": record.command, + "result": record.result, + "elapsed_s": record.elapsed_s, + }, + "images": [], + "artifacts": sorted(record.artifacts), + } + + images: list[bytes] = [] + for camera, candidates in _CAMERA_IMAGE_ARTIFACTS.items(): + for name in candidates: + if name not in record.artifacts: + continue + try: + image = state.load_bytes(name, step=nn) + except FileNotFoundError: + continue + images.append(image) + out["images"].append( + { + "role": "nav_view" if camera == "navview" else "calibration_frame", + "camera": camera, + "artifact": name, + } ) - except Exception as e: - logger.warning("failed to save success_criteria.md: %s", e) - self._primitives = primitives - self._publish_step(record) + break - def close(self) -> None: - """Flush the agent-side video buffer through ``EnvState``.""" - try: - frames = self._primitives.stop_recording() - if frames: - self._state.save("episode.mp4", frames, step=None, fps=20) - except Exception as e: - logger.warning("failed to save episode video: %s", e) + return out, images - def solved(self) -> bool: - """Return the success value from the final recorded environment state.""" - record = self._state.latest_record() - return bool(record is not None and record.extras.get("success", False)) - def write_recipe(self, recipe_tag: str) -> str: - """Write the RoboCasa recipe JSONL from the dumped state trace.""" - return robocasa_tools.write_recipe_from_states(self._state, recipe_tag) +_PRIMITIVE_ACTIONS = frozenset( + { + "move_to", + "move_delta", + "rotate_pitch", + "set_gripper", + "release", + "scripted_grasp", + "rldx_skill", + "rldx_arm", + "navigate_to", + "move_base", + "reset", + } +) + + +def write_recipe_from_states(state: EnvState, recipe_tag: str) -> str: + """Export non-error RoboCasa primitive commands from the state trace as JSONL.""" + commands = [] + for record in state.records(): + command = record.command + if not isinstance(command, dict): + continue + if command.get("action") not in _PRIMITIVE_ACTIONS: + continue + result = record.result + if isinstance(result, dict) and result.get("error"): + continue + commands.append(command) + recipe_name = f"{recipe_tag}_recipe.jsonl" + state.save(recipe_name, commands, step=None) + return recipe_name diff --git a/robots/robocasa/tools.py b/robots/robocasa/tools.py index 91796754f..87f907495 100644 --- a/robots/robocasa/tools.py +++ b/robots/robocasa/tools.py @@ -12,783 +12,646 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""RoboCasa tool schemas and handlers backed by the run's ``EnvState``. - -The state trace (``states.json`` manifest + per-step artifact files under -``/``) is owned by :class:`rpent.session.EnvState`. Tool handlers -that need to read it take a ``state: EnvState`` keyword argument (bound by the -toolkit via :func:`functools.partial`); state-advancing primitive tools capture -state automatically through :meth:`RoboCasaToolkit.get_env_state`. -""" +"""Native RoboCasa action and perception tools.""" from __future__ import annotations -from typing import TYPE_CHECKING +import os +from functools import partial +from typing import TYPE_CHECKING, Annotated, Literal import numpy as np +from pydantic import Field -from rpent.session import EnvState, StepRecord -from rpent.tools.toolkit import readonly +from rpent.tools import ToolContext, ToolResult, readonly, tool if TYPE_CHECKING: - from robots.robocasa.primitives import RoboCasaPrimitives - -# ---- TOOLS_SPEC: 15 Anthropic-shaped tool schemas ---- - -TOOLS_SPEC = [ - # ---- primitive tools (11): dispatched by the toolkit base to primitives. ---- - { - "name": "move_to", - "description": ( - "Scripted EEF servo to a world-frame XYZ target via the OSC " - "controller. Holds pitch/yaw orientation (use rotate_pitch to " - "reorient). gripper='hold' (DEFAULT) maintains current finger width " - "— carry-safe without crushing small objects. Pass +1 to close, " - "-1 to open. NEVER command a single move_to with |dxyz| > 0.30 — " - "OSC flips IK; split long traversal into 2-3 mid waypoints at carry z." - ), - "input_schema": { - "type": "object", - "properties": { - "xyz": { - "type": "array", - "description": "World-frame target [x, y, z] in meters", - "items": {"type": "number"}, - "minItems": 3, - "maxItems": 3, - }, - "gripper": { - "type": ["number", "string"], - "description": ( - "Gripper: +1 close, -1 open, or 'hold' to maintain " - "current finger width (default 'hold')" - ), - }, - "step_clip": { - "type": "number", - "description": "Per-step dxyz cap, m (default 0.02)", - }, - "max_steps": { - "type": "integer", - "description": "Step budget (default 200)", - }, - "tol": { - "type": "number", - "description": "Position tolerance, m (default 0.012)", - }, - }, - "required": ["xyz"], - }, - }, - { - "name": "move_delta", - "description": ( - "Relative EEF displacement from the current position. Computes " - "target = current_eef + dxyz and delegates to move_to. " - "Use for small adjustments (micro-align for grasp, approach). " - "gripper='hold' (DEFAULT) maintains current finger width." - ), - "input_schema": { - "type": "object", - "properties": { - "dxyz": { - "type": "array", - "description": "Relative displacement [dx, dy, dz] in meters", - "items": {"type": "number"}, - "minItems": 3, - "maxItems": 3, - }, - "gripper": { - "type": ["number", "string"], - "description": ( - "Gripper: +1 close, -1 open, or 'hold' (default 'hold')" - ), - }, - "step_clip": { - "type": "number", - "description": "Per-step dxyz cap, m (default 0.02)", - }, - "max_steps": { - "type": "integer", - "description": "Step budget (default 80)", - }, - }, - "required": ["dxyz"], - }, - }, - { - "name": "rotate_pitch", - "description": ( - "Tilt the wrist forward (axis-angle about the control X-axis). " - "This pitches the gripper down/up. Holds xyz fixed. " - "Use before threading the gripper into a narrow opening whose " - "front face normal is along world +/-y." - ), - "input_schema": { - "type": "object", - "properties": { - "target_pitch": { - "type": "number", - "description": ( - "Absolute pitch target, radians (clamped +/-1.5; default 0.6)" - ), - }, - "gripper": { - "type": "number", - "description": "Gripper command held during rotation (default +1)", - }, - "n": { - "type": "integer", - "description": "Number of env steps for the rotation (default 12)", - }, - }, - }, - }, - { - "name": "set_gripper", - "description": ( - "Hold the current EEF pose and drive the gripper command for " - "`steps` env steps. Use to firm up a grip mid-carry or to " - "actively open/close the gripper." - ), - "input_schema": { - "type": "object", - "properties": { - "gripper": { - "type": "number", - "description": "Gripper command: +1 close, -1 open (default +1)", - }, - "steps": { - "type": "integer", - "description": "Number of env steps to hold (default 10)", - }, - }, - }, - }, - { - "name": "release", - "description": ( - "Open the gripper for `steps` env steps while holding EEF in " - "place. Delegates to set_gripper(-1.0, steps=steps). " - "Use to drop a grasped object." - ), - "input_schema": { - "type": "object", - "properties": { - "steps": { - "type": "integer", - "description": "Number of env steps (default 10)", - }, - }, - }, - }, - { - "name": "scripted_grasp", - "description": ( - "Coarse scripted grasp sequence: open -> hover above target -> " - "descend -> close -> lift. A fallback when the VLA closed-loop " - "grasp is unavailable. For hard objects prefer rldx_arm. " - "approach_z and grasp_z_offset are RELATIVE offsets from the " - "target xyz." - ), - "input_schema": { - "type": "object", - "properties": { - "xyz": { - "type": "array", - "description": "World-frame grasp target [x, y, z] in meters", - "items": {"type": "number"}, - "minItems": 3, - "maxItems": 3, - }, - "approach_z": { - "type": "number", - "description": "Z offset above target before descent, m (default 0.10)", - }, - "grasp_z_offset": { - "type": "number", - "description": "Z offset at grasp point (default 0.0; negative = below target)", - }, - "step_clip": { - "type": "number", - "description": "Per-step dxyz cap during descent, m (default 0.02)", - }, - }, - "required": ["xyz"], - }, - }, - { - "name": "rldx_skill", - "description": ( - "RLDX VLA closed-loop skill — FULL base motion allowed. The VLA " - "drives both arm and mobile base. Use for full-body tasks where " - "the base must reposition (e.g. navigating to a counter while " - "reaching). Pass the complete live task_language verbatim; the " - "runtime always uses that environment language for RLDX. Do NOT " - "interrupt consecutive VLA calls with manual " - "primitives — that breaks VLA frame history continuity " - "(sets vla_desync=True)." - ), - "input_schema": { - "type": "object", - "properties": { - "prompt": { - "type": "string", - "description": "Complete live task_language, copied verbatim", - }, - "base_clip": { - "type": ["number", "null"], - "description": "Base motion magnitude cap (default null = no clamp)", - }, - "max_chunks": { - "type": "integer", - "description": "Action-chunk budget (default 70)", - }, - "force_reset": { - "type": "boolean", - "description": "Force VLA frame history reset (default False)", - }, - "n_action_steps": { - "type": "integer", - "description": "Actions per VLA chunk (default 8)", - }, - "settle_patience": { - "type": "integer", - "description": ( - "Settle step budget before declaring done " - "(default 999; do NOT set small)" - ), - }, - "settle_eps": { - "type": "number", - "description": "Settle position tolerance, m (default 0.012)", - }, - }, - "required": ["prompt"], - }, - }, - { - "name": "rldx_arm", - "description": ( - "RLDX VLA closed-loop skill — base CLAMPED to small motions " - "(base_clip=0.1 default). The VLA drives the arm for precise " - "micro-alignment (e.g. fine-tuning a grasp approach) but cannot " - "drive the base away. Pass the complete live task_language " - "verbatim; the runtime always uses that environment language for " - "RLDX. Do NOT interrupt consecutive VLA calls with manual primitives." - ), - "input_schema": { - "type": "object", - "properties": { - "prompt": { - "type": "string", - "description": "Complete live task_language, copied verbatim", - }, - "base_clip": { - "type": ["number", "null"], - "description": "Base motion magnitude cap (default 0.1 = small)", - }, - "max_chunks": { - "type": "integer", - "description": "Action-chunk budget (default 70)", - }, - "force_reset": { - "type": "boolean", - "description": "Force VLA frame history reset (default False)", - }, - "n_action_steps": { - "type": "integer", - "description": "Actions per VLA chunk (default 8)", - }, - "settle_patience": { - "type": "integer", - "description": ( - "Settle step budget before declaring done " - "(default 999; do NOT set small)" - ), - }, - "settle_eps": { - "type": "number", - "description": "Settle position tolerance, m (default 0.012)", - }, - }, - "required": ["prompt"], - }, - }, - { - "name": "navigate_to", - "description": ( - "Drive the mobile base toward a WORLD (x, y) target. " - "Online-calibrates the base forward-heading, then turns to face " - "+ drives forward closed-loop. Holds the arm in place. " - "gripper='hold' (DEFAULT) maintains current finger width while " - "driving (carry-safe). Use tol = expected approach distance " - "+ object radius." - ), - "input_schema": { - "type": "object", - "properties": { - "xy": { - "type": "array", - "description": "World-frame target [x, y] in meters (z ignored if provided)", - "items": {"type": "number"}, - "minItems": 2, - "maxItems": 2, - }, - "tol": { - "type": "number", - "description": "Distance threshold to stop, m (default 0.20)", - }, - "max_steps": { - "type": "integer", - "description": "Step budget (default 300)", - }, - "gripper": { - "type": ["number", "string"], - "description": ( - "Gripper while driving: +1 close, -1 open, " - "or 'hold' (default 'hold')" - ), - }, - }, - "required": ["xy"], - }, - }, - { - "name": "move_base", - "description": ( - "Raw base velocity commands in the robot's LOCAL frame. " - "+forward = drive forward, +lateral = strafe right, " - "+turn = rotate CCW (yaw). All values clamped [-1, 1]. " - "Use move_base for fine base adjustments near a target; " - "use navigate_to for long-range navigation. " - "gripper='hold' (DEFAULT) maintains finger width while driving." - ), - "input_schema": { - "type": "object", - "properties": { - "forward": { - "type": "number", - "description": "Forward velocity, [-1, 1] (default 0)", - }, - "lateral": { - "type": "number", - "description": "Lateral / strafe velocity, [-1, 1] (default 0)", - }, - "turn": { - "type": "number", - "description": "Yaw rotation velocity, [-1, 1] (default 0)", - }, - "steps": { - "type": "integer", - "description": "Number of env steps (default 10)", - }, - "gripper": { - "type": ["number", "string"], - "description": ( - "Gripper while driving: +1 close, -1 open, " - "or 'hold' (default 'hold')" - ), - }, - }, - }, - }, - { - "name": "reset", - "description": ( - "Restart the episode (new layout / object placement sampled). " - "Arm and base calibration are invalidated on reset. " - "DISABLED in no-reset / matched evaluation — the policy must " - "solve the scene in one shot. Only available in EXPLORE mode " - "when RLDX_ALLOW_RESET is enabled." - ), - "input_schema": { - "type": "object", - "properties": {}, - }, - }, - # ---- perception tools (4) -- module-level @readonly handlers ---- - { - "name": "view_env_state", - "description": ( - "Read step NN from states.json + the matching state images " - "in the output dir. If step is null, returns the latest entry. " - "Each entry contains the env state, robocasa_terminated flag, " - "task_progress, vla_desync status, and log. Embeds available " - "PNGs as multimodal image content blocks. Use calibration-frame " - "agentview images for pixel back-projection; use navview for " - "base navigation and floor walkability; use wrist for close-range " - "details near the gripper." - ), - "input_schema": { - "type": "object", - "properties": { - "step": { - "type": ["integer", "null"], - "description": "Step number; 0 = initial. Null = latest.", - }, - }, - }, - }, - { - "name": "back_project_batch", - "description": ( - "Back-project MULTIPLE pixels to world XYZ points in a single " - "call. Loads the world map once and queries all pixels — " - "replaces N separate back_project calls. " - "Returns each pixel's world_xyz plus a summary with " - "median_xyz across valid pixels.\n\n" - "USE THIS for robust object localization: sample 3-8 pixels " - "on the target object and read summary.median_xyz. " - "Maximum 50 pixels per call." - ), - "input_schema": { - "type": "object", - "properties": { - "pixels": { - "type": "array", - "description": "List of [row, col] pixel coordinates (max 50)", - "items": { - "type": "array", - "items": {"type": "integer"}, - "minItems": 2, - "maxItems": 2, - }, - "minItems": 1, - "maxItems": 50, - }, - "step": { - "type": ["integer", "null"], - "description": "Depth / world-map step to use (default latest).", - }, - "camera": { - "type": "string", - "enum": ["agentview", "navview", "wrist"], - "description": "Camera to back-project from (default agentview).", - }, - "resolution": { - "type": "string", - "enum": ["high", "low"], - "description": ( - "Coordinate system for pixels (default low). " - "Use 'low' for the standard 256x256 world map." - ), - }, - }, - "required": ["pixels"], - }, - }, - { - "name": "query_world_map", - "description": ( - "Query the world map by Z-range / XY region to find objects " - "at specific heights. Loads the world map once, filters pixels " - "by z_min <= z <= z_max, optionally restricts to x_range / " - "y_range, then clusters contiguous pixels into objects.\n\n" - "TYPICAL USES:\n" - "- z_min=0.85, z_max=0.95 -> countertop-height objects\n" - "- z_min=0.0, z_max=0.12, camera='navview' -> walkable floor\n" - "- z_min=0.85, z_max=0.95, x_range=[0,2], y_range=[-3,-1] -> " - "counter objects in a specific quadrant" - ), - "input_schema": { - "type": "object", - "properties": { - "z_min": { - "type": "number", - "description": "Minimum Z in meters (default 0.85 for counter height).", - }, - "z_max": { - "type": "number", - "description": "Maximum Z in meters (default 0.95 for counter height).", - }, - "x_range": { - "type": ["array", "null"], - "description": "Optional X range [min, max] in meters; null = no filter.", - "items": {"type": "number"}, - "minItems": 2, - "maxItems": 2, - }, - "y_range": { - "type": ["array", "null"], - "description": "Optional Y range [min, max] in meters; null = no filter.", - "items": {"type": "number"}, - "minItems": 2, - "maxItems": 2, - }, - "camera": { - "type": "string", - "enum": ["agentview", "navview", "wrist"], - "description": "Camera world map to query (default agentview).", - }, - "resolution": { - "type": "string", - "enum": ["high", "low"], - "description": "World map resolution (default low).", - }, - "min_cluster_size": { - "type": "integer", - "description": "Minimum pixels per cluster to report (default 10).", - }, - }, - }, - }, - { - "name": "finish", - "description": ( - "Declare the task finished. Call when robocasa_terminated " - "becomes True (success detected), or when genuinely stuck " - "after honest exploration. Provide a 1-3 sentence summary " - "of what worked and what failed." - ), - "input_schema": { - "type": "object", - "properties": { - "status": { - "type": "string", - "enum": ["success", "failure", "stuck"], - "description": "Task outcome classification.", - }, - "summary": { - "type": "string", - "description": "1-3 sentence summary of what worked / what failed.", - }, - }, - "required": ["status", "summary"], + from robots.robocasa.toolkit import RoboCasaRuntime + +OSC_ROT_SCALE = 0.5 # action 1.0 -> 0.5 rad + + +def _step_env(ctx: ToolContext[RoboCasaRuntime], action: np.ndarray) -> None: + """Check cancellation and record each physical step, including calibration.""" + ctx.check_cancelled() + ctx.robot.env.step(action) + image = ctx.robot.env.render_camera( + camera_name="agentview", height=256, width=256, depth=False + ) + ctx.record_frame(np.asarray(image, dtype=np.uint8)) + + +def _calibrate_pos_jacobian(ctx: ToolContext[RoboCasaRuntime], gripper=-1.0): + """Probe 3 unit arm-xyz actions, measure world dpos -> 3x3 jacobian J s.t. + world_dpos ~= J @ action_xyz. move_to inverts J to map desired world delta.""" + runtime = ctx.robot + cols = [] + for axis in range(3): + p0 = runtime.env.eef_pos.copy() + a = np.zeros(12) + a[11] = -1.0 + a[axis] = 0.4 + a[6] = gripper + for _ in range(3): + _step_env(ctx, a) + d = (runtime.env.eef_pos - p0) / (0.4 * 3) # world dpos per unit action + cols.append(d) + # settle back is not needed (closed-loop re-reads); keep going + runtime._pos_jac = np.stack(cols, axis=1) # 3x3: world_dpos = J @ a_xyz + return runtime._pos_jac + + +def _base_pos(ctx: ToolContext[RoboCasaRuntime]) -> np.ndarray: + return np.asarray( + ctx.robot.env.current_raw_obs["robot0_base_pos"], dtype=np.float64 + ) + + +def _calibrate_forward(ctx: ToolContext[RoboCasaRuntime], gripper=1.0): + """Drive forward briefly, measure the WORLD direction the base actually goes, + so navigate_to can steer regardless of the base->world frame offset.""" + from scipy.spatial.transform import Rotation as R + + runtime = ctx.robot + p0 = _base_pos(ctx) + y0 = float( + R.from_quat( + np.asarray(runtime.env.current_raw_obs["robot0_base_quat"]) + ).as_euler("xyz")[2] + ) + a = np.zeros(12) + a[11] = 1.0 + a[6] = float(np.clip(gripper, -1, 1)) + a[7] = 1.0 + for _ in range(6): + _step_env(ctx, a) + p1 = _base_pos(ctx) + disp = (p1 - p0)[:2] + if np.linalg.norm(disp) > 0.005: + runtime._fwd_offset = np.arctan2(disp[1], disp[0]) - y0 + else: + runtime._fwd_offset = 0.0 + return runtime._fwd_offset + + +def _resolve_grip(ctx: ToolContext[RoboCasaRuntime], gripper, target_q): + """Return the a[6] gripper command for a motion step. + gripper="hold" (DEFAULT for moves) -> SERVO the fingers back to `target_q`, + the width they had when the motion began. This is the carry-safe hold: the + gripper action is a CLOSE-VELOCITY command, so a sustained +1 keeps driving the + fingers shut and SQUEEZES a small object OUT (verified: bread qpos 0.0376 -> 0.0005 + during a +1 carry). Servoing to the grasped width holds the object without crushing + it and without letting it drift open. A numeric gripper (+1 close / -1 open) is an + EXPLICIT override and passes through unchanged.""" + runtime = ctx.robot + if isinstance(gripper, str): + cur = float(runtime.env.gripper_qpos[0]) + return float( + np.clip(60.0 * (cur - target_q), -1.0, 1.0) + ) # +a[6] closes (qpos↓) + return float(np.clip(gripper, -1, 1)) + + +def _move_to( + ctx: ToolContext[RoboCasaRuntime], + xyz: list[float] | np.ndarray, + gripper: float | str = "hold", + step_clip: float = 0.02, + max_steps: int = 200, + tol: float = 0.012, +) -> ToolResult: + runtime = ctx.robot + runtime._vla_desync = True + target = np.asarray(xyz, dtype=np.float64) + target_q = float(runtime.env.gripper_qpos[0]) # finger width to hold + if runtime._pos_jac is None: + _calibrate_pos_jacobian(ctx, gripper=_resolve_grip(ctx, gripper, target_q)) + Jinv = np.linalg.pinv(runtime._pos_jac) + for i in range(max_steps): + cur = runtime.env.eef_pos + err = target - cur + dist = float(np.linalg.norm(err)) + if dist < tol: + return ToolResult( + data={ + "ok": True, + "steps": i, + "final_dist": dist, + "eef": cur.tolist(), + "gripper_qpos": round(float(runtime.env.gripper_qpos[0]), 4), + } + ) + step_world = err if dist <= step_clip else err / dist * step_clip + a_xyz = np.clip(Jinv @ step_world, -1, 1) + a = np.zeros(12) + a[11] = -1.0 + a[0:3] = a_xyz + a[6] = _resolve_grip(ctx, gripper, target_q) + _step_env(ctx, a) + + cur = runtime.env.eef_pos + return ToolResult( + data={ + "ok": False, + "steps": max_steps, + "final_dist": float(np.linalg.norm(target - cur)), + "eef": cur.tolist(), + "gripper_qpos": round(float(runtime.env.gripper_qpos[0]), 4), }, - }, -] - - -# ---- state persistence (dump_state writes through the run's EnvState) ---- - -# Heavy npy artifacts pruned after the ``_keep_heavy`` window elapses (the -# agent localizes from the latest frame; old world/depth maps are dead weight -# that once filled the 100GB root and deadlocked everything). -_HEAVY_ARTIFACTS = ( - "agentview_depth.npz", - "agentview_world.npz", - "wrist_depth.npz", - "wrist_world.npz", - "agentview_world_high.npz", - "navview_world.npz", -) + error="move_to did not reach the target within max_steps", + ) -# Artifact base names exposed to the agent for each camera (high-res first). -_CAMERA_IMAGE_ARTIFACTS = { - "agentview": ("agentview_high.png", "agentview.png"), - "navview": ("navview.png",), - "wrist": ("wrist_high.png", "wrist.png"), -} -# Camera -> (low-res world map artifact, high-res world map artifact or None) -_CAMERA_WORLD_ARTIFACTS = { - "agentview": ("agentview_world.npz", "agentview_world_high.npz"), - "navview": ("navview_world.npz", None), - "wrist": ("wrist_world.npz", None), -} +@tool +def move_to( + xyz: Annotated[list[float], Field(min_length=3, max_length=3)], + gripper: float | str = "hold", + step_clip: float = 0.02, + max_steps: int = 200, + tol: float = 0.012, + *, + ctx: ToolContext[RoboCasaRuntime], +) -> ToolResult: + """Scripted EEF servo to a world-frame XYZ target via the OSC controller. Holds pitch/yaw orientation (use rotate_pitch to reorient). gripper='hold' (DEFAULT) maintains current finger width — carry-safe without crushing small objects. Pass +1 to close, -1 to open. NEVER command a single move_to with |dxyz| > 0.30 — OSC flips IK; split long traversal into 2-3 mid waypoints at carry z. + + Args: + xyz: World-frame target [x, y, z] in meters + gripper: Gripper: +1 close, -1 open, or 'hold' to maintain current finger width (default 'hold') + step_clip: Per-step dxyz cap, m (default 0.02) + max_steps: Step budget (default 200) + tol: Position tolerance, m (default 0.012) + """ + return _move_to(ctx, xyz, gripper, step_clip, max_steps, tol) + +@tool +def move_delta( + dxyz: Annotated[list[float], Field(min_length=3, max_length=3)], + gripper: float | str = "hold", + step_clip: float = 0.02, + max_steps: int = 80, + *, + ctx: ToolContext[RoboCasaRuntime], +) -> ToolResult: + """Relative EEF displacement from the current position. Computes target = current_eef + dxyz and delegates to move_to. Use for small adjustments (micro-align for grasp, approach). gripper='hold' (DEFAULT) maintains current finger width. + + Args: + dxyz: Relative displacement [dx, dy, dz] in meters + gripper: Gripper: +1 close, -1 open, or 'hold' (default 'hold') + step_clip: Per-step dxyz cap, m (default 0.02) + max_steps: Step budget (default 80) + """ + runtime = ctx.robot + return _move_to( + ctx, runtime.env.eef_pos + np.asarray(dxyz), gripper, step_clip, max_steps + ) -def dump_state( - primitives: RoboCasaPrimitives, - env_state: EnvState, - log: dict | None = None, -) -> StepRecord: - """Record one RoboCasa observation through its owned state record. - Appends a new :class:`StepRecord` (proprio + task + success + vla_desync) - via :meth:`EnvState.record_step`, saves the rendered RGB / depth / world - artifacts for that step, and prunes heavy npy artifacts that fell out of - the ``_keep_heavy`` window. +@tool +def rotate_pitch( + target_pitch: float = 0.6, + gripper: float = 1, + n: int = 12, + *, + ctx: ToolContext[RoboCasaRuntime], +) -> ToolResult: + """Tilt the wrist forward (axis-angle about the control X-axis). This pitches the gripper down/up. Holds xyz fixed. Use before threading the gripper into a narrow opening whose front face normal is along world +/-y. + + Args: + target_pitch: Absolute pitch target, radians (clamped +/-1.5; default 0.6) + gripper: Gripper command held during rotation (default +1) + n: Number of env steps for the rotation (default 12) """ - state_dict = primitives.current_state_dict() - log = log or {} - with env_state.record_step( - state=state_dict["state"], - terminated=state_dict["robocasa_terminated"], - truncated=False, - command=log.get("command"), - result=log.get("result"), - elapsed_s=log.get("elapsed_s"), - extras={ - "task_language": state_dict["task_language"], - "success": state_dict["success"], - "task_progress": state_dict["task_progress"], - "vla_desync": primitives._vla_desync, - }, - ) as step_idx: - _save_observation_artifacts(primitives, env_state, step_idx) - env_state.prune_artifacts( - _HEAVY_ARTIFACTS, step=step_idx, keep_last=primitives._keep_heavy - ) - return env_state.get(step_idx) - - -def _save_observation_artifacts( - primitives: RoboCasaPrimitives, - env_state: EnvState, - step_idx: int, -) -> None: - """Render and save all per-step observation artifacts for ``step_idx``.""" - env = primitives.env - hi_res = primitives.hi_res - - # ---- agentview + wrist: rgb, depth, world map, camera meta ---- - for cam, image_name, depth_name, world_name, meta_name in ( - ( - "agentview", - "agentview.png", - "agentview_depth.npz", - "agentview_world.npz", - "agentview_metadata.json", - ), - ( - "wrist", - "wrist.png", - "wrist_depth.npz", - "wrist_world.npz", - "wrist_metadata.json", - ), - ): - rgb, depth = env.render_camera(cam, depth=True) - env_state.save(image_name, rgb, step=step_idx) - env_state.save(depth_name, depth.astype(np.float32), step=step_idx) - env_state.save(world_name, env.world_map(cam).astype(np.float32), step=step_idx) - if cam not in primitives._cam_meta_cache or cam == "wrist": - primitives._cam_meta_cache[cam] = env.get_camera_meta(cam) - env_state.save(meta_name, primitives._cam_meta_cache[cam], step=step_idx) - - # ---- hi-res agentview (SAM grounding / fine localize) ---- - if hi_res: - hrgb, _ = env.render_camera("agentview", hi_res, hi_res, depth=True) - env_state.save("agentview_high.png", hrgb, step=step_idx) - env_state.save( - "agentview_world_high.npz", - env.world_map("agentview", hi_res, hi_res).astype(np.float16), - step=step_idx, - ) + runtime = ctx.robot + runtime._vla_desync = True + per = float(np.clip(target_pitch, -1.5, 1.5)) / n + action = np.zeros(12) + action[11] = -1.0 + action[3] = np.clip(per / OSC_ROT_SCALE, -1, 1) + action[6] = float(np.clip(gripper, -1, 1)) + for _ in range(n): + _step_env(ctx, action) + return ToolResult(data={"ok": True, "eef": runtime.env.eef_pos.tolist()}) + + +def _set_gripper( + ctx: ToolContext[RoboCasaRuntime], gripper: float = 1, steps: int = 10 +) -> ToolResult: + runtime = ctx.robot + runtime._vla_desync = True + a = np.zeros(12) + a[11] = -1.0 + a[6] = float(np.clip(gripper, -1, 1)) + for _ in range(steps): + _step_env(ctx, a) + return ToolResult( + data={"ok": True, "gripper_qpos": runtime.env.gripper_qpos.tolist()} + ) + + +@tool +def set_gripper( + gripper: float = 1, steps: int = 10, *, ctx: ToolContext[RoboCasaRuntime] +) -> ToolResult: + """Hold the current EEF pose and drive the gripper command for `steps` env steps. Use to firm up a grip mid-carry or to actively open/close the gripper. + + Args: + gripper: Gripper command: +1 close, -1 open (default +1) + steps: Number of env steps to hold (default 10) + """ + return _set_gripper(ctx, gripper, steps) - # ---- navview: base-mounted forward-down floor camera (follows the base) ---- - nrgb, _ = env.render_camera("navview", depth=True) - nworld = env.world_map("navview").astype(np.float32) - env_state.save("navview.png", nrgb, step=step_idx) - env_state.save("navview_world.npz", nworld, step=step_idx) - floor = (nworld[:, :, 2] < 0.12) & (nworld[:, :, 2] > -0.2) - overlay = nrgb.copy() - overlay[floor] = [0, 255, 0] - env_state.save("navview_floor.png", overlay, step=step_idx) +@tool +def release(steps: int = 10, *, ctx: ToolContext[RoboCasaRuntime]) -> ToolResult: + """Open the gripper for `steps` env steps while holding EEF in place. Delegates to set_gripper(-1.0, steps=steps). Use to drop a grasped object. -# ---- tool handlers (@readonly perception tools take state: EnvState) ---- + Args: + steps: Number of env steps (default 10) + """ + return _set_gripper(ctx, -1.0, steps=steps) -@readonly -def view_env_state(step: int | None = None, *, state: EnvState) -> dict: - """Read one recorded state with embedded camera / navview / wrist images.""" - try: - record = state.get(step if step is not None else -1) - except Exception as exc: - return {"error": f"state step not available: {exc}"} +@tool +def scripted_grasp( + xyz: Annotated[list[float], Field(min_length=3, max_length=3)], + approach_z: float = 0.10, + grasp_z_offset: float = 0.0, + step_clip: float = 0.02, + *, + ctx: ToolContext[RoboCasaRuntime], +) -> ToolResult: + """Coarse scripted grasp sequence: open -> hover above target -> descend -> close -> lift. A fallback when the VLA closed-loop grasp is unavailable. For hard objects prefer rldx_arm. approach_z and grasp_z_offset are RELATIVE offsets from the target xyz. + + Args: + xyz: World-frame grasp target [x, y, z] in meters + approach_z: Z offset above target before descent, m (default 0.10) + grasp_z_offset: Z offset at grasp point (default 0.0; negative = below target) + step_clip: Per-step dxyz cap during descent, m (default 0.02) + """ + runtime = ctx.robot + t = np.asarray(xyz, dtype=np.float64) + _set_gripper(ctx, -1.0, steps=4) + r = _move_to(ctx, t + [0, 0, approach_z], gripper=-1.0, step_clip=step_clip) + if r.is_error: + r.data["stage"] = "approach" + return r + r = _move_to( + ctx, t + [0, 0, grasp_z_offset], gripper=-1.0, step_clip=0.012, tol=0.01 + ) + if r.is_error: + r.data["stage"] = "descent" + return r + _set_gripper(ctx, +1.0, steps=14) + r = _move_to(ctx, t + [0, 0, approach_z + 0.05], gripper="hold", step_clip=0.015) + if r.is_error: + r.data["stage"] = "lift" + return r + return ToolResult( + data={ + "ok": True, + "gripper_qpos": runtime.env.gripper_qpos.tolist(), + "eef": runtime.env.eef_pos.tolist(), + } + ) - nn = record.step_idx - extras = record.extras - out: dict = { - "step": nn, - "task_progress": extras.get("task_progress", {}), - "task_language": extras.get("task_language", ""), - "state": record.state, - "robocasa_terminated": record.terminated, - "vla_desync": extras.get("vla_desync", False), - "success": extras.get("success", False), - "log": { - "command": record.command, - "result": record.result, - "elapsed_s": record.elapsed_s, - }, - "images": [], - } - for kind, candidates in ( - ("_image_cam_bytes", _CAMERA_IMAGE_ARTIFACTS["agentview"]), - ("_image_nav_bytes", _CAMERA_IMAGE_ARTIFACTS["navview"]), - ("_image_wrist_bytes", _CAMERA_IMAGE_ARTIFACTS["wrist"]), - ): - for name in candidates: - if name not in record.artifacts: - continue - try: - out[kind] = state.load_bytes(name, step=nn) - except FileNotFoundError: - continue - label = { - "_image_cam_bytes": ("calibration_frame", "agentview"), - "_image_nav_bytes": ("nav_view", "navview"), - "_image_wrist_bytes": ("calibration_frame", "wrist"), - }[kind] - out["images"].append( - { - "role": label[0], - "camera": label[1], - "artifact": name, +@tool +def move_base( + forward: float = 0, + lateral: float = 0, + turn: float = 0, + steps: int = 10, + gripper: float | str = "hold", + *, + ctx: ToolContext[RoboCasaRuntime], +) -> ToolResult: + """Raw base velocity commands in the robot's LOCAL frame. +forward = drive forward, +lateral = strafe right, +turn = rotate CCW (yaw). All values clamped [-1, 1]. Use move_base for fine base adjustments near a target; use navigate_to for long-range navigation. gripper='hold' (DEFAULT) maintains finger width while driving. + + Args: + forward: Forward velocity, [-1, 1] (default 0) + lateral: Lateral / strafe velocity, [-1, 1] (default 0) + turn: Yaw rotation velocity, [-1, 1] (default 0) + steps: Number of env steps (default 10) + gripper: Gripper while driving: +1 close, -1 open, or 'hold' (default 'hold') + """ + runtime = ctx.robot + runtime._vla_desync = True + target_q = float(runtime.env.gripper_qpos[0]) + a = np.zeros(12) + a[11] = 1.0 + a[7:10] = [ + np.clip(forward, -1, 1), + np.clip(lateral, -1, 1), + np.clip(turn, -1, 1), + ] + bp0 = _base_pos(ctx) + for _ in range(steps): + a[6] = _resolve_grip(ctx, gripper, target_q) + _step_env(ctx, a) + bp1 = _base_pos(ctx) + return ToolResult( + data={ + "ok": True, + "base_moved": (bp1 - bp0).tolist(), + "base_pos": bp1.tolist(), + } + ) + + +@tool +def navigate_to( + xy: Annotated[list[float], Field(min_length=2, max_length=2)], + tol: float = 0.20, + max_steps: int = 300, + gripper: float | str = "hold", + *, + ctx: ToolContext[RoboCasaRuntime], +) -> ToolResult: + """Drive the mobile base toward a WORLD (x, y) target. Online-calibrates the base forward-heading, then turns to face + drives forward closed-loop. Holds the arm in place. gripper='hold' (DEFAULT) maintains current finger width while driving (carry-safe). Use tol = expected approach distance + object radius. + + Args: + xy: World-frame target [x, y] in meters (z ignored if provided) + tol: Distance threshold to stop, m (default 0.20) + max_steps: Step budget (default 300) + gripper: Gripper while driving: +1 close, -1 open, or 'hold' (default 'hold') + """ + from scipy.spatial.transform import Rotation as R + + runtime = ctx.robot + runtime._vla_desync = True + target = np.asarray(xy[:2], dtype=np.float64) + target_q = float(runtime.env.gripper_qpos[0]) + if runtime._fwd_offset is None: + _calibrate_forward(ctx, _resolve_grip(ctx, gripper, target_q)) + start = _base_pos(ctx)[:2].copy() + for i in range(max_steps): + bp = _base_pos(ctx) + to = target - bp[:2] + dist = float(np.linalg.norm(to)) + if dist < tol: + runtime._pos_jac = None # base moved -> recalibrate arm + moved = float(np.linalg.norm(bp[:2] - start)) + return ToolResult( + data={ + "ok": True, + "steps": i, + "final_dist": dist, + "moved": moved, + "start_pos": start.tolist(), + "base_pos": bp.tolist(), } ) - break + world_dir = np.arctan2(to[1], to[0]) + yaw = float( + R.from_quat( + np.asarray(runtime.env.current_raw_obs["robot0_base_quat"]) + ).as_euler("xyz")[2] + ) + cur_forward = yaw + runtime._fwd_offset + dyaw = (world_dir - cur_forward + np.pi) % (2 * np.pi) - np.pi + a = np.zeros(12) + a[11] = 1.0 + a[6] = _resolve_grip(ctx, gripper, target_q) + if abs(dyaw) > 0.30: # turn to face the target + a[9] = float(np.sign(dyaw)) + else: # drive forward + small steer + a[7] = 1.0 + a[9] = float(np.clip(dyaw * 1.5, -0.4, 0.4)) + _step_env(ctx, a) + bp = _base_pos(ctx) + runtime._pos_jac = None + moved = float(np.linalg.norm(bp[:2] - start)) + # stuck = ran out of steps having barely moved (rammed a fixture, no path-planning) + return ToolResult( + data={ + "ok": False, + "steps": max_steps, + "final_dist": float(np.linalg.norm(target - bp[:2])), + "moved": moved, + "stuck": moved < 0.12, + "start_pos": start.tolist(), + "base_pos": bp.tolist(), + }, + error="navigate_to did not reach the target within max_steps", + ) - return out +@tool +def reset(*, ctx: ToolContext[RoboCasaRuntime]) -> ToolResult: + """Restart the episode (new layout / object placement sampled). Arm and base calibration are invalidated on reset. DISABLED in no-reset / matched evaluation — the policy must solve the scene in one shot. Only available in EXPLORE mode when RLDX_ALLOW_RESET is enabled.""" + runtime = ctx.robot + if not runtime._allow_reset: + return ToolResult( + error="reset is DISABLED in this run (no-reset/matched evaluation). Solve the scene in one shot; do not restart the episode." + ) + ctx.check_cancelled() + runtime.reset() + return ToolResult( + data={"ok": True, "reset": True, "eef": runtime.env.eef_pos.tolist()} + ) + + +def _run_rldx_skill( + ctx: ToolContext[RoboCasaRuntime], + prompt: str, + base_clip: float | None, + max_chunks: int, + force_reset: bool, + n_action_steps: int, + settle_patience: int, + settle_eps: float, +) -> ToolResult: + """Use the full live task and apply environment budget overrides.""" + runtime = ctx.robot + max_chunks = int(os.environ.get("RLDX_MAX_CHUNKS", max_chunks)) + n_action_steps = int(os.environ.get("RLDX_ACTION_STEPS_PER_CHUNK", n_action_steps)) + settle_patience = int(os.environ.get("RLDX_SETTLE_PATIENCE", settle_patience)) + for name, value in ( + ("max_chunks", max_chunks), + ("n_action_steps", n_action_steps), + ("settle_patience", settle_patience), + ): + if value < 1: + return ToolResult(error=f"{name} must be positive; VLA was not executed") + task_lang = ( + runtime.env.current_raw_obs.get("language") or runtime.env.get_task_language() + ) + if not task_lang: + return ToolResult( + data={"effective_prompt": "", "prompt_overridden": False}, + error="RoboCasa task language is unavailable; VLA was not executed", + ) + ctx.check_cancelled() + reset_history = force_reset or runtime._vla_desync + runtime._vla_desync = False + result = runtime._rldx.run( + task_lang, + max_chunks, + n_action_steps, + base_clip=base_clip, + settle_patience=settle_patience, + settle_eps=settle_eps, + force_reset=reset_history, + step_env=partial(_step_env, ctx), + check_cancelled=ctx.check_cancelled, + ) + result["effective_prompt"] = task_lang + result["effective_max_chunks"] = max_chunks + result["effective_n_action_steps"] = n_action_steps + result["effective_settle_patience"] = settle_patience + result["prompt_overridden"] = prompt != task_lang + if prompt != task_lang: + result["requested_prompt"] = prompt + return ToolResult(data=result) + + +@tool +def rldx_skill( + prompt: str, + base_clip: float | None = None, + max_chunks: int = 70, + force_reset: bool = False, + n_action_steps: int = 8, + settle_patience: int = 999, + settle_eps: float = 0.012, + *, + ctx: ToolContext[RoboCasaRuntime], +) -> ToolResult: + """RLDX VLA closed-loop skill — FULL base motion allowed. The VLA drives both arm and mobile base. Use for full-body tasks where the base must reposition (e.g. navigating to a counter while reaching). Pass the complete live task_language verbatim; the runtime always uses that environment language for RLDX. Do NOT interrupt consecutive VLA calls with manual primitives — that breaks VLA frame history continuity (sets vla_desync=True). + + Args: + prompt: Complete live task_language, copied verbatim + base_clip: Base motion magnitude cap (default null = no clamp) + max_chunks: Action-chunk budget (default 70) + force_reset: Force VLA frame history reset (default False) + n_action_steps: Actions per VLA chunk (default 8) + settle_patience: Settle step budget before declaring done (default 999; do NOT set small) + settle_eps: Settle position tolerance, m (default 0.012) + """ + return _run_rldx_skill( + ctx, + prompt, + base_clip, + max_chunks, + force_reset, + n_action_steps, + settle_patience, + settle_eps, + ) + + +@tool +def rldx_arm( + prompt: str, + base_clip: float | None = 0.1, + max_chunks: int = 70, + force_reset: bool = False, + n_action_steps: int = 8, + settle_patience: int = 999, + settle_eps: float = 0.012, + *, + ctx: ToolContext[RoboCasaRuntime], +) -> ToolResult: + """RLDX VLA closed-loop skill — base CLAMPED to small motions (base_clip=0.1 default). The VLA drives the arm for precise micro-alignment (e.g. fine-tuning a grasp approach) but cannot drive the base away. Pass the complete live task_language verbatim; the runtime always uses that environment language for RLDX. Do NOT interrupt consecutive VLA calls with manual primitives. + + Args: + prompt: Complete live task_language, copied verbatim + base_clip: Base motion magnitude cap (default 0.1 = small) + max_chunks: Action-chunk budget (default 70) + force_reset: Force VLA frame history reset (default False) + n_action_steps: Actions per VLA chunk (default 8) + settle_patience: Settle step budget before declaring done (default 999; do NOT set small) + settle_eps: Settle position tolerance, m (default 0.012) + """ + return _run_rldx_skill( + ctx, + prompt, + base_clip, + max_chunks, + force_reset, + n_action_steps, + settle_patience, + settle_eps, + ) + + +# Camera -> (low-resolution world map, optional high-resolution world map). +_CAMERA_WORLD_ARTIFACTS = { + "agentview": ("agentview_world.npz", "agentview_world_high.npz"), + "navview": ("navview_world.npz", None), + "wrist": ("wrist_world.npz", None), +} + + +@tool +@readonly +def view_env_state( + step: int | None = None, *, ctx: ToolContext[RoboCasaRuntime] +) -> ToolResult: + """Read step NN from states.json + the matching state images in the output dir. If step is null, returns the latest entry. Each entry contains the env state, robocasa_terminated flag, task_progress, vla_desync status, and log. Embeds available PNGs as multimodal image content blocks. Use calibration-frame agentview images for pixel back-projection; use navview for base navigation and floor walkability; use wrist for close-range details near the gripper. + + Args: + step: Step number; 0 = initial. Null = latest. + """ + from robots.robocasa.toolkit import build_observation + try: + record = ctx.state.get(step if step is not None else -1) + except Exception as exc: + return ToolResult(error=f"state step not available: {exc}") + data, images = build_observation(ctx.state, record) + return ToolResult(data=data, images=images) + + +@tool @readonly def back_project_batch( - pixels: list[list[int]], + pixels: Annotated[ + list[Annotated[list[int], Field(min_length=2, max_length=2)]], + Field(min_length=1, max_length=50), + ], step: int | None = None, - camera: str = "agentview", - resolution: str = "low", + camera: Literal["agentview", "navview", "wrist"] = "agentview", + resolution: Literal["high", "low"] = "low", *, - state: EnvState, -) -> dict: - """Back-project multiple pixels to world XYZ in a single call. + ctx: ToolContext[RoboCasaRuntime], +) -> ToolResult: + """Back-project one or more pixels to world XYZ points in a single call. Loads the world map once and queries all pixels. Returns each pixel's world_xyz plus a summary with median_xyz across valid pixels. + + USE THIS for robust object localization: sample 3-8 pixels on the target object and read summary.median_xyz. Maximum 50 pixels per call. - Loads the precomputed world map once and queries all *pixels*, returning - each result individually plus a summary with the median of valid points. + Args: + pixels: List of [row, col] pixel coordinates (max 50) + step: Depth / world-map step to use (default latest). + camera: Camera to back-project from (default agentview). + resolution: Coordinate system for pixels (default low). Use 'low' for the standard 256x256 world map. """ - camera = camera or "agentview" - resolution = resolution or "low" - if camera not in _CAMERA_WORLD_ARTIFACTS: - return {"error": f"bad camera '{camera}' (use agentview, navview, or wrist)"} + state = ctx.state low_name, hi_name = _CAMERA_WORLD_ARTIFACTS[camera] source_artifact = hi_name if resolution == "high" else low_name if source_artifact is None: - return {"error": f"{camera} has no {resolution}-resolution world map"} + return ToolResult(error=f"{camera} has no {resolution}-resolution world map") try: record = state.get(step if step is not None else -1) except Exception as exc: - return {"error": f"state step not available: {exc}"} + return ToolResult(error=f"state step not available: {exc}") nn = record.step_idx if source_artifact not in record.artifacts: - return { - "error": f"{camera} {resolution}-resolution world map not recorded for step {nn}" - } + return ToolResult( + error=f"{camera} {resolution}-resolution world map not recorded for step {nn}" + ) try: world_map = state.load(source_artifact, step=nn) except Exception as exc: - return {"error": f"{source_artifact} not found for step {nn}: {exc}"} + return ToolResult(error=f"{source_artifact} not found for step {nn}: {exc}") results = [] valid_xyzs = [] for pixel in pixels: - if not isinstance(pixel, (list, tuple)) or len(pixel) != 2: - results.append( - { - "pixel": pixel, - "world_xyz": None, - "valid": False, - "error": "pixel must be [row, col]", - } - ) - continue - row, col = int(pixel[0]), int(pixel[1]) + row, col = pixel h, w = world_map.shape[:2] if row < 0 or row >= h or col < 0 or col >= w: results.append( @@ -837,50 +700,67 @@ def back_project_batch( round(float(median[2]), 4), ] - return { - "results": results, - "summary": summary, - "step": nn, - "camera": camera, - "resolution": resolution, - } + return ToolResult( + data={ + "results": results, + "summary": summary, + "step": nn, + "camera": camera, + "resolution": resolution, + } + ) +@tool @readonly def query_world_map( z_min: float = 0.85, z_max: float = 0.95, - x_range: list[float] | None = None, - y_range: list[float] | None = None, - camera: str = "agentview", - resolution: str = "low", + x_range: Annotated[list[float], Field(min_length=2, max_length=2)] | None = None, + y_range: Annotated[list[float], Field(min_length=2, max_length=2)] | None = None, + camera: Literal["agentview", "navview", "wrist"] = "agentview", + resolution: Literal["high", "low"] = "low", min_cluster_size: int = 10, *, - state: EnvState, -) -> dict: - """Query the world map by z-range and/or region to find objects.""" - if camera not in _CAMERA_WORLD_ARTIFACTS: - return {"error": f"bad camera '{camera}' (use agentview, navview, or wrist)"} + ctx: ToolContext[RoboCasaRuntime], +) -> ToolResult: + """Query the world map by Z-range / XY region to find objects at specific heights. Loads the world map once, filters pixels by z_min <= z <= z_max, optionally restricts to x_range / y_range, then clusters contiguous pixels into objects. + + TYPICAL USES: + - z_min=0.85, z_max=0.95 -> countertop-height objects + - z_min=0.0, z_max=0.12, camera='navview' -> walkable floor + - z_min=0.85, z_max=0.95, x_range=[0,2], y_range=[-3,-1] -> counter objects in a specific quadrant + + Args: + z_min: Minimum Z in meters (default 0.85 for counter height). + z_max: Maximum Z in meters (default 0.95 for counter height). + x_range: Optional X range [min, max] in meters; null = no filter. + y_range: Optional Y range [min, max] in meters; null = no filter. + camera: Camera world map to query (default agentview). + resolution: World map resolution (default low). + min_cluster_size: Minimum pixels per cluster to report (default 10). + """ + state = ctx.state low_name, hi_name = _CAMERA_WORLD_ARTIFACTS[camera] source_artifact = hi_name if resolution == "high" else low_name if source_artifact is None: - return {"error": f"{camera} has no {resolution}-resolution world map"} + return ToolResult(error=f"{camera} has no {resolution}-resolution world map") try: record = state.get(-1) except Exception: - return {"error": "no state trace available"} + return ToolResult(error="no state trace available") nn = record.step_idx if source_artifact not in record.artifacts: - return { - "error": f"{camera} {resolution}-resolution world map not found for step {nn}" - } + return ToolResult( + error=f"{camera} {resolution}-resolution world map not found for step {nn}" + ) try: world_map = state.load(source_artifact, step=nn) except Exception: - return { - "error": f"{camera} {resolution}-resolution world map not found for step {nn}" - } + return ToolResult( + error=f"{camera} {resolution}-resolution world map not found for step {nn}" + ) z = world_map[:, :, 2] mask = (z >= z_min) & (z <= z_max) & np.isfinite(z) @@ -894,10 +774,12 @@ def query_world_map( ys, xs = np.where(mask) total_pixels = len(ys) if total_pixels < min_cluster_size: - return { - "clusters": [], - "summary": {"total_clusters": 0, "total_pixels_matched": 0}, - } + return ToolResult( + data={ + "clusters": [], + "summary": {"total_clusters": 0, "total_pixels_matched": 0}, + } + ) h, w = world_map.shape[:2] grid_cells = max(8, min(32, h // 32)) @@ -947,53 +829,47 @@ def query_world_map( ) clusters.sort(key=lambda c: -c["pixel_count"]) - return { - "clusters": clusters[:20], - "summary": { - "total_clusters": len(clusters[:20]), - "total_pixels_matched": total_pixels, - }, - } - + return ToolResult( + data={ + "clusters": clusters[:20], + "summary": { + "total_clusters": len(clusters[:20]), + "total_pixels_matched": total_pixels, + }, + } + ) -@readonly -def finish(status: str, summary: str) -> dict: - """Declare the task finished.""" - return {"_finish": True, "status": status, "summary": summary} - - -# ---- recipe export ---- - -_PRIMITIVE_ACTIONS = frozenset( - { - "move_to", - "move_delta", - "rotate_pitch", - "set_gripper", - "release", - "scripted_grasp", - "rldx_skill", - "rldx_arm", - "navigate_to", - "move_base", - "reset", - } -) +@tool +def finish( + status: Literal["success", "failure", "stuck"], + summary: str, + *, + ctx: ToolContext[RoboCasaRuntime], +) -> ToolResult: + """Declare the task finished. Call when robocasa_terminated becomes True (success detected), or when genuinely stuck after honest exploration. Provide a 1-3 sentence summary of what worked and what failed. -def write_recipe_from_states(state: EnvState, recipe_tag: str) -> str: - """Export non-error RoboCasa primitive commands from the state trace as JSONL.""" - commands = [] - for record in state.records(): - command = record.command - if not isinstance(command, dict): - continue - if command.get("action") not in _PRIMITIVE_ACTIONS: - continue - result = record.result - if isinstance(result, dict) and result.get("error"): - continue - commands.append(command) - recipe_name = f"{recipe_tag}_recipe.jsonl" - state.save(recipe_name, commands, step=None) - return recipe_name + Args: + status: Task outcome classification. + summary: 1-3 sentence summary of what worked / what failed. + """ + return ToolResult(data={"_finish": True, "status": status, "summary": summary}) + + +ROBOCASA_TOOLS = ( + finish, + move_to, + move_delta, + rotate_pitch, + set_gripper, + release, + scripted_grasp, + move_base, + navigate_to, + rldx_skill, + rldx_arm, + reset, + view_env_state, + back_project_batch, + query_world_map, +) diff --git a/robots/robotwin/primitives.py b/robots/robotwin/primitives.py deleted file mode 100644 index fc1b40919..000000000 --- a/robots/robotwin/primitives.py +++ /dev/null @@ -1,428 +0,0 @@ -# Copyright 2026 The RPent Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""RoboTwin primitives built on RLinf environment APIs.""" - -from __future__ import annotations - -from collections.abc import Callable -from typing import Any - -import numpy as np - -from robots.robotwin.env_client import RoboTwinEnvClient -from robots.robotwin.robot_spec import MODEL_SPEC, ROBOTWIN_CAMERA_NAMES -from robots.robotwin.vla_client import LingBotVLAClient - - -def _qmult(left: np.ndarray, right: np.ndarray) -> np.ndarray: - w1, x1, y1, z1 = left - w2, x2, y2, z2 = right - return np.asarray( - [ - w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, - w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, - w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, - w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, - ], - dtype=np.float64, - ) - - -class RoboTwinPrimitives: - """Compose RoboTwin operations from the RLinf environment API.""" - - def __init__( - self, - *, - env: RoboTwinEnvClient, - model: LingBotVLAClient, - seed: int, - check_cancelled: Callable[[], None], - seed_mode: str = "exact", - ): - if seed_mode != "exact": - raise ValueError("standard RoboTwin integration requires seed_mode='exact'") - self.env = env - self.model = model - self.seed = int(seed) - self._check_cancelled = check_cancelled - self.policy_actions = 0 - self.native_actions = 0 - self._recording = False - self._frames: list[np.ndarray] = [] - - def start_recording(self) -> None: - self._recording = True - self._frames = [] - - def record_frame(self, rgb: Any) -> None: - self._frames.append(np.ascontiguousarray(np.asarray(rgb))) - - def recorded_frame_count(self) -> int: - return len(self._frames) - - def stop_recording(self) -> list[np.ndarray]: - frames = list(self._frames) - self._recording = False - self._frames = [] - return frames - - def frame_slice(self, start: int) -> list[np.ndarray]: - return list(self._frames[int(start) :]) - - def reset( - self, - *, - instruction: str | None = None, - feasibility_precheck: bool = True, - ) -> dict[str, Any]: - """Reset the RoboTwin episode and return the native info plus success.""" - del instruction, feasibility_precheck - _, info = self.env.reset() - return {**info, "success": True} - - @staticmethod - def _completion( - *, requested: int, executed: int, status: dict[str, Any] - ) -> dict[str, Any]: - step_lim = status.get("step_lim") - budget_exhausted = step_lim is not None and int( - status.get("take_action_cnt", 0) - ) >= int(step_lim) - completed = executed == requested - if status.get("eval_success") is True: - stop_reason = "native_success" - elif budget_exhausted: - stop_reason = "budget_exhausted" - elif completed: - stop_reason = "completed" - else: - stop_reason = "runtime_failure" - return { - "completed": completed, - "requested_steps": requested, - "executed_steps": executed, - "stop_reason": stop_reason, - } - - def _build_lingbot_observation(self) -> dict[str, Any]: - """Assemble the rgb-only observation the LingBot policy infers on.""" - views = {} - for camera_name in ROBOTWIN_CAMERA_NAMES: - views[camera_name] = { - "rgb": np.asarray(self.env.render_camera(camera_name)) - } - return { - "views": views, - "robot_state": self.env.last_info["robot_state"], - "task_language": self.env.get_task_language(), - } - - @staticmethod - def _validate_qpos_updates_request(updates: Any) -> list[dict[str, Any]]: - if not isinstance(updates, list) or not updates: - raise ValueError("qpos updates must contain at least one update") - normalized = [] - for update in updates: - if not isinstance(update, dict): - raise TypeError("qpos update must be a mapping") - arm = update.get("arm") - if arm not in ("left", "right"): - raise ValueError("arm must be 'left' or 'right'") - if update.get("arm_qpos") is None and update.get("gripper") is None: - raise ValueError("qpos update must set arm_qpos and/or gripper") - item: dict[str, Any] = {"arm": arm} - if update.get("arm_qpos") is not None: - arm_qpos = np.asarray(update["arm_qpos"], dtype=np.float64) - if arm_qpos.shape != (6,) or not np.isfinite(arm_qpos).all(): - raise ValueError("arm_qpos must be finite and have shape (6,)") - item["arm_qpos"] = arm_qpos - if update.get("gripper") is not None: - gripper = float(update["gripper"]) - if not np.isfinite(gripper): - raise ValueError("gripper must be finite") - item["gripper"] = gripper - normalized.append(item) - return normalized - - def apply_qpos_updates( - self, - updates: list[dict[str, Any]], - ) -> dict[str, Any]: - """Compose qpos waypoint actions from the latest robot state.""" - if self.env.terminated or self.env.truncated: - raise RuntimeError("RoboTwin common episode is terminal; reset is required") - updates = self._validate_qpos_updates_request(updates) - executed = 0 - episode_status: dict[str, Any] | None = None - for update in updates: - state = self.env.last_info["robot_state"] - qpos_target14 = np.asarray(state.get("qpos_target14"), dtype=np.float64) - if qpos_target14.shape != (14,) or not np.isfinite(qpos_target14).all(): - raise ValueError( - "RoboTwin robot_state.qpos_target14 must be finite and have " - "shape (14,)" - ) - action = qpos_target14.copy() - offset = 0 if update["arm"] == "left" else 7 - if "arm_qpos" in update: - action[offset : offset + 6] = update["arm_qpos"] - if "gripper" in update: - action[offset + 6] = update["gripper"] - obs, _, _, _, info = self.env.step(action, action_type="qpos") - if self._recording and isinstance(obs, dict) and "main_images" in obs: - self.record_frame(obs["main_images"]) - executed += int(info.get("executed_actions", 0)) - episode_status = info["episode_status"] - if self.env.terminated or self.env.truncated: - break - if episode_status is None: - raise RuntimeError("RoboTwin qpos composition executed no updates") - return { - "action_type": "qpos", - "requested_actions": len(updates), - "executed_actions": executed, - "episode_status": episode_status, - } - - def lingbot_act( - self, *, chunks: int = 4, use_length: int = 50, prompt: str | None = None - ) -> dict[str, Any]: - """Infer and execute up to the requested number of LingBot EEF action chunks.""" - if int(chunks) < 1: - raise ValueError("chunks must be at least 1") - if int(use_length) != MODEL_SPEC.use_length: - raise ValueError( - f"RoboTwin LingBot requires use_length={MODEL_SPEC.use_length}" - ) - executed = 0 - requested = int(chunks) * MODEL_SPEC.use_length - native_prompt = None - for _ in range(int(chunks)): - self._check_cancelled() - status = self.env.last_info["episode_status"] - step_lim = status.get("step_lim") - budget_exhausted = step_lim is not None and int( - status.get("take_action_cnt", 0) - ) >= int(step_lim) - if status.get("eval_success") is True or budget_exhausted: - break - observation = self._build_lingbot_observation() - native_prompt = observation["task_language"] - actions = self.model.infer(observation)[: MODEL_SPEC.use_length] - self._check_cancelled() - payload, _, _, _, info = self.env.chunk_step( - actions, - action_type="ee", - return_all_frames=self._recording - and self.env.execution_capabilities.get("chunk_step_all_frames") - is True, - ) - if self._recording and isinstance(payload, dict) and "frames" in payload: - for frame in payload["frames"]: - self.record_frame(frame) - count = int(info.get("executed_actions", 0)) - executed += count - self.policy_actions += count - self.native_actions += count - self._check_cancelled() - status = self.env.last_info["episode_status"] - return { - **self._completion( - requested=requested, - executed=executed, - status=status, - ), - "success": True, - "prompt": native_prompt, - "agent_prompt_ignored": prompt is not None, - "ignored_agent_prompt": prompt, - "episode_status": status, - } - - def move_to( - self, - *, - arm: str, - xyz: list[float], - quat: list[float] | None = None, - gripper: float | None = None, - substeps: int = 25, - _primitive_name: str = "move_to", - ) -> dict[str, Any]: - """Plan and execute a qpos path to a target end-effector pose.""" - del _primitive_name - if int(substeps) < 0: - raise ValueError("substeps must be non-negative") - self._check_cancelled() - state = self.env.last_info["robot_state"] - if quat is None: - key = "left_eef_pose" if arm == "left" else "right_eef_pose" - quat = np.asarray(state[key], dtype=np.float64)[3:].tolist() - target = np.asarray([*xyz, *quat], dtype=np.float64) - planned = self.env.plan_arm_path(arm, target) - self._check_cancelled() - if planned["status"] != "Success" or planned.get("position") is None: - return { - "completed": False, - "requested_steps": 0, - "executed_steps": 0, - "stop_reason": "plan_failed", - "success": False, - "plan_status": planned["status"], - "hint": "target may be unreachable or in collision", - } - path = np.asarray(planned["position"], dtype=np.float64) - if substeps == 1: - path = path[-1:] - elif substeps >= 2 and len(path) > substeps: - indices = np.linspace(0, len(path) - 1, substeps).astype(int) - path = path[indices] - updates = [ - {"arm": arm, "arm_qpos": waypoint, "gripper": gripper} for waypoint in path - ] - execution = self.apply_qpos_updates(updates) - executed = int(execution.get("executed_actions", 0)) - self.native_actions += executed - self._check_cancelled() - status = execution["episode_status"] - final = self.env.last_info["robot_state"] - key = "left_eef_pose" if arm == "left" else "right_eef_pose" - final_pose = np.asarray(final[key], dtype=np.float64) - return { - **execution, - **self._completion( - requested=len(updates), - executed=executed, - status=status, - ), - "success": True, - "plan_status": planned["status"], - "waypoints": len(path), - "final_eef_xyz": final_pose[:3].tolist(), - "final_dist_m": float( - np.linalg.norm(final_pose[:3] - np.asarray(xyz, dtype=np.float64)) - ), - } - - def rotate_wrist( - self, - *, - arm: str, - delta_yaw_deg: float, - gripper: float | None = None, - substeps: int = 25, - ) -> dict[str, Any]: - """Rotate an arm about world z by the requested yaw delta.""" - state = self.env.last_info["robot_state"] - key = "left_eef_pose" if arm == "left" else "right_eef_pose" - pose = np.asarray(state[key], dtype=np.float64) - yaw = np.deg2rad(float(delta_yaw_deg)) - world_z = np.asarray([np.cos(yaw / 2), 0.0, 0.0, np.sin(yaw / 2)]) - result = self.move_to( - arm=arm, - xyz=pose[:3].tolist(), - quat=_qmult(world_z, pose[3:]).tolist(), - gripper=gripper, - substeps=substeps, - _primitive_name="rotate_wrist", - ) - result["requested_delta_yaw_deg"] = float(delta_yaw_deg) - return result - - def set_gripper( - self, - *, - arm: str, - val: float, - steps: int = 10, - _primitive_name: str = "set_gripper", - ) -> dict[str, Any]: - """Interpolate the gripper command to a target value over multiple steps.""" - del _primitive_name - if int(steps) < 1: - raise ValueError("steps must be at least 1") - self._check_cancelled() - state = self.env.last_info["robot_state"] - current = float(state[f"{arm}_gripper"]) - step_count = int(steps) - target = float(val) - values = [ - current + (target - current) * index / step_count - for index in range(1, step_count + 1) - ] - updates = [{"arm": arm, "gripper": float(value)} for value in values] - execution = self.apply_qpos_updates(updates) - executed = int(execution.get("executed_actions", 0)) - self.native_actions += executed - self._check_cancelled() - status = execution["episode_status"] - now = self.env.last_info["robot_state"] - return { - **execution, - **self._completion( - requested=len(updates), - executed=executed, - status=status, - ), - "success": True, - "gripper_val": float(now[f"{arm}_gripper"]), - } - - def release(self, *, arm: str, val: float = 1.0, steps: int = 10) -> dict[str, Any]: - """Open the gripper to the requested release value.""" - return self.set_gripper( - arm=arm, - val=val, - steps=steps, - _primitive_name="release", - ) - - def status(self) -> dict[str, Any]: - """Return the native episode status plus action counters.""" - return { - **self.env.last_info["episode_status"], - "policy_actions": self.policy_actions, - "native_actions": self.native_actions, - } - - def finish(self, *, status: str, summary: str) -> dict[str, Any]: - """Finish the Planner run and verify success against native episode status.""" - requested_success = status.lower() == "success" - try: - native = self.status() - except Exception as error: # The terminal tool must still stop the Planner. - return { - "_finish": True, - "status": "error", - "summary": summary, - "requested_status": status, - "requested_success": requested_success, - "runtime_error": f"{type(error).__name__}: {error}", - } - verified_success = native.get("eval_success") is True - reported_status = ( - "success" - if verified_success - else ("failure" if requested_success else status) - ) - return { - "_finish": True, - "status": reported_status, - "summary": summary, - "requested_success": requested_success, - "success": verified_success, - "episode_status": native, - } diff --git a/robots/robotwin/robot_spec.py b/robots/robotwin/robot_spec.py index 3d85137b0..97b82ca43 100644 --- a/robots/robotwin/robot_spec.py +++ b/robots/robotwin/robot_spec.py @@ -50,7 +50,7 @@ ) #: Env-side camera names exposed by the RoboTwin EnvServer, in fixed order. -#: Shared across the env client, primitives, and toolkit. +#: Shared across the env client, tools, and toolkit. ROBOTWIN_CAMERA_NAMES = ( "head", "left_wrist", @@ -208,6 +208,7 @@ def get_toolkit( ) return RoboTwinToolkit( runtime_kwargs=runtime_kwargs, + output_dir=config.output_dir, dashboard_events=dashboard_events, memory=memory, ) diff --git a/robots/robotwin/toolkit.py b/robots/robotwin/toolkit.py index 4c5a025fd..f4a0cefe6 100644 --- a/robots/robotwin/toolkit.py +++ b/robots/robotwin/toolkit.py @@ -12,231 +12,147 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""RPent tools for the RLinf RoboTwin robot.""" +"""RoboTwin runtime, tool composition, and observation persistence.""" from __future__ import annotations -from functools import partial +from pathlib import Path from typing import TYPE_CHECKING, Any import numpy as np from robots.robotwin import tools -from robots.robotwin.primitives import RoboTwinPrimitives from robots.robotwin.robot_spec import ROBOTWIN_CAMERA_NAMES -from rpent.dashboard.events import DashboardEventSink -from rpent.session import EnvState -from rpent.tools.toolkit import Toolkit, readonly -from rpent.utils.logging import get_output_dir +from rpent.dashboard.events import DashboardEventSink, StepRecordEvent +from rpent.session import EnvState, StepRecord +from rpent.tools import Toolkit, ToolResult +from rpent.utils.logging import get_logger if TYPE_CHECKING: - from rpent.memory.manager import MemoryManager + from robots.robotwin.env_client import RoboTwinEnvClient + from robots.robotwin.vla_client import LingBotVLAClient + from rpent.memory import MemoryManager -# State-advancing RoboTwin primitives eligible for the recipe. ``reset``, -# ``render``, and read-only tools are intentionally excluded so the recipe -# records only commands that actually move the robot. -_RECIPE_ACTIONS = { - "lingbot_act", - "move_to", - "rotate_wrist", - "set_gripper", - "release", -} +logger = get_logger("robotwin_toolkit") -def _world_from_depth( - depth_metric: np.ndarray, camera_meta: dict[str, Any] -) -> np.ndarray: - """Back-project metric depth into the RoboTwin world frame.""" - depth = np.asarray(depth_metric, dtype=np.float64) - if depth.ndim != 2: - raise ValueError(f"RoboTwin depth must have shape [H,W], got {depth.shape}") - - intrinsic = np.asarray(camera_meta.get("intrinsic_K"), dtype=np.float64) - cam2world = np.asarray(camera_meta.get("cam2world_gl"), dtype=np.float64) - if intrinsic.shape != (3, 3): - raise ValueError("RoboTwin camera intrinsic_K must have shape (3,3)") - if cam2world.shape != (4, 4): - raise ValueError("RoboTwin camera cam2world_gl must have shape (4,4)") - if not np.isfinite(intrinsic).all() or not np.isfinite(cam2world).all(): - raise ValueError("RoboTwin camera calibration must contain only finite values") +class RoboTwinRuntime: + """Clients, fixed episode configuration, and cumulative action counts.""" - height, width = depth.shape - if camera_meta.get("height") != height or camera_meta.get("width") != width: - raise ValueError( - "RoboTwin depth shape does not match camera metadata: " - f"depth={depth.shape}, metadata=" - f"({camera_meta.get('height')}, {camera_meta.get('width')})" - ) - fx, fy = intrinsic[0, 0], intrinsic[1, 1] - cx, cy = intrinsic[0, 2], intrinsic[1, 2] - rows, cols = np.mgrid[0:height, 0:width] - camera_points = np.stack( - [ - (cols - cx) * depth / fx, - -(rows - cy) * depth / fy, - -depth, - ], - axis=-1, - ) - world = camera_points @ cam2world[:3, :3].T + cam2world[:3, 3] - return world.astype(np.float32) + def __init__( + self, + *, + env: RoboTwinEnvClient, + model: LingBotVLAClient, + seed: int, + seed_mode: str = "exact", + ) -> None: + if seed_mode != "exact": + raise ValueError("standard RoboTwin integration requires seed_mode='exact'") + self.env = env + self.model = model + self.seed = int(seed) + self.policy_actions = 0 + self.native_actions = 0 + def status(self) -> dict[str, Any]: + return { + **self.env.last_info["episode_status"], + "policy_actions": self.policy_actions, + "native_actions": self.native_actions, + } -class RoboTwinToolkit(Toolkit): - """Common RPent tools plus RoboTwin primitives.""" - _SPECS = {spec["name"]: spec for spec in tools.TOOLS_SPEC} +class RoboTwinToolkit(Toolkit[RoboTwinRuntime]): + """Native tools and observations for one RoboTwin planner session.""" def __init__( self, *, runtime_kwargs: dict[str, Any], + output_dir: Path | str, dashboard_events: DashboardEventSink, memory: MemoryManager, - ): - state = EnvState(get_output_dir()) + ) -> None: + runtime = RoboTwinRuntime(**runtime_kwargs) + state = EnvState(output_dir) super().__init__( - dashboard_events=dashboard_events, state=state, memory=memory, + robot=runtime, + output_dir=output_dir, + tools=tools.ROBOTWIN_TOOLS, + dashboard_events=dashboard_events, ) - self._latest_status: dict[str, Any] = {} - self._primitives = RoboTwinPrimitives( - check_cancelled=self.raise_if_cancelled, - **runtime_kwargs, - ) - self._primitives.start_recording() - self._action_frame_cursor = self._primitives.recorded_frame_count() - reset_result = { - **self._primitives.env.last_reset_info, - "success": True, - } - self._register_robotwin_tools() - initial = self.get_env_state( - command={"action": "reset"}, - result=reset_result, - elapsed_s=0.0, - ) - record = self._state.latest_record() - if record is not None: - self._publish_step(record) - initial_state = initial.get("state") - if isinstance(initial_state, dict): - self._latest_status = initial_state.get( - "episode_status", self._latest_status - ) - - def _register_robotwin_tools(self) -> None: - self._tools.pop("finish", None) - self.add_tool( - "view_env_state", - self._SPECS["view_env_state"], - partial(tools.view_env_state, state=self._state), - ) - self.add_tool( - "sample_world_xyz", - self._SPECS["sample_world_xyz"], - partial(tools.sample_world_xyz, self._state), - ) - self.add_tool( - "query_world_map", - self._SPECS["query_world_map"], - partial(tools.query_world_map, self._state), + # The env client is already reset to the requested seed during startup. + record = dump_state( + runtime, + state, + log={ + "command": {"action": "reset"}, + "result": {**runtime.env.last_reset_info, "success": True}, + "elapsed_s": 0.0, + }, ) - for name in ( - "render", - "lingbot_act", - "move_to", - "rotate_wrist", - "set_gripper", - "release", - ): - self.add_tool(name, self._SPECS[name], partial(self._step, name)) - self.add_tool("finish", self._SPECS["finish"], self._finish) - - @readonly - def _finish(self, *, status: str, summary: str) -> dict[str, Any]: - return self._primitives.finish(status=status, summary=summary) - - def _capture_full_observation(self) -> dict[str, Any]: - """Assemble the full observation (rgb + depth + camera_meta + world_xyz). - - This is the dump/recording path consumed by ``tools.dump_observation`` - and the ``sample_world_xyz`` agent tool. It deliberately fetches depth - and camera_meta so ``world_xyz`` can be back-projected and saved as an - artifact -- distinct from the rgb-only observation built for LingBot - inference in ``RoboTwinPrimitives._build_lingbot_observation``. - """ - env = self._primitives.env - views: dict[str, dict[str, Any]] = {} - for camera_name in ROBOTWIN_CAMERA_NAMES: - rendered = env.render_camera(camera_name, depth=True) - if not isinstance(rendered, (list, tuple)) or len(rendered) != 2: - raise TypeError( - "RoboTwin render_camera(depth=True) must return (rgb, depth)" - ) - rgb, depth = rendered - camera_meta = env.get_camera_meta(camera_name) - views[camera_name] = { - "rgb": np.asarray(rgb), - "depth": np.asarray(depth, dtype=np.float32), - "world_xyz": _world_from_depth(depth, camera_meta), - "camera_meta": camera_meta, - } - return { - "views": views, - "robot_state": env.last_info["robot_state"], - "task_name": env.server_meta["task_name"], - "task_language": env.get_task_language(), - "depth_unit": "metres", - "world_frame": "world", - } + try: + self._dashboard_events.emit(StepRecordEvent(record=record, env_state=state)) + except Exception: + logger.exception("Dashboard failed to publish step %s", record.step_idx) + self._action_frame_cursor = 0 - def get_env_state( + def _capture_observation( self, *, command: dict[str, Any], - result: dict[str, Any], + result: ToolResult, elapsed_s: float, - ) -> dict[str, Any]: + ) -> tuple[dict[str, Any], list[bytes]]: frame_start = self._action_frame_cursor - self._action_frame_cursor = self._primitives.recorded_frame_count() - status = self._primitives.status() - self._latest_status = status - observation = self._capture_full_observation() - record = tools.dump_observation( - observation, - env_state=self._state, - status=status, - log={ - "command": command, - "result": result, - "elapsed_s": elapsed_s, - }, + self._action_frame_cursor = len(self._frames) + logged_result = result.to_dict() + record = dump_state( + self._robot, + self._state, + log={"command": command, "result": logged_result, "elapsed_s": elapsed_s}, ) if self._dashboard_events.enabled: - frames = self._primitives.frame_slice(frame_start) - if frames: - self._state.save( - f"action_{command['action']}.mp4", - frames, - step=record.step_idx, - fps=20, + try: + frames = self._frames[frame_start:] + if frames: + self._state.save( + f"action_{command['action']}.mp4", + frames, + step=record.step_idx, + fps=20, + ) + except Exception as exc: + logger.warning( + "failed to save action clip for step %s: %s", record.step_idx, exc ) - return tools.view_env_state(record.step_idx, state=self._state) + record = self._state.get(record.step_idx) + data, images = build_observation(self._state, record) + if result.is_error: + data["log"]["result"] = { + key: value for key, value in logged_result.items() if key != "error" + } + data["agent_elapsed_s"] = elapsed_s + return data, images - def close(self) -> None: - """Flush the per-step frame buffer into ``episode.mp4`` (LIBERO parity).""" - frames = self._primitives.stop_recording() - if frames: - self._state.save("episode.mp4", frames, step=None, fps=20) + def solved(self) -> bool: + """Use recorded native success, independently of the requested finish status.""" + record = self._state.latest_record() + return bool( + record is not None + and record.state["episode_status"].get("eval_success") is True + ) - def _step(self, name: str, **kwargs) -> dict[str, Any]: - self.raise_if_cancelled() - if name == "render": - return {"success": True} - return getattr(self._primitives, name)(**kwargs) + def close(self) -> None: + """Save this robot's accumulated episode frames.""" + try: + if self._frames: + self._state.save("episode.mp4", self._frames, step=None, fps=20) + except Exception as exc: + logger.warning("failed to save episode video: %s", exc) def write_recipe(self, recipe_tag: str) -> str: """Export state-advancing RoboTwin primitives with no error and no @@ -258,3 +174,142 @@ def write_recipe(self, recipe_tag: str) -> str: if saved is None: raise RuntimeError(f"failed to save RoboTwin recipe artifact: {name}") return str(self._state.artifact_path(name, step=None)) + + +def _world_from_depth( + depth_metric: np.ndarray, camera_meta: dict[str, Any] +) -> np.ndarray: + """Back-project metric depth into the RoboTwin world frame.""" + depth = np.asarray(depth_metric, dtype=np.float64) + if depth.ndim != 2: + raise ValueError(f"RoboTwin depth must have shape [H,W], got {depth.shape}") + + intrinsic = np.asarray(camera_meta.get("intrinsic_K"), dtype=np.float64) + cam2world = np.asarray(camera_meta.get("cam2world_gl"), dtype=np.float64) + if intrinsic.shape != (3, 3): + raise ValueError("RoboTwin camera intrinsic_K must have shape (3,3)") + if cam2world.shape != (4, 4): + raise ValueError("RoboTwin camera cam2world_gl must have shape (4,4)") + if not np.isfinite(intrinsic).all() or not np.isfinite(cam2world).all(): + raise ValueError("RoboTwin camera calibration must contain only finite values") + + height, width = depth.shape + if camera_meta.get("height") != height or camera_meta.get("width") != width: + raise ValueError( + "RoboTwin depth shape does not match camera metadata: " + f"depth={depth.shape}, metadata=" + f"({camera_meta.get('height')}, {camera_meta.get('width')})" + ) + fx, fy = intrinsic[0, 0], intrinsic[1, 1] + cx, cy = intrinsic[0, 2], intrinsic[1, 2] + rows, cols = np.mgrid[0:height, 0:width] + camera_points = np.stack( + [ + (cols - cx) * depth / fx, + -(rows - cy) * depth / fy, + -depth, + ], + axis=-1, + ) + world = camera_points @ cam2world[:3, :3].T + cam2world[:3, 3] + return world.astype(np.float32) + + +def _artifact_name(view: str, field: str) -> str: + suffix = { + "rgb": ".png", + "depth": ".npy", + "world_xyz": ".npy", + "camera_meta": ".json", + }[field] + return f"{view}_{field}{suffix}" + + +def dump_state( + runtime: RoboTwinRuntime, + env_state: EnvState, + log: dict[str, Any], +) -> StepRecord: + """Persist synchronized RGB, metric depth, and world maps for all three views.""" + env = runtime.env + views = {} + for camera in ROBOTWIN_CAMERA_NAMES: + rgb, depth = env.render_camera(camera, depth=True) + camera_meta = env.get_camera_meta(camera) + views[camera] = { + "rgb": np.asarray(rgb), + "depth": np.asarray(depth, dtype=np.float32), + "world_xyz": _world_from_depth(depth, camera_meta), + "camera_meta": camera_meta, + } + step_idx = 0 if env_state.latest_step is None else env_state.latest_step + 1 + state = { + "step_idx": step_idx, + "task_name": env.server_meta["task_name"], + "task_language": env.get_task_language(), + "robot_state": env.last_info["robot_state"], + "episode_status": runtime.status(), + "artifacts": { + camera: { + field: str( + env_state.artifact_path( + _artifact_name(camera, field), step=step_idx + ) + ) + for field in view + } + for camera, view in views.items() + }, + "view_specs": { + camera: { + "coordinate_space": camera, + "image_shape": list(view["rgb"].shape[:2]), + "pixel_order": "row_col", + } + for camera, view in views.items() + }, + } + with env_state.record_step( + state=state, + terminated=env.terminated, + truncated=env.truncated, + command=log["command"], + result=log["result"], + elapsed_s=log["elapsed_s"], + extras={"task_language": state["task_language"]}, + ) as recorded_step: + for camera, view in views.items(): + for field, value in view.items(): + env_state.save(_artifact_name(camera, field), value, step=recorded_step) + return env_state.get(step_idx) + + +def build_observation( + state: EnvState, record: StepRecord +) -> tuple[dict[str, Any], list[bytes]]: + """Return recorded data and head, left wrist, and right wrist PNGs in order.""" + data = { + "step": record.step_idx, + "terminated": record.terminated, + "truncated": record.truncated, + "state": record.state, + "artifacts": sorted(record.artifacts), + "task_language": record.extras["task_language"], + "log": { + "command": record.command, + "result": record.result, + "elapsed_s": record.elapsed_s, + }, + } + images = [] + for camera in ROBOTWIN_CAMERA_NAMES: + name = _artifact_name(camera, "rgb") + if name in record.artifacts: + try: + images.append(state.load_bytes(name, step=record.step_idx)) + except FileNotFoundError: + pass + return data, images + + +_RECIPE_ACTIONS = {"lingbot_act", "move_to", "rotate_wrist", "set_gripper", "release"} diff --git a/robots/robotwin/tools.py b/robots/robotwin/tools.py index d3058afd4..61f4e561e 100644 --- a/robots/robotwin/tools.py +++ b/robots/robotwin/tools.py @@ -12,33 +12,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""RoboTwin tool schemas and observation artifact helpers.""" +"""Native RoboTwin action and perception tools.""" from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Annotated, Any, Literal import numpy as np +from pydantic import Field -from rpent.session import EnvState, StepRecord -from rpent.tools.toolkit import readonly +from robots.robotwin.robot_spec import MODEL_SPEC, ROBOTWIN_CAMERA_NAMES +from rpent.session import EnvState +from rpent.tools import ToolContext, ToolResult, readonly, tool +if TYPE_CHECKING: + from robots.robotwin.toolkit import RoboTwinRuntime -def _tool_error(code: str, message: str, **details: Any) -> dict[str, Any]: - return { - "success": False, - "error": {"code": code, "message": message, **details}, - } - -def _artifact_name(view: str, field: str) -> str: - suffix = { - "rgb": ".png", - "depth": ".npy", - "world_xyz": ".npy", - "camera_meta": ".json", - }[field] - return f"{view}_{field}{suffix}" +def _tool_error(code: str, message: str, **details: Any) -> ToolResult: + return ToolResult(data={"success": False, "code": code, **details}, error=message) def _load_world_xyz( @@ -46,111 +38,82 @@ def _load_world_xyz( *, view: str, step: int | None, -) -> tuple[dict[str, Any] | None, np.ndarray | None, dict[str, Any] | None]: +) -> tuple[dict[str, Any], np.ndarray] | ToolResult: """Load one persisted agent-visible world map without touching the env.""" - requested_step = -1 if step is None else int(step) + from robots.robotwin.toolkit import _artifact_name + + requested_step = -1 if step is None else step try: record = env_state.get(requested_step) - except Exception: - return ( - None, - None, - _tool_error( - "state_not_found", - "The requested RoboTwin state artifact does not exist.", - step=requested_step, - ), + except (LookupError, ValueError): + return _tool_error( + "state_not_found", + "The requested RoboTwin state artifact does not exist.", + step=requested_step, ) state = record.state - actual_step = record.step_idx - views = state.get("artifacts", {}) + views = state["artifacts"] if view not in views: - return ( - None, - None, - _tool_error( - "view_not_found", - "The requested view is unavailable in this state.", - view=view, - available_views=sorted(views) if isinstance(views, dict) else [], - ), + return _tool_error( + "view_not_found", + "The requested view is unavailable in this state.", + view=view, + available_views=sorted(views), ) world_name = _artifact_name(view, "world_xyz") if world_name not in record.artifacts: - return ( - None, - None, - _tool_error( - "world_xyz_not_found", - "The requested view has no persisted world map.", - view=view, - step_idx=actual_step, - ), + return _tool_error( + "world_xyz_not_found", + "The requested view has no persisted world map.", + view=view, + step_idx=record.step_idx, ) try: - world = env_state.load(world_name, step=actual_step) + world = env_state.load(world_name, step=record.step_idx) except Exception as error: - return ( - None, - None, - _tool_error( - "world_xyz_invalid", - "The persisted world map cannot be read.", - detail=str(error), - ), + return _tool_error( + "world_xyz_invalid", + "The persisted world map cannot be read.", + detail=str(error), ) if world.ndim != 3 or world.shape[2] != 3: - return ( - None, - None, - _tool_error( - "world_xyz_shape", - "A RoboTwin world map must have shape [H,W,3].", - actual_shape=list(world.shape), - ), + return _tool_error( + "world_xyz_shape", + "A RoboTwin world map must have shape [H,W,3].", + actual_shape=list(world.shape), ) - return state, np.asarray(world), None + return state, world +@tool @readonly def sample_world_xyz( - env_state: EnvState, - *, view: str, - pixels: list[list[int]], + pixels: Annotated[ + list[Annotated[list[int], Field(min_length=2, max_length=2)]], + Field(min_length=1, max_length=256), + ], step: int | None = None, - neighborhood: int = 1, -) -> dict[str, Any]: - """Return deterministic median world coordinates around image pixels.""" - state, world, error = _load_world_xyz(env_state, view=view, step=step) - if error is not None: - return error - assert state is not None and world is not None - radius = int(neighborhood) - if radius < 0 or radius > 32: - return _tool_error( - "invalid_neighborhood", - "neighborhood must be an integer from 0 through 32.", - ) - if not isinstance(pixels, list) or not pixels or len(pixels) > 256: - return _tool_error( - "invalid_pixels", - "pixels must contain between 1 and 256 [row,col] pairs.", - ) + neighborhood: Annotated[ + int, Field(ge=0, le=32, json_schema_extra={"default": 1}) + ] = 1, + *, + ctx: ToolContext[RoboTwinRuntime], +) -> ToolResult: + """Read persisted same-frame world xyz around [row,col] pixels. The view is also the pixel coordinate space: use the exact view whose RGB supplied the pixels. The current state's view_specs gives each view's [height,width]. This is read-only and does not render or move the robot. + + Args: + view: Artifact view and pixel coordinate space. It must match the RGB image used to choose pixels. + """ + loaded = _load_world_xyz(ctx.state, view=view, step=step) + if isinstance(loaded, ToolResult): + return loaded + state, world = loaded + radius = neighborhood height, width = world.shape[:2] samples: list[dict[str, Any]] = [] for pixel in pixels: - if ( - not isinstance(pixel, (list, tuple)) - or len(pixel) != 2 - or not all(isinstance(value, (int, np.integer)) for value in pixel) - ): - return _tool_error( - "invalid_pixel", - "Every pixel must be an integer [row,col] pair.", - pixel=pixel, - ) - row, col = (int(pixel[0]), int(pixel[1])) + row, col = pixel if row < 0 or row >= height or col < 0 or col >= width: return _tool_error( "pixel_out_of_bounds", @@ -187,45 +150,45 @@ def sample_world_xyz( "valid_coordinates": finite_counts.tolist(), } ) - return { - "success": True, - "step_idx": state["step_idx"], - "view": view, - "coordinate_space": view, - "image_shape": [height, width], - "pixel_order": "row_col", - "coordinate_order": "xyz", - "frame": "world", - "unit": "metre", - "neighborhood": radius, - "samples": samples, - } + return ToolResult( + data={ + "success": True, + "step_idx": state["step_idx"], + "view": view, + "coordinate_space": view, + "image_shape": [height, width], + "pixel_order": "row_col", + "coordinate_order": "xyz", + "frame": "world", + "unit": "metre", + "neighborhood": radius, + "samples": samples, + } + ) +@tool @readonly def query_world_map( - env_state: EnvState, - *, view: str, - bbox: list[int], + bbox: Annotated[list[int], Field(min_length=4, max_length=4)], step: int | None = None, - max_points: int = 256, -) -> dict[str, Any]: - """Return deterministic row-major samples and statistics for one bbox.""" - state, world, error = _load_world_xyz(env_state, view=view, step=step) - if error is not None: - return error - assert state is not None and world is not None - if ( - not isinstance(bbox, (list, tuple)) - or len(bbox) != 4 - or not all(isinstance(value, (int, np.integer)) for value in bbox) - ): - return _tool_error( - "invalid_bbox", - "bbox must be [row_start,col_start,row_end,col_end].", - ) - row_start, col_start, row_end, col_end = map(int, bbox) + max_points: Annotated[ + int, Field(ge=1, le=4096, json_schema_extra={"default": 256}) + ] = 256, + *, + ctx: ToolContext[RoboTwinRuntime], +) -> ToolResult: + """Read deterministic world-xyz samples from a half-open [row_start,col_start,row_end,col_end] region. The view is also the bbox coordinate space and must match the source RGB artifact; view_specs gives [height,width]. This is read-only. + + Args: + view: Artifact view and bbox coordinate space. It must match the RGB image used to choose the bbox. + """ + loaded = _load_world_xyz(ctx.state, view=view, step=step) + if isinstance(loaded, ToolResult): + return loaded + state, world = loaded + row_start, col_start, row_end, col_end = bbox height, width = world.shape[:2] if not (0 <= row_start < row_end <= height and 0 <= col_start < col_end <= width): return _tool_error( @@ -233,18 +196,13 @@ def query_world_map( "bbox must be a non-empty half-open region inside this view's " "world map. Use the exact artifact view whose RGB supplied the " "bbox coordinates.", - bbox=list(map(int, bbox)), + bbox=bbox, shape=[height, width], view=view, coordinate_space=view, valid_bbox=[0, 0, height, width], ) - limit = int(max_points) - if limit < 1 or limit > 4096: - return _tool_error( - "invalid_max_points", - "max_points must be an integer from 1 through 4096.", - ) + limit = max_points region = world[row_start:row_end, col_start:col_end] valid_mask = np.isfinite(region).all(axis=2) local_rows, local_cols = np.nonzero(valid_mask) @@ -252,7 +210,7 @@ def query_world_map( return _tool_error( "no_valid_world_points", "The requested region contains no finite world coordinates.", - bbox=list(map(int, bbox)), + bbox=bbox, ) xyz = region[local_rows, local_cols] if len(xyz) > limit: @@ -269,338 +227,403 @@ def query_world_map( } for index in indices ] + return ToolResult( + data={ + "success": True, + "step_idx": state["step_idx"], + "view": view, + "coordinate_space": view, + "image_shape": [height, width], + "bbox": [row_start, col_start, row_end, col_end], + "bbox_interval": "half_open", + "pixel_order": "row_col", + "coordinate_order": "xyz", + "frame": "world", + "unit": "metre", + "valid_points": int(len(xyz)), + "returned_points": len(points), + "xyz_min": np.min(xyz, axis=0).tolist(), + "xyz_max": np.max(xyz, axis=0).tolist(), + "xyz_median": np.median(xyz, axis=0).tolist(), + "points": points, + } + ) + + +def _qmult(left: np.ndarray, right: np.ndarray) -> np.ndarray: + w1, x1, y1, z1 = left + w2, x2, y2, z2 = right + return np.asarray( + [ + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, + w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, + w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, + w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, + ], + dtype=np.float64, + ) + + +def _completion( + *, requested: int, executed: int, status: dict[str, Any] +) -> dict[str, Any]: + step_lim = status.get("step_lim") + budget_exhausted = step_lim is not None and int( + status.get("take_action_cnt", 0) + ) >= int(step_lim) + completed = executed == requested + if status.get("eval_success") is True: + stop_reason = "native_success" + elif budget_exhausted: + stop_reason = "budget_exhausted" + elif completed: + stop_reason = "completed" + else: + stop_reason = "runtime_failure" return { - "success": True, - "step_idx": state["step_idx"], - "view": view, - "coordinate_space": view, - "image_shape": [height, width], - "bbox": [row_start, col_start, row_end, col_end], - "bbox_interval": "half_open", - "pixel_order": "row_col", - "coordinate_order": "xyz", - "frame": "world", - "unit": "metre", - "valid_points": int(len(xyz)), - "returned_points": len(points), - "xyz_min": np.min(xyz, axis=0).tolist(), - "xyz_max": np.max(xyz, axis=0).tolist(), - "xyz_median": np.median(xyz, axis=0).tolist(), - "points": points, + "completed": completed, + "requested_steps": requested, + "executed_steps": executed, + "stop_reason": stop_reason, } -def dump_observation( - observation: dict[str, Any], +def _apply_qpos_updates( + ctx: ToolContext[RoboTwinRuntime], updates: list[dict[str, Any]] +) -> dict[str, Any]: + """Compose each waypoint from fresh state and account for it before cancellation.""" + runtime = ctx.robot + env = runtime.env + executed = 0 + for update in updates: + ctx.check_cancelled() + action = np.asarray( + env.last_info["robot_state"]["qpos_target14"], dtype=np.float64 + ).copy() + offset = 0 if update["arm"] == "left" else 7 + if "arm_qpos" in update: + action[offset : offset + 6] = update["arm_qpos"] + if update.get("gripper") is not None: + action[offset + 6] = update["gripper"] + obs, _, _, _, info = env.step(action, action_type="qpos") + count = int(info["executed_actions"]) + executed += count + runtime.native_actions += count + ctx.record_frame(obs["main_images"]) + ctx.check_cancelled() + if env.terminated or env.truncated: + break + return { + "action_type": "qpos", + "requested_actions": len(updates), + "executed_actions": executed, + "episode_status": env.last_info["episode_status"], + } + + +@tool +def lingbot_act( + chunks: Annotated[int, Field(ge=1, json_schema_extra={"default": 4})] = 4, + use_length: Annotated[Literal[50], Field(json_schema_extra={"default": 50})] = 50, + prompt: str | None = None, *, - env_state: EnvState, - status: dict[str, Any], - log: dict[str, Any] | None, -) -> StepRecord: - """Persist one agent-visible observation without simulator oracle state.""" - step_idx = 0 if env_state.latest_step is None else env_state.latest_step + 1 - paths: dict[str, dict[str, str]] = {} - view_specs: dict[str, dict[str, Any]] = {} - for view_name, view in observation["views"].items(): - view_paths: dict[str, str] = {} - if "rgb" in view: - name = _artifact_name(view_name, "rgb") - view_paths["rgb"] = str(env_state.artifact_path(name, step=step_idx)) - for field in ("depth", "world_xyz"): - if field in view: - name = _artifact_name(view_name, field) - view_paths[field] = str(env_state.artifact_path(name, step=step_idx)) - if "camera_meta" in view: - name = _artifact_name(view_name, "camera_meta") - view_paths["camera_meta"] = str( - env_state.artifact_path(name, step=step_idx) - ) - paths[view_name] = view_paths - shape_source = next( - ( - np.asarray(view[field]) - for field in ("rgb", "world_xyz", "depth") - if field in view - ), - None, + ctx: ToolContext[RoboTwinRuntime], +) -> ToolResult: + """Run LingBot-VLA eef16 actions using the native task instruction. The optional prompt is recorded but never sent to the policy.""" + runtime = ctx.robot + env = runtime.env + executed = 0 + requested = chunks * MODEL_SPEC.use_length + native_prompt = None + for _ in range(chunks): + ctx.check_cancelled() + status = env.last_info["episode_status"] + step_lim = status.get("step_lim") + budget_exhausted = step_lim is not None and int( + status["take_action_cnt"] + ) >= int(step_lim) + if status["eval_success"] is True or budget_exhausted: + break + observation = { + "views": { + camera: {"rgb": np.asarray(env.render_camera(camera))} + for camera in ROBOTWIN_CAMERA_NAMES + }, + "robot_state": env.last_info["robot_state"], + "task_language": env.get_task_language(), + } + native_prompt = observation["task_language"] + actions = runtime.model.infer(observation)[: MODEL_SPEC.use_length] + ctx.check_cancelled() + all_frames = env.execution_capabilities.get("chunk_step_all_frames") is True + payload, _, _, _, info = env.chunk_step( + actions, + action_type="ee", + return_all_frames=all_frames, ) - if shape_source is not None and shape_source.ndim >= 2: - view_specs[view_name] = { - "coordinate_space": view_name, - "image_shape": [ - int(shape_source.shape[0]), - int(shape_source.shape[1]), - ], - "pixel_order": "row_col", - } - - state = { - "step_idx": step_idx, - "task_name": observation["task_name"], - "task_language": observation["task_language"], - "robot_state": observation["robot_state"], - "episode_status": status, - "artifacts": paths, - "view_specs": view_specs, - "log": log, - } - eval_success = status.get("eval_success") is True - with env_state.record_step( - state=state, - terminated=eval_success, - truncated=False, - command=(log or {}).get("command"), - result=(log or {}).get("result"), - elapsed_s=(log or {}).get("elapsed_s"), - extras={"task_language": observation.get("task_language")}, - ) as recorded_step: - for view_name, view in observation["views"].items(): - for field in ("rgb", "depth", "world_xyz", "camera_meta"): - if field in view: - env_state.save( - _artifact_name(view_name, field), - view[field], - step=recorded_step, - ) - return env_state.get(step_idx) + count = int(info["executed_actions"]) + executed += count + runtime.policy_actions += count + runtime.native_actions += count + if all_frames: + for frame in payload["frames"]: + ctx.record_frame(frame) + else: + ctx.record_frame(payload["main_images"]) + ctx.check_cancelled() + status = env.last_info["episode_status"] + return ToolResult( + data={ + **_completion(requested=requested, executed=executed, status=status), + "success": True, + "prompt": native_prompt, + "agent_prompt_ignored": prompt is not None, + "ignored_agent_prompt": prompt, + "episode_status": status, + } + ) -@readonly -def view_env_state(step: int = -1, *, state: EnvState) -> dict[str, Any]: - try: - record = state.get(step) - except Exception as error: - return {"error": f"state step not available: {error}"} - result: dict[str, Any] = { - "step": record.step_idx, - "terminated": record.terminated, - "truncated": record.truncated, - "state": record.state, - "artifacts": sorted(record.artifacts), - "task_language": record.extras.get("task_language"), - } - result["log"] = { - "command": record.command, - "result": record.result, - "elapsed_s": record.elapsed_s, - } - for slot, views in ( - ("_image_bytes", ("head",)), - ("_image_cam_bytes", ("left_wrist",)), - ("_image_wrist_bytes", ("right_wrist",)), +def _move_to( + ctx: ToolContext[RoboTwinRuntime], + *, + arm: str, + xyz: list[float], + quat: list[float] | None, + gripper: float | None, + substeps: int, +) -> ToolResult: + ctx.check_cancelled() + env = ctx.robot.env + state = env.last_info["robot_state"] + if quat is None: + quat = np.asarray(state[f"{arm}_eef_pose"], dtype=np.float64)[3:].tolist() + target = np.asarray([*xyz, *quat], dtype=np.float64) + planned = env.plan_arm_path(arm, target) + ctx.check_cancelled() + if ( + planned["status"] != "Success" + or planned.get("position") is None + or not len(planned["position"]) ): - name = next( - ( - _artifact_name(view, "rgb") - for view in views - if _artifact_name(view, "rgb") in record.artifacts - ), - None, + return ToolResult( + data={ + "completed": False, + "requested_steps": 0, + "executed_steps": 0, + "stop_reason": "plan_failed", + "success": False, + "plan_status": planned["status"], + "hint": "target may be unreachable or in collision", + }, + error="Arm path planning failed.", ) - if name is not None: - try: - result[slot] = state.load_bytes(name, step=record.step_idx) - except FileNotFoundError: - pass + path = np.asarray(planned["position"], dtype=np.float64) + if path.ndim != 2 or path.shape[1] != 6: + raise ValueError("RoboTwin planned path must have shape [N,6]") + if substeps == 1: + path = path[-1:] + elif substeps >= 2 and len(path) > substeps: + indices = np.linspace(0, len(path) - 1, substeps).astype(int) + path = path[indices] + updates = [ + {"arm": arm, "arm_qpos": waypoint, "gripper": gripper} for waypoint in path + ] + execution = _apply_qpos_updates(ctx, updates) + final_pose = np.asarray( + env.last_info["robot_state"][f"{arm}_eef_pose"], dtype=np.float64 + ) + return ToolResult( + data={ + **execution, + **_completion( + requested=len(updates), + executed=execution["executed_actions"], + status=execution["episode_status"], + ), + "success": True, + "plan_status": planned["status"], + "waypoints": len(path), + "final_eef_xyz": final_pose[:3].tolist(), + "final_dist_m": float( + np.linalg.norm(final_pose[:3] - np.asarray(xyz, dtype=np.float64)) + ), + } + ) + + +@tool +def move_to( + arm: Literal["left", "right"], + xyz: Annotated[list[float], Field(min_length=3, max_length=3)], + quat: Annotated[list[float], Field(min_length=4, max_length=4)] | None = None, + gripper: float | None = None, + substeps: Annotated[int, Field(ge=0, json_schema_extra={"default": 25})] = 25, + *, + ctx: ToolContext[RoboTwinRuntime], +) -> ToolResult: + """Plan and move one arm to a world-frame xyz and wxyz orientation. The native planner returns qpos waypoints executed with fresh state.""" + return _move_to( + ctx, arm=arm, xyz=xyz, quat=quat, gripper=gripper, substeps=substeps + ) + + +@tool +def rotate_wrist( + arm: Literal["left", "right"], + delta_yaw_deg: float, + gripper: float | None = None, + substeps: Annotated[int, Field(ge=0, json_schema_extra={"default": 25})] = 25, + *, + ctx: ToolContext[RoboTwinRuntime], +) -> ToolResult: + """Rotate one EEF about world Z by a relative angle in degrees.""" + state = ctx.robot.env.last_info["robot_state"] + pose = np.asarray(state[f"{arm}_eef_pose"], dtype=np.float64) + yaw = np.deg2rad(delta_yaw_deg) + world_z = np.asarray([np.cos(yaw / 2), 0.0, 0.0, np.sin(yaw / 2)]) + result = _move_to( + ctx, + arm=arm, + xyz=pose[:3].tolist(), + quat=_qmult(world_z, pose[3:]).tolist(), + gripper=gripper, + substeps=substeps, + ) + result.data["requested_delta_yaw_deg"] = delta_yaw_deg return result -TOOLS_SPEC = [ - { - "name": "view_env_state", - "description": ( - "Read one EnvState step and its synchronized RoboTwin observation " - "artifacts. Step -1 selects the latest entry. Embeds the head, left " - "wrist, and right wrist RGB images when available." - ), - "input_schema": { - "type": "object", - "properties": { - "step": { - "type": "integer", - "default": -1, - "description": "Step number; 0 = initial, -1 = latest.", - } - }, - }, - }, - { - "name": "render", - "description": "Capture a fresh synchronized RoboTwin agent observation.", - "input_schema": {"type": "object", "properties": {}}, - }, - { - "name": "sample_world_xyz", - "description": ( - "Read persisted same-frame world xyz around [row,col] pixels. " - "The view is also the pixel coordinate space: use the exact view " - "whose RGB supplied the pixels. The current state's view_specs " - "gives each view's [height,width]. This is read-only and does not " - "render or move the robot." - ), - "input_schema": { - "type": "object", - "properties": { - "view": { - "type": "string", - "description": ( - "Artifact view and pixel coordinate space. It must match " - "the RGB image used to choose pixels." - ), - }, - "pixels": { - "type": "array", - "minItems": 1, - "maxItems": 256, - "items": { - "type": "array", - "items": {"type": "integer"}, - "minItems": 2, - "maxItems": 2, - }, - }, - "step": {"type": ["integer", "null"]}, - "neighborhood": { - "type": "integer", - "minimum": 0, - "maximum": 32, - "default": 1, - }, - }, - "required": ["view", "pixels"], - }, - }, - { - "name": "query_world_map", - "description": ( - "Read deterministic world-xyz samples from a half-open " - "[row_start,col_start,row_end,col_end] region. The view is also " - "the bbox coordinate space and must match the source RGB artifact; " - "view_specs gives [height,width]. This is read-only." - ), - "input_schema": { - "type": "object", - "properties": { - "view": { - "type": "string", - "description": ( - "Artifact view and bbox coordinate space. It must match " - "the RGB image used to choose the bbox." - ), - }, - "bbox": { - "type": "array", - "items": {"type": "integer"}, - "minItems": 4, - "maxItems": 4, - }, - "step": {"type": ["integer", "null"]}, - "max_points": { - "type": "integer", - "minimum": 1, - "maximum": 4096, - "default": 256, - }, - }, - "required": ["view", "bbox"], - }, - }, - { - "name": "lingbot_act", - "description": ( - "Run LingBot-VLA eef16 actions using the native task instruction. " - "The optional prompt is recorded but never sent to the policy." - ), - "input_schema": { - "type": "object", - "properties": { - "chunks": {"type": "integer", "minimum": 1, "default": 4}, - "use_length": {"type": "integer", "const": 50, "default": 50}, - "prompt": {"type": ["string", "null"]}, - }, - }, - }, - { - "name": "move_to", - "description": ( - "Plan and move one arm to a world-frame xyz and wxyz orientation. " - "The native planner returns qpos waypoints executed with fresh state." - ), - "input_schema": { - "type": "object", - "properties": { - "arm": {"type": "string", "enum": ["left", "right"]}, - "xyz": { - "type": "array", - "items": {"type": "number"}, - "minItems": 3, - "maxItems": 3, - }, - "quat": { - "type": ["array", "null"], - "items": {"type": "number"}, - "minItems": 4, - "maxItems": 4, - }, - "gripper": {"type": ["number", "null"]}, - "substeps": {"type": "integer", "minimum": 0, "default": 25}, - }, - "required": ["arm", "xyz"], - }, - }, - { - "name": "rotate_wrist", - "description": "Rotate one EEF about world Z by a relative angle in degrees.", - "input_schema": { - "type": "object", - "properties": { - "arm": {"type": "string", "enum": ["left", "right"]}, - "delta_yaw_deg": {"type": "number"}, - "gripper": {"type": ["number", "null"]}, - "substeps": {"type": "integer", "minimum": 0, "default": 25}, - }, - "required": ["arm", "delta_yaw_deg"], - }, - }, - { - "name": "set_gripper", - "description": "Linearly move one normalized gripper to val over 10 actions.", - "input_schema": { - "type": "object", - "properties": { - "arm": {"type": "string", "enum": ["left", "right"]}, - "val": {"type": "number", "minimum": 0, "maximum": 1}, - "steps": {"type": "integer", "minimum": 1, "default": 10}, - }, - "required": ["arm", "val"], - }, - }, - { - "name": "release", - "description": "Open one gripper to 1.0 over 10 native actions.", - "input_schema": { - "type": "object", - "properties": { - "arm": {"type": "string", "enum": ["left", "right"]}, - "val": {"type": "number", "default": 1.0}, - "steps": {"type": "integer", "minimum": 1, "default": 10}, - }, - "required": ["arm"], - }, - }, - { - "name": "finish", - "description": ( - "Stop the run. A fresh native status query is authoritative; requesting " - "success cannot override TASK_ENV.eval_success." - ), - "input_schema": { - "type": "object", - "properties": { - "status": {"type": "string"}, - "summary": {"type": "string"}, - }, - "required": ["status", "summary"], - }, - }, -] +def _set_gripper( + ctx: ToolContext[RoboTwinRuntime], *, arm: str, val: float, steps: int +) -> ToolResult: + ctx.check_cancelled() + env = ctx.robot.env + current = float(env.last_info["robot_state"][f"{arm}_gripper"]) + values = [ + current + (val - current) * index / steps for index in range(1, steps + 1) + ] + execution = _apply_qpos_updates( + ctx, [{"arm": arm, "gripper": value} for value in values] + ) + return ToolResult( + data={ + **execution, + **_completion( + requested=steps, + executed=execution["executed_actions"], + status=execution["episode_status"], + ), + "success": True, + "gripper_val": float(env.last_info["robot_state"][f"{arm}_gripper"]), + } + ) + + +@tool +def set_gripper( + arm: Literal["left", "right"], + val: Annotated[float, Field(ge=0, le=1)], + steps: Annotated[int, Field(ge=1, json_schema_extra={"default": 10})] = 10, + *, + ctx: ToolContext[RoboTwinRuntime], +) -> ToolResult: + """Linearly move one normalized gripper to val over the requested number of actions.""" + return _set_gripper(ctx, arm=arm, val=val, steps=steps) + + +@tool +def release( + arm: Literal["left", "right"], + val: Annotated[float, Field(json_schema_extra={"default": 1.0})] = 1.0, + steps: Annotated[int, Field(ge=1, json_schema_extra={"default": 10})] = 10, + *, + ctx: ToolContext[RoboTwinRuntime], +) -> ToolResult: + """Open one gripper to 1.0 over 10 native actions by default.""" + return _set_gripper(ctx, arm=arm, val=val, steps=steps) + + +@tool +def finish( + status: str, summary: str, *, ctx: ToolContext[RoboTwinRuntime] +) -> ToolResult: + """Stop the run using native episode status as authority. Requesting success cannot override TASK_ENV.eval_success. + + Args: + status: Requested task outcome. + summary: Summary of what worked and what failed. + """ + requested_success = status.lower() == "success" + try: + native = ctx.robot.status() + except Exception as error: # The terminal tool must still stop the Planner. + return ToolResult( + data={ + "_finish": True, + "status": "error", + "summary": summary, + "requested_status": status, + "requested_success": requested_success, + "runtime_error": f"{type(error).__name__}: {error}", + } + ) + verified_success = native.get("eval_success") is True + if verified_success: + reported_status = "success" + elif requested_success: + reported_status = "failure" + else: + reported_status = status + return ToolResult( + data={ + "_finish": True, + "status": reported_status, + "summary": summary, + "requested_success": requested_success, + "success": verified_success, + "episode_status": native, + } + ) + + +@tool +@readonly +def view_env_state( + step: Annotated[int, Field(json_schema_extra={"default": -1})] = -1, + *, + ctx: ToolContext[RoboTwinRuntime], +) -> ToolResult: + """Read one EnvState step and its synchronized RoboTwin observation artifacts. Step -1 selects the latest entry. Embeds the head, left wrist, and right wrist RGB images when available. + + Args: + step: Step number; 0 = initial, -1 = latest. + """ + from robots.robotwin.toolkit import build_observation + + try: + record = ctx.state.get(step) + except (LookupError, ValueError) as error: + return ToolResult(error=f"state step not available: {error}") + data, images = build_observation(ctx.state, record) + return ToolResult(data=data, images=images) + + +@tool +def render(*, ctx: ToolContext[RoboTwinRuntime]) -> ToolResult: + """Capture a fresh synchronized RoboTwin agent observation.""" + return ToolResult(data={"success": True}) + + +ROBOTWIN_TOOLS = ( + view_env_state, + render, + sample_world_xyz, + query_world_map, + lingbot_act, + move_to, + rotate_wrist, + set_gripper, + release, + finish, +) diff --git a/rpent/dashboard/server.py b/rpent/dashboard/server.py index cf487cb31..7c5dc4918 100644 --- a/rpent/dashboard/server.py +++ b/rpent/dashboard/server.py @@ -46,11 +46,7 @@ InteractionUnavailableError, UnknownDashboardMessageError, ) -from rpent.dashboard.state import ( - DashboardState, - PrimitiveArgumentError, - PrimitiveConfigError, -) +from rpent.dashboard.state import DashboardState from rpent.utils.daemon import pick_free_port from rpent.utils.logging import get_logger @@ -246,16 +242,10 @@ def api_execute_primitive( tool_result = self._state.execute_primitive(name, arguments) except InteractionUnavailableError as exc: return JSONResponse({"error": str(exc)}, status_code=409) - except PrimitiveArgumentError as exc: - return JSONResponse({"error": str(exc)}, status_code=422) - except PrimitiveConfigError as exc: - logger.error("Dashboard primitive configuration error: %s", exc) - return JSONResponse({"error": str(exc)}, status_code=500) except ValueError as exc: return JSONResponse({"error": str(exc)}, status_code=403) - result = tool_result.result - if isinstance(result, dict) and result.get("error") is not None: - error = " ".join(str(result["error"]).split())[:500] + if tool_result.is_error: + error = " ".join(tool_result.error.split())[:500] return JSONResponse( {"error": error or "primitive execution failed"}, status_code=422, diff --git a/rpent/dashboard/state.py b/rpent/dashboard/state.py index 6e73819ae..43ba1a7be 100644 --- a/rpent/dashboard/state.py +++ b/rpent/dashboard/state.py @@ -24,8 +24,6 @@ from typing import TYPE_CHECKING, Any, Literal import numpy as np -from jsonschema.exceptions import SchemaError, ValidationError -from jsonschema.validators import validator_for from rpent.dashboard.events import ( DashboardEvent, @@ -62,48 +60,6 @@ logger = get_logger("dashboard_state") -class PrimitiveArgumentError(ValueError): - """Dashboard primitive arguments do not satisfy the Toolkit schema.""" - - -class PrimitiveConfigError(RuntimeError): - """A Toolkit primitive has no usable Dashboard input schema.""" - - -def _primitive_validator(name: str, schema: Any) -> Any: - """Return the validator class for one well-formed primitive schema.""" - if not isinstance(schema, dict): - raise PrimitiveConfigError(f"primitive {name!r} has no valid input schema") - validator_type = validator_for(schema) - try: - validator_type.check_schema(schema) - except SchemaError as exc: - raise PrimitiveConfigError( - f"primitive {name!r} has an invalid input schema" - ) from exc - return validator_type - - -def _validate_primitive_arguments( - name: str, - arguments: dict[str, Any], - schema: Any, -) -> None: - """Validate untrusted Dashboard arguments against one Toolkit schema.""" - validator_type = _primitive_validator(name, schema) - - try: - validator_type(schema).validate(arguments) - except ValidationError as exc: - location = "$" + "".join( - f"[{part}]" if isinstance(part, int) else f".{part}" - for part in exc.absolute_path - ) - raise PrimitiveArgumentError( - f"invalid arguments for {name} at {location}: {exc.message}" - ) from exc - - def _to_json_safe(value: Any) -> Any: """Normalize common non-JSON values (numpy / Path) used in Dashboard timeline data.""" if isinstance(value, np.ndarray): @@ -269,20 +225,12 @@ def primitive_specs(self) -> list[dict[str, Any]]: "TaskRun primitives are not available" ) assert toolkit is not None - by_name = {spec.get("name"): spec for spec in toolkit.get_tools_spec()} - primitives = [] - for name in allowlist: - spec = by_name.get(name) - if spec is None: - continue - schema = spec.get("input_schema") - try: - _primitive_validator(name, schema) - except PrimitiveConfigError as exc: - logger.warning("omitting Dashboard primitive: %s", exc) - continue - primitives.append({"name": name, "input_schema": schema}) - return primitives + by_name = {tool.name: tool for tool in toolkit.list_tools()} + return [ + {"name": name, "input_schema": by_name[name].input_schema} + for name in allowlist + if name in by_name + ] def execute_primitive( self, @@ -299,15 +247,8 @@ def execute_primitive( assert toolkit is not None if name not in self._primitive_allowlist: raise ValueError(f"primitive is not allowed: {name}") - available = {spec.get("name"): spec for spec in toolkit.get_tools_spec()} - spec = available.get(name) - if spec is None: + if name not in {tool.name for tool in toolkit.list_tools()}: raise ValueError(f"primitive is not available: {name}") - _validate_primitive_arguments( - name, - arguments, - spec.get("input_schema"), - ) toolkit_key = id(toolkit) self._active_primitive_calls[toolkit_key] = ( self._active_primitive_calls.get(toolkit_key, 0) + 1 diff --git a/rpent/memory/manager.py b/rpent/memory/manager.py index 3ae2d566d..1c8606f85 100644 --- a/rpent/memory/manager.py +++ b/rpent/memory/manager.py @@ -20,12 +20,12 @@ import os import re import shutil -from collections.abc import Callable from pathlib import Path from typing import Any import yaml +from rpent.utils.config import get_repo_root from rpent.utils.logging import get_logger logger = get_logger("memory") @@ -146,45 +146,42 @@ def root(self) -> Path: """Resolved corpus root.""" return self._root - def get_common_tool_bindings( - self, - ) -> dict[str, tuple[dict[str, Any], Callable[..., Any]]]: - """Return memory-aware bindings for shared file tools.""" - from functools import partial - - from rpent.memory import tools as memory_tools - from rpent.tools import common - - handlers = { - "read_text_file": partial( - memory_tools.read_text_file, - memory_root=self._root, - memory_access=self._memory_access, - cell_tag=self._inbox_cell_tag, - ), - "write_text_file": partial( - memory_tools.write_text_file, - memory_root=self._root, - memory_access=self._memory_access, - cell_tag=self._inbox_cell_tag, - ), - "list_dir": partial( - memory_tools.list_dir, - memory_root=self._root, - memory_access=self._memory_access, - cell_tag=self._inbox_cell_tag, - ), - } - bindings: dict[str, tuple[dict[str, Any], Callable[..., Any]]] = {} - for spec in common.TOOLS_SPEC: - name = spec["name"] - handler = handlers.get(name) - if handler is None: - continue - tool_spec = dict(spec) - tool_spec["description"] += memory_tools.MEMORY_BOUNDARY_NOTE - bindings[name] = (tool_spec, handler) - return bindings + def authorize_read(self, path: str | Path) -> Path: + """Resolve a path and enforce the current run's memory read permissions.""" + return self._authorize(path, write=False) + + def authorize_write(self, path: str | Path) -> Path: + """Resolve a path and enforce the current run's memory write permissions.""" + return self._authorize(path, write=True) + + def _authorize(self, path: str | Path, *, write: bool) -> Path: + repo_root = get_repo_root() + resolved = (repo_root / path).resolve() + if not resolved.is_relative_to(self._root): + if resolved.is_relative_to((repo_root / "memory").resolve()): + raise PermissionError( + f"access to another robot's memory is denied: {path}" + ) + return resolved + + parts = resolved.relative_to(self._root).parts + if self._memory_access == "inbox_write" and parts[:3] == ( + "_internal", + "inbox", + self._inbox_cell_tag, + ): + return resolved + if write: + raise PermissionError(f"writing to memory is denied in this mode: {path}") + # Keep published results and root-level Markdown readable, including + # hand-maintained notes without frontmatter. + if ( + not parts + or parts[0] in {"global", "suite", "task_only", "results"} + or (len(parts) == 1 and parts[0].endswith(".md")) + ): + return resolved + raise PermissionError(f"reading this memory path is denied: {path}") def merge_memory( self, diff --git a/rpent/memory/tools.py b/rpent/memory/tools.py deleted file mode 100644 index 90aed8473..000000000 --- a/rpent/memory/tools.py +++ /dev/null @@ -1,167 +0,0 @@ -# Copyright 2026 The RPent Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Memory-aware file-tool handlers for the common tools. - -These wrap the shared IO in :mod:`rpent.tools.common` with two access -checks: robot isolation and run-mode permission. Each handler -only validates the path and checks access, then delegates the actual IO. -""" - -from __future__ import annotations - -from pathlib import Path - -from rpent.tools.toolkit import readonly -from rpent.utils.config import get_repo_root - -# Published memory subtrees an eval run may read. ``results`` retains -# compatibility with the immutable RoboCasa task-only corpus. -_READABLE_SCOPES = {"global", "suite", "task_only", "results"} - -# Suffix appended to shared file-tool descriptions when registered through -# the memory access boundary. -MEMORY_BOUNDARY_NOTE = ( - " Published memory is read-only. During exploration, you may write only " - "to your current memory inbox. Memory for other robots is unavailable." -) - - -def _resolve_memory_path(path: str) -> Path: - """Resolve a file-tool path before memory boundary checks.""" - p = Path(path) - if not p.is_absolute(): - p = get_repo_root() / p - return p.resolve() - - -def _classify_memory_path(resolved: Path, *, memory_root: Path) -> str: - """Classify a resolved path as current, foreign, or non_memory.""" - if resolved.is_relative_to(memory_root): - return "current" - memory_namespace = get_repo_root() / "memory" - if resolved.is_relative_to(memory_namespace): - return "foreign" - return "non_memory" - - -def _check_memory_access( - path: str, - *, - memory_root: Path, - access: str, - memory_access: str, - cell_tag: str | None, -) -> None: - """Enforce current-memory read/write permissions. - - Other robots' memory is always rejected. Within the current robot's memory, - read_only exposes published subtrees read-only and rejects writes; - inbox_write also allows the current cell's inbox. Empty and non-memory - paths are skipped. - """ - if not path: - return - resolved = _resolve_memory_path(path) - bucket = _classify_memory_path(resolved, memory_root=memory_root) - if bucket == "non_memory": - return - if bucket == "foreign": - raise PermissionError(f"access to another robot's memory is denied: {path}") - parts = resolved.relative_to(memory_root).parts - if not parts: - if access == "read": - return - raise PermissionError(f"writing to memory is denied in this mode: {path}") - - top = parts[0] - own_inbox = ( - len(parts) >= 3 - and parts[0] == "_internal" - and parts[1] == "inbox" - and parts[2] == cell_tag - ) - - if access == "write": - if memory_access == "inbox_write" and own_inbox: - return - raise PermissionError(f"writing to memory is denied in this mode: {path}") - - if top in _READABLE_SCOPES or (len(parts) == 1 and top.endswith(".md")): - return - if memory_access == "inbox_write" and own_inbox: - return - raise PermissionError(f"reading this memory path is denied: {path}") - - -@readonly -def read_text_file( - path: str, - *, - memory_root: Path, - memory_access: str, - cell_tag: str | None, - max_chars: int = 40000, -) -> dict: - _check_memory_access( - path, - memory_root=memory_root, - access="read", - memory_access=memory_access, - cell_tag=cell_tag, - ) - from rpent.tools import common - - return common.read_text_file(path, max_chars) - - -@readonly -def write_text_file( - path: str, - content: str, - *, - memory_root: Path, - memory_access: str, - cell_tag: str | None, -) -> dict: - _check_memory_access( - path, - memory_root=memory_root, - access="write", - memory_access=memory_access, - cell_tag=cell_tag, - ) - from rpent.tools import common - - return common.write_text_file(path, content) - - -@readonly -def list_dir( - path: str = "", - *, - memory_root: Path, - memory_access: str, - cell_tag: str | None, -) -> dict: - _check_memory_access( - path, - memory_root=memory_root, - access="read", - memory_access=memory_access, - cell_tag=cell_tag, - ) - from rpent.tools import common - - return common.list_dir(path) diff --git a/rpent/planner/api_loop.py b/rpent/planner/api_loop.py index fee704307..c4d59f07b 100644 --- a/rpent/planner/api_loop.py +++ b/rpent/planner/api_loop.py @@ -17,20 +17,17 @@ The loop wraps the agent's :class:`~rpent.tools.toolkit.Toolkit` as pydantic-ai function tools and drives :class:`pydantic_ai.Agent` runs, streaming each turn so progress is logged in real time. Task completion is -signalled by the robot-provided ``finish`` tool, whose result carries ``_finish``. +signalled by the robot-provided ``finish`` tool, whose accepted result is stored on the toolkit. """ from __future__ import annotations import asyncio -import base64 import contextlib import dataclasses import json import queue from collections import deque -from collections.abc import Callable -from pathlib import Path from typing import Any from pydantic_ai import Agent, BinaryContent, ModelSettings, Tool, ToolReturn @@ -57,8 +54,11 @@ ) from rpent.dashboard.interaction import DashboardInteractionPort, DashboardMessage from rpent.dashboard.planner_control import DashboardPlannerControl -from rpent.planner.base import REASONING_EFFORTS, PlannerResult -from rpent.session import EnvState +from rpent.planner.base import ( + REASONING_EFFORTS, + PlannerResult, + execute_tool, +) from rpent.tools.toolkit import Toolkit from rpent.utils.logging import get_logger @@ -138,7 +138,7 @@ def solve( except asyncio.TimeoutError: toolkit.cancel_active_and_wait() return PlannerResult( - finish_result=None, + finish_result=toolkit.finish_result, messages=[{"role": "user", "content": user_message}], stats={}, error=f"API planner timed out after {self._timeout_s}s", @@ -158,6 +158,7 @@ async def _solve( interactive = input_queue is not None messages: list[dict[str, Any]] = [{"role": "user", "content": user_message}] observer = _ApiRunObserver( + toolkit=toolkit, dashboard_events=self._dashboard_events, messages=messages, max_turns=max_turns, @@ -315,6 +316,7 @@ def emit_user(text: str, *, initial: bool = False) -> None: defer_message_ack=True, ) observer = _ApiRunObserver( + toolkit=toolkit, dashboard_events=self._dashboard_events, messages=messages, max_turns=max_turns, @@ -380,13 +382,16 @@ def _build_agent(self, system_prompt: str, toolkit: Toolkit) -> Agent: class _ApiRunObserver: """Record model/tool events shared by terminal and Dashboard runs.""" + toolkit: Toolkit dashboard_events: DashboardEventSink messages: list[dict[str, Any]] max_turns: int turns: int = 0 tool_calls: int = 0 - finish_result: dict[str, Any] | None = None - pending_finish: dict[str, Any] | None = None + + @property + def finish_result(self) -> dict[str, Any] | None: + return self.toolkit.finish_result def observe_response( self, @@ -425,19 +430,17 @@ def observe_tool(self, event: Any, usage: RunUsage) -> bool: {"type": "tool_call", "tool": part.tool_name, "args": args} ) ) - if part.tool_name == "finish": - self.pending_finish = {"_finish": True, **args} elif isinstance(event, FunctionToolResultEvent): completed = True message = _serialize_tool_result(event) self.messages.append(message) _log_tool_result(message) part = event.part - is_error = bool(getattr(part, "is_error", False)) - if self.pending_finish is not None: - if not is_error and "finish refused" not in str(message): - self.finish_result = self.pending_finish - self.pending_finish = None + metadata = getattr(part, "metadata", None) or {} + is_error = ( + bool(metadata.get("is_error")) + or getattr(part, "outcome", "success") != "success" + ) self.dashboard_events.emit( TranscriptEvent( { @@ -523,6 +526,8 @@ async def interrupt(self) -> int: for message_id in discarded_message_ids: self._control.message_discarded(message_id) if run_task is None or run_task.done(): + if self._observer.finish_result is not None: + self._control.end() return interrupted run_task.cancel() try: @@ -531,6 +536,8 @@ async def interrupt(self) -> int: finally: if self._run_task is run_task: self._run_task = None + if self._observer.finish_result is not None: + self._control.end() return interrupted async def close(self) -> None: @@ -717,144 +724,54 @@ def _api_error_text(error: Exception, *, no_images: bool) -> str: def _build_tools(toolkit: Toolkit, *, no_images: bool = False) -> list[Tool]: - """Build the API-only image reader plus pydantic-ai toolkit wrappers.""" - image_reader = _make_image_reader(toolkit.state, no_images=no_images) - # sequential=True serializes a turn's tool calls so the toolkit's - # single-operation lock never rejects an overlapping call. - tools: list[Tool] = [Tool(image_reader, name="read_image", sequential=True)] - for spec in toolkit.get_tools_spec(): - name = spec["name"] - tools.append( - Tool.from_schema( - function=_make_tool_function(toolkit, name, no_images=no_images), - name=name, - description=spec.get("description", ""), - json_schema=spec.get("input_schema") - or {"type": "object", "properties": {}}, - takes_ctx=False, - sequential=True, - ) + """Expose the frozen native toolset, including the API-only image reader.""" + # Preserve the API's original ordering, with its image reader first. + tools = sorted(toolkit.list_tools(), key=lambda tool: tool.name != "read_image") + return [ + Tool.from_schema( + function=_make_tool_function(toolkit, tool.name, no_images=no_images), + name=tool.name, + description=( + "Acknowledge an image artifact without sending bytes to the model." + if no_images and tool.name == "read_image" + else tool.description + ), + json_schema=tool.input_schema, + takes_ctx=False, + sequential=True, ) - return tools - - -def _make_image_reader( - state: EnvState, - *, - no_images: bool, -) -> Callable[[str, int], ToolReturn | dict[str, str] | str]: - if no_images: - - def read_image_tool(name: str, step: int = -1) -> str: - return read_image_text_only(name, step, state=state) - - read_image_tool.__name__ = "read_image" - read_image_tool.__doc__ = read_image_text_only.__doc__ - return read_image_tool - - def read_image_tool(name: str, step: int = -1) -> ToolReturn | dict[str, str]: - return read_image(name, step, state=state) - - read_image_tool.__name__ = "read_image" - read_image_tool.__doc__ = read_image.__doc__ - return read_image_tool - - -def read_image( - name: str, step: int = -1, *, state: EnvState -) -> ToolReturn | dict[str, str]: - """Read a step-scoped image artifact as visual input. - - Artifact failures are returned as structured tool errors so a bad - model-supplied name or step does not abort the agent run. - """ - try: - resolved_step, path = _resolve_image_artifact(state, name, step) - content = BinaryContent( - data=state.load_bytes(name, step=resolved_step), - media_type=_image_media_type(path), - ) - except Exception as e: - return {"error": str(e)} - return ToolReturn( - return_value={"artifact": name, "step": resolved_step}, - content=[content], - ) - - -def read_image_text_only( - name: str, step: int = -1, *, state: EnvState -) -> str | dict[str, str]: - """Acknowledge an image artifact without sending bytes to the model.""" - try: - resolved_step, _ = _resolve_image_artifact(state, name, step) - except Exception as e: - return {"error": str(e)} - return ( - f"Image artifact {name!r} exists at step {resolved_step}, but image " - "input is disabled (--no-images, text-only model). Reason from textual " - "state instead: view_env_state, back_project, and numeric tool results." - ) - - -def _resolve_image_artifact( - state: EnvState, - name: str, - step: int, -) -> tuple[int, Path]: - record = state.get(step) - path = state.artifact_path(name, step=record.step_idx) - if name not in record.artifacts or not path.is_file(): - raise FileNotFoundError( - f"image artifact {name!r} is not available at step {step}" - ) - if path.suffix.lower() not in {".png", ".jpg", ".jpeg"}: - raise ValueError(f"artifact {name!r} is not an image") - return record.step_idx, path - - -def _image_media_type(path: Path) -> str: - return "image/jpeg" if path.suffix.lower() in {".jpg", ".jpeg"} else "image/png" + for tool in tools + ] def _make_tool_function(toolkit: Toolkit, name: str, *, no_images: bool = False): - """Return a callable that dispatches one tool call to the toolkit.""" - - def _call(**kwargs: Any) -> Any: - result = toolkit.execute_tool(name, kwargs) - text, images = _content_blocks_to_pydantic(result.content_blocks) - if images and not no_images: - return ToolReturn(return_value=text, content=images) + """Dispatch once through the native executor; business errors are results.""" + + async def _call(**kwargs: Any) -> Any: + result = await execute_tool(toolkit, name, kwargs) + text = result.to_text() + images = ( + [] + if no_images + else [ + BinaryContent(data=data, media_type="image/png") + for data in result.images + ] + ) + if images or result.is_error: + # SDK metadata is for the event recorder only. Returning the native + # error text also retains failure observations without SDK retries. + return ToolReturn( + return_value=text, + content=images, + metadata={"is_error": result.is_error}, + ) return text _call.__name__ = name return _call -def _content_blocks_to_pydantic( - blocks: list[dict[str, Any]], -) -> tuple[str, list[BinaryContent]]: - """Split Anthropic-shaped content blocks into text and image content.""" - text_parts: list[str] = [] - images: list[BinaryContent] = [] - for block in blocks: - block_type = block.get("type") - if block_type == "text": - text_parts.append(block.get("text", "")) - elif block_type == "image": - source = block.get("source") or {} - data = source.get("data") - if source.get("type") == "base64" and data: - images.append( - BinaryContent( - data=base64.b64decode(data), - media_type=source.get("media_type", "image/png"), - ) - ) - text = "\n\n".join(part for part in text_parts if part) or "{}" - return text, images - - def _serialize_response(response: ModelResponse) -> dict[str, Any]: """Render one assistant turn as a serialisable transcript message.""" content: list[dict[str, Any]] = [] diff --git a/rpent/planner/base.py b/rpent/planner/base.py index 654a3cd9d..bfd6a8af7 100644 --- a/rpent/planner/base.py +++ b/rpent/planner/base.py @@ -16,6 +16,7 @@ from __future__ import annotations +import asyncio import os import queue from pathlib import Path @@ -23,7 +24,7 @@ from rpent.dashboard.events import DashboardEventSink from rpent.dashboard.interaction import DashboardInteractionPort -from rpent.tools.toolkit import Toolkit +from rpent.tools import Toolkit, ToolResult from rpent.utils.config import ( get_memory_dir, get_repo_root, @@ -50,6 +51,11 @@ def strip_mcp_prefix(name: str) -> str: return name.removeprefix(MCP_TOOL_PREFIX) +async def execute_tool(toolkit: Toolkit, name: str, arguments: dict) -> ToolResult: + """Dispatch native tool execution off the event loop.""" + return await asyncio.to_thread(toolkit.execute_tool, name, arguments) + + class PlannerResult: """Result returned by a planner invocation.""" @@ -99,7 +105,7 @@ def solve( user_message: Initial user message (task description, first steps). toolkit: The full :class:`~rpent.tools.toolkit.Toolkit` (common + robot tools). Backends derive ``tools_spec`` via - ``toolkit.get_tools_spec()`` and dispatch calls via + ``toolkit.list_tools()`` and dispatch calls via ``toolkit.execute_tool()``. max_turns: Maximum LLM turns before giving up. input_queue: Optional queue of user-typed lines for interactive steering. diff --git a/rpent/planner/claude_code.py b/rpent/planner/claude_code.py index 1c2676085..dc6e7e4d1 100644 --- a/rpent/planner/claude_code.py +++ b/rpent/planner/claude_code.py @@ -49,6 +49,7 @@ add_mcp_prefix, strip_mcp_prefix, ) +from rpent.planner.utils.http_mcp_server import build_mcp_server, list_mcp_tools from rpent.tools.toolkit import Toolkit from rpent.utils.config import get_repo_root from rpent.utils.logging import get_logger, init_output_dir @@ -145,6 +146,7 @@ async def _solve_async( output_path.parent.mkdir(parents=True, exist_ok=True) raw_stream_path = output_path.with_suffix(output_path.suffix + ".stream.jsonl") recorder = _Recorder( + toolkit=toolkit, max_turns=max_turns, dashboard_events=self._dashboard_events, ) @@ -343,9 +345,7 @@ def _build_options(self, sdk: Any, *, toolkit: Toolkit, max_turns: int) -> Any: part for part in self._allowed_tools.replace(",", " ").split() if part ] builtins = [name for name in allowed if "__" not in name] - allowed.extend( - add_mcp_prefix(str(spec["name"])) for spec in toolkit.get_tools_spec() - ) + allowed.extend(add_mcp_prefix(tool.name) for tool in list_mcp_tools(toolkit)) thinking = {"type": "disabled"} if self._reasoning_effort == "none" else None effort = None if self._reasoning_effort == "none" else self._reasoning_effort @@ -359,7 +359,6 @@ def _build_options(self, sdk: Any, *, toolkit: Toolkit, max_turns: int) -> Any: allowed_tools=list(dict.fromkeys(allowed)), mcp_servers={ "rpent": _build_rpent_server( - sdk, toolkit=toolkit, ), }, @@ -569,13 +568,13 @@ class _Recorder: ``recorder.error``; transport-level errors are written beside the transcript. """ + toolkit: Toolkit max_turns: int dashboard_events: DashboardEventSink turns: int = 0 _seen_assistant_ids: set[str] = field(default_factory=set) tool_calls: int = 0 tool_names: dict[str, str] = field(default_factory=dict) - pending_finish: dict[str, dict[str, Any]] = field(default_factory=dict) usage: dict[str, int] = field( default_factory=lambda: { "total_input_tokens": 0, @@ -585,7 +584,6 @@ class _Recorder: } ) total_cost_usd: float | None = None - finish_result: dict[str, Any] | None = None error: str | None = None #: Set by the interactive loop before a user-initiated ``interrupt`` so the #: next result (which the CLI may flag ``is_error``) is not mistaken for a @@ -594,6 +592,10 @@ class _Recorder: # -- public ------------------------------------------------------------ + @property + def finish_result(self) -> dict[str, Any] | None: + return self.toolkit.finish_result + def stats(self) -> dict[str, int | float | None]: return { "turns_used": self.turns, @@ -665,8 +667,6 @@ def _assistant(self, message: Any) -> str: name = strip_mcp_prefix(str(_get(block, "name", "tool"))) self.tool_names[tool_id] = name tool_input = _get(block, "input", {}) or {} - if name == "finish" and isinstance(tool_input, dict): - self.pending_finish[tool_id] = dict(tool_input) lines.append(f"[tool->] {name}: {_short_json(tool_input, limit=500)}\n") self.dashboard_events.emit( TranscriptEvent( @@ -714,10 +714,6 @@ def _tool_result_content( summary["images"] = image_count if is_error: summary["is_error"] = bool(is_error) - # Promote the finish payload once the tool result lands successfully. - pending = self.pending_finish.pop(tool_use_id, None) - if pending is not None and not is_error and self.finish_result is None: - self.finish_result = {"_finish": True, **pending} self.dashboard_events.emit( TranscriptEvent( { @@ -788,63 +784,10 @@ def _set_usage(self, usage: Any) -> None: # --------------------------------------------------------------------------- -def _build_rpent_server(sdk: Any, *, toolkit: Toolkit) -> Any: - sdk_tools = [] - tool_execution_lock = asyncio.Lock() - for spec in toolkit.get_tools_spec(): - name = str(spec["name"]) - description = str(spec.get("description", "")) - input_schema = spec.get("input_schema", {"type": "object"}) - - async def run_tool( - args: dict[str, Any], - *, - tool_name: str = name, - ) -> dict[str, Any]: - async with tool_execution_lock: - result = await asyncio.to_thread( - toolkit.execute_tool, - tool_name, - args or {}, - ) - return _tool_result_to_mcp(result) - - run_tool.__name__ = f"rpent_{name}" - sdk_tools.append(sdk.tool(name, description, input_schema)(run_tool)) - - return sdk.create_sdk_mcp_server(name="rpent", version="0.1.0", tools=sdk_tools) - - -def _tool_result_to_mcp(tr: Any) -> dict[str, Any]: - # The toolkit already formatted the result into Anthropic content blocks; - # translate those into the MCP content shape (text + image). - blocks = getattr(tr, "content_blocks", None) - if blocks is None: - return {"content": [{"type": "text", "text": str(tr)}]} - - content: list[dict[str, Any]] = [] - for block in blocks: - block_type = _get(block, "type") - if block_type == "text": - content.append({"type": "text", "text": _get(block, "text", "")}) - elif block_type == "image": - src = _get(block, "source", {}) - content.append( - { - "type": "image", - "data": _get(src, "data", ""), - "mimeType": _get(src, "media_type", "image/png"), - } - ) - - response: dict[str, Any] = {"content": content} - # Surface toolkit-level failures as MCP errors. Without this every result - # looks successful to the SDK, and `_Recorder` promotes a `finish` call to - # the run's finish_result even when the handler rejected it. - result_dict = getattr(tr, "result", None) - if isinstance(result_dict, dict) and result_dict.get("error"): - response["is_error"] = True - return response +def _build_rpent_server(*, toolkit: Toolkit) -> dict[str, Any]: + # A native MCP Server avoids the SDK helper's extra JSON Schema validator, + # which would reject values the Args Model intentionally converts. + return {"type": "sdk", "name": "rpent", "instance": build_mcp_server(toolkit)} # --------------------------------------------------------------------------- diff --git a/rpent/planner/codex.py b/rpent/planner/codex.py index 14eb87ac7..c7c4b9623 100644 --- a/rpent/planner/codex.py +++ b/rpent/planner/codex.py @@ -48,7 +48,11 @@ ) from rpent.dashboard.interaction import DashboardInteractionPort from rpent.dashboard.planner_control import DashboardPlannerControl -from rpent.planner.base import REASONING_EFFORTS, PlannerResult, strip_mcp_prefix +from rpent.planner.base import ( + REASONING_EFFORTS, + PlannerResult, + strip_mcp_prefix, +) from rpent.planner.utils.http_mcp_server import HttpMcpServer from rpent.tools.toolkit import Toolkit from rpent.utils.config import get_repo_root @@ -150,6 +154,7 @@ def solve( ) output_path, raw_stream_path, last_message_path = self._output_paths() recorder = _Recorder( + toolkit=toolkit, max_turns=max_turns, dashboard_events=self._dashboard_events, ) @@ -326,6 +331,7 @@ def _steer() -> None: daemon=True, ).start() + finish_interrupted = False try: for event in turn.stream(): _write_jsonl(raw_f, _message_to_json(event)) @@ -335,6 +341,21 @@ def _steer() -> None: out_f.write(rendered) out_f.flush() logger.info(rendered.strip()) + if ( + recorder.finish_result is not None + and not finish_interrupted + and _get(event, "method") != "turn/completed" + ): + finish_interrupted = True + with contextlib.suppress(Exception): + turn.interrupt() + _write_jsonl( + raw_f, + { + "type": "toolkit_finish", + "finish": recorder.finish_result, + }, + ) finally: if stop_steer is not None: stop_steer.set() @@ -359,7 +380,9 @@ async def _solve_dashboard( """Run a controllable sequence of turns on one Codex thread.""" output_path, raw_stream_path, last_message_path = self._output_paths() recorder = _Recorder( - max_turns=max_turns, dashboard_events=self._dashboard_events + toolkit=toolkit, + max_turns=max_turns, + dashboard_events=self._dashboard_events, ) chunks: list[str] = [] error: str | None = None @@ -431,6 +454,10 @@ def emit_user(text: str, *, initial: bool = False) -> None: logger.warning(cleanup_error) error = error or cleanup_error await session.close() + _write_jsonl( + raw_f, + {"type": "toolkit_finish", "finish": recorder.finish_result}, + ) finally: mcp_server.stop() @@ -554,6 +581,7 @@ async def close(self) -> None: async def _consume_turn(self, turn: Any, done: asyncio.Event) -> None: limit_reached = False + finish_interrupted = False try: async for event in turn.stream(): self._emit_event(event) @@ -578,6 +606,13 @@ async def _consume_turn(self, turn: Any, done: asyncio.Event) -> None: await turn.interrupt() if method != "turn/completed": + if ( + self._recorder.finish_result is not None + and not finish_interrupted + ): + finish_interrupted = True + with contextlib.suppress(Exception): + await turn.interrupt() continue status = _status(_get(payload, "turn")) self._turn = None @@ -615,6 +650,7 @@ async def _consume_turn(self, turn: Any, done: asyncio.Event) -> None: class _Recorder: """Pure adapter: consume Codex SDK events, emit text + accumulate stats.""" + toolkit: Toolkit max_turns: int dashboard_events: DashboardEventSink turns: int = 0 @@ -628,9 +664,12 @@ class _Recorder: } ) final_response: str | None = None - finish_result: dict[str, Any] | None = None error: str | None = None + @property + def finish_result(self) -> dict[str, Any] | None: + return self.toolkit.finish_result + def stats(self) -> dict[str, int]: return {"turns_used": self.turns, "tool_calls": self.tool_calls, **self.usage} @@ -691,7 +730,6 @@ def _render_item(self, item: Any) -> str: self.tool_calls += 1 if item_type in {"mcpToolCall", "dynamicToolCall"}: name = strip_mcp_prefix(str(_get(item, "tool", item_type))) - self._maybe_capture_finish(name, item) elif item_type == "commandExecution": name = str(_get(item, "command", item_type)) else: @@ -751,26 +789,6 @@ def _set_usage(self, usage: Any) -> None: ) ) - def _maybe_capture_finish(self, name: str, item: Any) -> None: - if self.finish_result is not None: - return - if name.lower() != "finish": - return - status = _status(item) - if status and status != "completed": - return - if _get(item, "error") not in (None, ""): - return - data = _jsonable(item) - args = data.get("arguments") if isinstance(data, dict) else None - if isinstance(args, str): - try: - args = json.loads(args) - except Exception: - args = None - if isinstance(args, dict): - self.finish_result = {"_finish": True, **args} - # --------------------------------------------------------------------------- # Codex config overrides diff --git a/rpent/planner/utils/http_mcp_server.py b/rpent/planner/utils/http_mcp_server.py index df4daddd1..5e50825d4 100644 --- a/rpent/planner/utils/http_mcp_server.py +++ b/rpent/planner/utils/http_mcp_server.py @@ -33,6 +33,7 @@ from __future__ import annotations import asyncio +import base64 import socket import threading from typing import Any @@ -43,6 +44,8 @@ from mcp.server.lowlevel import Server from mcp.server.streamable_http_manager import StreamableHTTPSessionManager +from rpent.planner.base import execute_tool +from rpent.tools import Tool, ToolResult from rpent.tools.toolkit import Toolkit from rpent.utils.logging import get_logger @@ -51,33 +54,6 @@ SERVER_NAME = "rpent" -def _toolkit_to_mcp_content( - tr: Any, -) -> tuple[list[types.TextContent | types.ImageContent], bool]: - """Translate a :class:`ToolResult` into MCP content blocks + isError.""" - blocks = getattr(tr, "content_blocks", None) - if blocks is None: - return [types.TextContent(type="text", text=str(tr))], False - - out: list[types.TextContent | types.ImageContent] = [] - for block in blocks: - block_type = block.get("type") - if block_type == "text": - out.append(types.TextContent(type="text", text=block.get("text", ""))) - elif block_type == "image": - src = block.get("source", {}) - out.append( - types.ImageContent( - type="image", - data=src.get("data", ""), - mimeType=src.get("media_type", "image/png"), - ) - ) - result_dict = getattr(tr, "result", None) - is_error = isinstance(result_dict, dict) and bool(result_dict.get("error")) - return out, is_error - - def _strip_mcp_prefix(name: str) -> str: """``mcp__rpent__mcp_list_dir`` -> ``mcp_list_dir`` ; passthrough.""" prefix = f"mcp__{SERVER_NAME}__" @@ -86,36 +62,66 @@ def _strip_mcp_prefix(name: str) -> str: return name -def _build_asgi_app(toolkit: Toolkit) -> Any: - """Build a raw ASGI3 app wrapping an MCP ``Server`` + streamable HTTP.""" +def list_mcp_tools(toolkit: Toolkit) -> tuple[Tool, ...]: + """Codex and Claude use built-in image readers; read_image is API-only.""" + return tuple(tool for tool in toolkit.list_tools() if tool.name != "read_image") + + +def mcp_result(result: ToolResult) -> dict[str, Any]: + """Return MCP text/PNG blocks and the native error flag.""" + return { + "content": [ + {"type": "text", "text": result.to_text()}, + *[ + { + "type": "image", + "data": base64.b64encode(data).decode("ascii"), + "mimeType": "image/png", + } + for data in result.images + ], + ], + "isError": result.is_error, + } + + +def build_mcp_server(toolkit: Toolkit) -> Server: + """Build the MCP service shared by Claude and Codex with native validation.""" mcp_app: Server = Server(SERVER_NAME, version="0.1.0") tool_execution_lock = asyncio.Lock() + exported_tools = list_mcp_tools(toolkit) + exported_names = {tool.name for tool in exported_tools} @mcp_app.list_tools() async def _list_tools() -> list[types.Tool]: tools: list[types.Tool] = [] - for spec in toolkit.get_tools_spec(): + for tool in exported_tools: tools.append( types.Tool( - name=str(spec["name"]), - description=str(spec.get("description", "")), - inputSchema=spec.get("input_schema", {"type": "object"}), + name=tool.name, + description=tool.description, + inputSchema=tool.input_schema, ) ) return tools - @mcp_app.call_tool() + @mcp_app.call_tool(validate_input=False) async def _call_tool(name: str, arguments: dict[str, Any]) -> types.CallToolResult: lookup = _strip_mcp_prefix(name) - async with tool_execution_lock: - tr = await asyncio.get_running_loop().run_in_executor( - None, toolkit.execute_tool, lookup, arguments or {} - ) - content, is_error = _toolkit_to_mcp_content(tr) - return types.CallToolResult(content=content, isError=is_error) + if lookup not in exported_names: + result = ToolResult(error=f"Unknown tool: {lookup}") + else: + async with tool_execution_lock: + result = await execute_tool(toolkit, lookup, arguments or {}) + return types.CallToolResult(**mcp_result(result)) + return mcp_app + + +def _build_asgi_app(toolkit: Toolkit) -> Any: + """Build a raw ASGI3 app wrapping the native MCP service.""" session_manager = StreamableHTTPSessionManager( - app=mcp_app, + app=build_mcp_server(toolkit), stateless=True, json_response=True, ) diff --git a/rpent/session/base.py b/rpent/session/base.py index b93dfa3ba..ec78e484f 100644 --- a/rpent/session/base.py +++ b/rpent/session/base.py @@ -96,7 +96,7 @@ class EnvState: """Own a session'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 = Path(output_dir).resolve() self.reset() # -- private file resolution ----------------------------------------- diff --git a/rpent/tools/__init__.py b/rpent/tools/__init__.py index d7b56d893..ebab5e17f 100644 --- a/rpent/tools/__init__.py +++ b/rpent/tools/__init__.py @@ -12,11 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Agent tool declarations, handlers, and result serialization.""" +"""Native tool protocol and execution.""" -from rpent.tools.toolkit import Toolkit, ToolResult +from rpent.tools.base import ( + Tool, + ToolCancelled, + ToolContext, + ToolResult, + readonly, + tool, +) +from rpent.tools.toolkit import Toolkit __all__ = [ - "Toolkit", + "Tool", + "ToolCancelled", + "ToolContext", "ToolResult", + "Toolkit", + "readonly", + "tool", ] diff --git a/rpent/tools/base.py b/rpent/tools/base.py new file mode 100644 index 000000000..300086ace --- /dev/null +++ b/rpent/tools/base.py @@ -0,0 +1,230 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Native tool protocol.""" + +from __future__ import annotations + +import inspect +import json +import sys +import threading +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path +from types import SimpleNamespace +from typing import ( + TYPE_CHECKING, + Annotated, + Any, + Generic, + ParamSpec, + TypeVar, + get_type_hints, +) + +import numpy as np +from docstring_parser import DocstringStyle, parse +from pydantic import BaseModel, ConfigDict, Field, create_model +from pydantic.json_schema import GenerateJsonSchema +from pydantic_core import core_schema + +if TYPE_CHECKING: + from rpent.memory import MemoryManager + from rpent.session import EnvState + +ParamsT = ParamSpec("ParamsT") +RobotT = TypeVar("RobotT") +MAX_TOOL_TEXT_BYTES = 60000 + + +class _ToolJsonSchema(GenerateJsonSchema): + """Keep the published tool schema independent of Python-only metadata. + + Runtime defaults come from the function signature. Only defaults declared + with Field(json_schema_extra=...) are advertised, matching existing tools. + """ + + def field_title_should_be_set(self, schema: Any) -> bool: + return False + + def model_schema(self, schema: core_schema.ModelSchema) -> dict[str, Any]: + result = super().model_schema(schema) + result.pop("title", None) + result.pop("description", None) + return result + + def default_schema(self, schema: core_schema.WithDefaultSchema) -> dict[str, Any]: + return self.generate_inner(schema["schema"]) + + def dict_schema(self, schema: core_schema.DictSchema) -> dict[str, Any]: + result = super().dict_schema(schema) + # JSON Schema allows arbitrary properties when this keyword is absent. + if result.get("additionalProperties") is True: + result.pop("additionalProperties") + return result + + def nullable_schema(self, schema: core_schema.NullableSchema) -> dict[str, Any]: + inner = self.generate_inner(schema["schema"]) + if isinstance(inner.get("type"), str) and not any( + key in inner for key in ("enum", "const", "allOf", "anyOf", "oneOf", "not") + ): + if inner["type"] != "null": + inner["type"] = [inner["type"], "null"] + return inner + return {"anyOf": [inner, {"type": "null"}]} + + +@dataclass +class ToolResult: + """Tool data, PNG images, and an optional error.""" + + data: dict[str, Any] = field(default_factory=dict) + images: list[bytes] = field(default_factory=list) + error: str | None = None + + @property + def is_error(self) -> bool: + return self.error is not None + + def to_dict(self) -> dict[str, Any]: + """Combine result fields into the public JSON payload.""" + payload = dict(self.data) + if self.error is not None: + payload["error"] = self.error + return payload + + def to_text(self) -> str: + """Encode the original tool payload without truncating internal state.""" + # ASCII output makes character counts equal to UTF-8 byte counts. + text = json.dumps( + self.to_dict(), indent=2, allow_nan=False, ensure_ascii=True, default=str + ) + if len(text) <= MAX_TOOL_TEXT_BYTES: + return text + suffix = "\n[truncated]" + return text[: MAX_TOOL_TEXT_BYTES - len(suffix)] + suffix + + +class ToolCancelled(Exception): + """Raised when a tool reaches a safe cancellation boundary.""" + + +@dataclass(frozen=True) +class ToolContext(Generic[RobotT]): + """References to this run's resources and this invocation's cancellation.""" + + state: EnvState + memory: MemoryManager + robot: RobotT + output_dir: Path + record_frame: Callable[[np.ndarray], None] + _cancel_event: threading.Event = field(repr=False) + + def check_cancelled(self) -> None: + if self._cancel_event.is_set(): + raise ToolCancelled("Tool call cancelled.") + + +@dataclass(frozen=True) +class Tool(Generic[ParamsT]): + """A handler and its generated parameter model, fixed before execution. + + readonly skips automatic observation capture. + """ + + name: str + description: str + args_schema: type[BaseModel] + handler: Callable[ParamsT, ToolResult] + readonly: bool = False + + @property + def input_schema(self) -> dict[str, Any]: + """Generate the model-facing schema from the same model used for validation.""" + return self.args_schema.model_json_schema( + schema_generator=_ToolJsonSchema, + union_format="primitive_type_array", + ) + + +def _parameter_model( + handler: Callable, + descriptions: dict[str, str], + namespace: dict[str, Any], +) -> type[BaseModel]: + """Derive validation from public parameters without evaluating injected types.""" + parameters = inspect.signature(handler).parameters + public = {name: param for name, param in parameters.items() if name != "ctx"} + annotations = get_type_hints( + SimpleNamespace( + __annotations__={name: p.annotation for name, p in public.items()} + ), + globalns=handler.__globals__, + localns=namespace, + include_extras=True, + ) + fields: dict[str, Any] = {} + for name, param in public.items(): + annotation = annotations[name] + if name in descriptions: + annotation = Annotated[annotation, Field(description=descriptions[name])] + default = ... if param.default is inspect.Parameter.empty else param.default + fields[name] = (annotation, default) + return create_model( + f"{handler.__name__}Parameters", + __module__=handler.__module__, + __config__=ConfigDict(extra="forbid", allow_inf_nan=False), + **fields, + ) + + +def readonly(handler: Callable[ParamsT, ToolResult]) -> Callable[ParamsT, ToolResult]: + """Skip automatic environment observation capture; file writes remain allowed. + + Place this marker below @tool. Calls execute one at a time. + """ + setattr(handler, "_rpent_readonly", True) + return handler + + +def tool(function: Callable[ParamsT, ToolResult], /) -> Tool[ParamsT]: + """Declare typed parameters and constraints in the signature, prose in Args. + + The function name becomes the tool name. Google-style docstrings supply the + tool and parameter descriptions. + Every handler declares a required keyword-only ctx, injected by the executor + and excluded from the schema. + Place @tool above @readonly. By default calls are exclusive; robot tools + other than finish capture observations afterward. @readonly disables capture. + """ + # Resolve annotations in factories/tests as well as at module scope, without + # retaining the caller's frame or namespace in the resulting Tool. + namespace = dict(sys._getframe(1).f_locals) + + doc = parse(inspect.getdoc(function), style=DocstringStyle.GOOGLE) + descriptions = { + param.arg_name: param.description + for param in doc.params + if param.description is not None and param.arg_name != "ctx" + } + return Tool( + name=function.__name__, + description="\n\n".join( + part for part in (doc.short_description, doc.long_description) if part + ), + args_schema=_parameter_model(function, descriptions, namespace), + handler=function, + readonly=getattr(function, "_rpent_readonly", False), + ) diff --git a/rpent/tools/common.py b/rpent/tools/common.py deleted file mode 100644 index fec24b6e5..000000000 --- a/rpent/tools/common.py +++ /dev/null @@ -1,168 +0,0 @@ -# Copyright 2026 The RPent Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Common physical agent tools.""" - -from __future__ import annotations - -import os -from pathlib import Path - -from rpent.tools.toolkit import readonly -from rpent.utils.config import get_repo_root -from rpent.utils.logging import get_output_dir - -TOOLS_SPEC: list[dict] = [ - { - "name": "read_text_file", - "description": ( - "Read a UTF-8 text file. Use for past recipe JSONLs, audit JSONs, " - "and memory files. Large files are truncated." - ), - "input_schema": { - "type": "object", - "properties": { - "path": { - "type": "string", - "description": "Absolute or repo-relative path", - }, - "max_chars": { - "type": "integer", - "description": "Max chars (default 40000)", - }, - }, - "required": ["path"], - }, - }, - { - "name": "write_text_file", - "description": ( - "Write a UTF-8 text file (creates parent dirs). Use this to save " - "the working recipe JSONL and the final audit JSON at the end of " - "a successful run." - ), - "input_schema": { - "type": "object", - "properties": { - "path": {"type": "string"}, - "content": {"type": "string"}, - }, - "required": ["path", "content"], - }, - }, - { - "name": "list_dir", - "description": ( - "List files in a directory (non-recursive). Default = {{output_dir}}. " - "Use to inspect the working directory or discover existing resource files." - ), - "input_schema": { - "type": "object", - "properties": { - "path": {"type": "string", "description": "Default: {{output_dir}}"}, - }, - }, - }, - { - "name": "finish", - "description": ( - "Call when the task is complete or unrecoverable. Halts the agent " - "loop. Save any artifacts (recipe, audit) BEFORE calling finish." - ), - "input_schema": { - "type": "object", - "properties": { - "status": { - "type": "string", - "description": "Outcome, e.g. 'success', 'failure', or 'stuck'.", - }, - "summary": { - "type": "string", - "description": "Short natural-language summary of the run.", - }, - }, - "required": ["status", "summary"], - }, - }, -] - - -def _resolve(path: str) -> Path: - p = Path(path) - if not p.is_absolute(): - p = get_repo_root() / p - return p - - -def _truncate(text: str, max_chars: int) -> str: - if len(text) <= max_chars: - return text - return ( - text[:max_chars] - + f"\n\n[TRUNCATED — file is {len(text)} chars, showed first {max_chars}]" - ) - - -# Low-level file IO; planner-facing access is wrapped by MemoryManager. -@readonly -def read_text_file(path: str, max_chars: int = 40000) -> dict: - p = _resolve(path) - if not p.exists(): - return {"error": f"file not found: {p}"} - if p.is_dir(): - return {"error": f"is a directory: {p}"} - try: - text = p.read_text(errors="replace") - except Exception as e: - return {"error": str(e)} - return {"path": str(p), "size": len(text), "content": _truncate(text, max_chars)} - - -@readonly -def write_text_file(path: str, content: str) -> dict: - p = _resolve(path) - p.parent.mkdir(parents=True, exist_ok=True) - p.write_text(content) - return {"path": str(p), "bytes_written": len(content.encode("utf-8"))} - - -@readonly -def list_dir(path: str = "") -> dict: - # Default to the current output dir (so parallel agents see their own). - p = _resolve(path) if path else get_output_dir() - if not p.exists(): - return {"error": f"directory not found: {p}"} - files = sorted(os.listdir(p)) - return {"path": str(p), "count": len(files), "files": files} - - -@readonly -def finish(status: str, summary: str) -> dict: - """Signal that the run is complete. Halts the agent loop. - - The ``_finish`` sentinel is what each planner detects to stop the - tool-calling loop — see ``event.part.tool_name == "finish"`` in - :meth:`rpent.planner.api_loop.ApiAgentLoop._solve` and the - ``pending_finish`` bookkeeping in - :class:`rpent.planner.claude_code._Recorder`. - """ - return {"_finish": True, "status": status, "summary": summary} - - -TOOL_HANDLERS: dict = { - "read_text_file": read_text_file, - "write_text_file": write_text_file, - "list_dir": list_dir, - "finish": finish, -} diff --git a/rpent/tools/common_tools.py b/rpent/tools/common_tools.py new file mode 100644 index 000000000..2cd3488cf --- /dev/null +++ b/rpent/tools/common_tools.py @@ -0,0 +1,104 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Native file and image tools shared by robot toolkits.""" + +from __future__ import annotations + +from typing import Annotated + +from pydantic import Field + +from rpent.tools.base import ToolContext, ToolResult, readonly, tool + + +@tool +@readonly +def read_text_file( + path: str, max_chars: int = 40000, *, ctx: ToolContext +) -> ToolResult: + """Read a UTF-8 text file. Use for past recipe JSONLs, audit JSONs, and memory files. Large files are truncated. Published memory is read-only. During exploration, you may write only to your current memory inbox. Memory for other robots is unavailable. + + Args: + path: Absolute or repo-relative path + max_chars: Max chars (default 40000) + """ + resolved = ctx.memory.authorize_read(path) + text = resolved.read_text(encoding="utf-8", errors="replace") + content = text + if len(text) > max_chars: + content = ( + text[:max_chars] + + f"\n\n[TRUNCATED — file is {len(text)} chars, showed first {max_chars}]" + ) + return ToolResult( + data={"path": str(resolved), "size": len(text), "content": content} + ) + + +@tool +def write_text_file(path: str, content: str, *, ctx: ToolContext) -> ToolResult: + """Write a UTF-8 text file (creates parent dirs). Use this to save the working recipe JSONL and the final audit JSON at the end of a successful run. Published memory is read-only. During exploration, you may write only to your current memory inbox. Memory for other robots is unavailable.""" + resolved = ctx.memory.authorize_write(path) + resolved.parent.mkdir(parents=True, exist_ok=True) + resolved.write_text(content, encoding="utf-8") + return ToolResult( + data={"path": str(resolved), "bytes_written": len(content.encode("utf-8"))} + ) + + +@tool +@readonly +def read_image( + name: str, + step: Annotated[int, Field(ge=-1)] = -1, + *, + ctx: ToolContext, +) -> ToolResult: + """Read a step-scoped image artifact as visual input. + + Artifact failures are returned as structured tool errors so a bad + model-supplied name or step does not abort the agent run. + """ + try: + record = ctx.state.get(step) + path = ctx.state.artifact_path(name, step=record.step_idx) + except Exception as e: + return ToolResult(error=str(e)) + path = ctx.memory.authorize_read(path) + data = {"artifact": name, "step": record.step_idx} + if name not in record.artifacts or not path.is_file(): + return ToolResult( + data={**data}, + error=f"Image artifact {name!r} is not available at step {step}.", + ) + if path.suffix.lower() != ".png": + return ToolResult(data={**data}, error=f"Artifact {name!r} is not a PNG image.") + return ToolResult(data=data, images=[path.read_bytes()]) + + +@tool +@readonly +def list_dir(path: str = "", *, ctx: ToolContext) -> ToolResult: + """List files in a directory (non-recursive). Use to inspect the working directory or discover existing resource files. Published memory is read-only. During exploration, you may write only to your current memory inbox. Memory for other robots is unavailable. + + Args: + path: Directory path. Defaults to the current task's output directory. + """ + resolved = ctx.memory.authorize_read(path or ctx.output_dir) + files = sorted(item.name for item in resolved.iterdir()) + return ToolResult(data={"path": str(resolved), "count": len(files), "files": files}) + + +COMMON_TOOLS = (read_text_file, write_text_file, list_dir, read_image) diff --git a/rpent/tools/toolkit.py b/rpent/tools/toolkit.py index be1aee184..83d8d2720 100644 --- a/rpent/tools/toolkit.py +++ b/rpent/tools/toolkit.py @@ -12,31 +12,37 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Base class for agent tools. - -``Toolkit`` is the agent-facing tool container. Subclasses can register tools -during ``__init__`` via :meth:`Toolkit.add_tool`; the planner calls the tools through :meth:`Toolkit.get_tools_spec` and -:meth:`Toolkit.execute_tool`. -""" +"""Native tool execution with one active invocation per toolkit.""" from __future__ import annotations -import base64 import json import threading import time -import traceback -from collections.abc import Callable from dataclasses import dataclass, field -from functools import partial -from typing import TYPE_CHECKING, Any, ClassVar - -from rpent.dashboard.events import DashboardEventSink, StepRecordEvent -from rpent.utils.templates import substitute - -if TYPE_CHECKING: - from rpent.memory.manager import MemoryManager - from rpent.session import EnvState, StepRecord +from pathlib import Path +from typing import Any, Generic + +import numpy as np +from pydantic import ValidationError + +from rpent.dashboard.events import ( + DashboardEventSink, + NullDashboardEventSink, + StepRecordEvent, +) +from rpent.memory import MemoryManager +from rpent.session import EnvState +from rpent.tools.base import ( + RobotT, + Tool, + ToolContext, + ToolResult, +) +from rpent.tools.common_tools import COMMON_TOOLS +from rpent.utils.logging import get_logger + +logger = get_logger("tools") @dataclass(slots=True) @@ -45,308 +51,164 @@ class _ToolOperation: done_event: threading.Event = field(default_factory=threading.Event) -class ToolCancelled(Exception): - """Raised when an environment reaches a safe cancellation boundary.""" - - -def _truncate_utf8(text: str, max_bytes: int, *, marker: str = "") -> str: - """Truncate text to a valid UTF-8 byte budget, including its marker.""" - encoded = text.encode("utf-8") - if len(encoded) <= max_bytes: - return text - if max_bytes <= 0: - return "" +class Toolkit(Generic[RobotT]): + """A fixed tool collection and its execution resources for one planner session. - marker_bytes = marker.encode("utf-8") - if len(marker_bytes) > max_bytes: - return marker_bytes[:max_bytes].decode("utf-8", errors="ignore") - body = encoded[: max_bytes - len(marker_bytes)].decode( - "utf-8", - errors="ignore", - ) - return body + marker - - -def readonly(func): - """Mark a tool handler as not advancing environment state. - - Tool handlers capture a fresh observation (:meth:`Toolkit.get_env_state`) - by default. Apply this marker to observational and file/IO tools that do - not move the robot or otherwise change the environment. - """ - func._readonly = True - return func - - -def _is_readonly(handler: Callable[..., Any]) -> bool: - """Whether ``handler`` was marked with :func:`readonly`.""" - target = handler - while isinstance(target, partial): - target = target.func - target = getattr(target, "__func__", target) - return bool(getattr(target, "_readonly", False)) - - -@dataclass -class ToolResult: - """Result of executing one tool call. - - Carries the raw result dict (for logging and finish-signal detection) - alongside the Anthropic-shaped content blocks the LLM consumes. - """ - - name: str - result: dict[str, Any] - call_id: str | None = None - - content_blocks: list[dict[str, Any]] = field( - default_factory=list, init=False, repr=False - ) - is_finish: bool = field(default=False, init=False) - - #: Max bytes of the text block emitted in :attr:`content_blocks`. - MAX_TEXT_BYTES_IN_RESULT: ClassVar[int] = 60000 - - def __post_init__(self) -> None: - self.content_blocks = self._build_content_blocks() - self.is_finish = bool( - isinstance(self.result, dict) and self.result.get("_finish") - ) - - def _build_content_blocks(self) -> list[dict[str, Any]]: - """Build Anthropic-shaped content blocks (text + optional images). - - Strips image byte payloads from the text block and emits them as - separate base64 image blocks so the LLM receives the state images as - multimodal content. - """ - result = self.result - if not isinstance(result, dict): - return [ - { - "type": "text", - "text": _truncate_utf8( - str(result), - self.MAX_TEXT_BYTES_IN_RESULT, - ), - } - ] - - result_for_text = dict(result) - image = result_for_text.pop("_image_bytes", None) - image_cam = result_for_text.pop("_image_cam_bytes", None) - image_nav = result_for_text.pop("_image_nav_bytes", None) - image_wrist = result_for_text.pop("_image_wrist_bytes", None) - text = json.dumps(result_for_text, indent=2, default=str) - text = _truncate_utf8( - text, - self.MAX_TEXT_BYTES_IN_RESULT, - marker="\n[truncated]", - ) - - blocks: list[dict[str, Any]] = [{"type": "text", "text": text}] - - def _add_image_bytes(data_bytes: bytes) -> None: - data = base64.b64encode(data_bytes).decode("utf-8") - blocks.append( - { - "type": "image", - "source": { - "type": "base64", - "media_type": "image/png", - "data": data, - }, - } - ) - - if image: - _add_image_bytes(image) - if image_cam: - _add_image_bytes(image_cam) - if image_nav: - _add_image_bytes(image_nav) - if image_wrist: - _add_image_bytes(image_wrist) - return blocks - - -class Toolkit: - """Base toolkit: registers common tools and dispatches tool calls. - - Subclasses extend ``__init__`` (calling ``super().__init__()`` first) - and register additional tools with :meth:`add_tool`. Robot-specific - subclasses receive their env/model/etc. as constructor arguments and - build the underlying env Primitives in ``__init__``; the toolkit - base class only contributes the common file/IO tools. Override - :meth:`close` to release robot-side primitives / servers at the end of the run. + Robot tools submit RGB frames through ctx.record_frame(). Robot toolkits + save their action clips, episode video, and replay recipes. """ def __init__( self, *, - dashboard_events: DashboardEventSink, - state: Any = None, - memory: "MemoryManager", + state: EnvState, + memory: MemoryManager, + robot: RobotT, + output_dir: str | Path, + tools: tuple[Tool, ...], + dashboard_events: DashboardEventSink | None = None, ) -> None: - self._tools: dict[ - str, - tuple[dict[str, Any], Callable[..., Any]], - ] = {} - self._dashboard_events = dashboard_events self._state = state self._memory = memory + self._robot = robot + self._task_output_dir = Path(output_dir).resolve() + self._dashboard_events = dashboard_events or NullDashboardEventSink() + self._tools: dict[str, Tool] = { + item.name: item for item in (*COMMON_TOOLS, *tools) + } self._operation_lock = threading.Lock() self._active_operation: _ToolOperation | None = None - self._register_common_tools() - - # ------------------------------------------------------------------ - # Registration - # ------------------------------------------------------------------ - - def add_tool( - self, - name: str, - spec: dict[str, Any], - handler: Callable[..., Any], - ) -> None: - """Register one tool under ``name`` with its schema and handler. - - Args: - name: Tool name as the LLM sees it (e.g. ``"read_text_file"``). - spec: Anthropic-shaped tool schema dict (``name``, - ``description``, ``input_schema``). - handler: Callable invoked with the tool's input kwargs; returns - a result dict. Decorate read-only handlers with - :func:`readonly`; all other handlers capture state. - """ - self._tools[name] = (spec, handler) - - def _register_common_tools(self) -> None: - """Register the file/IO tools shared by every run.""" - from rpent.tools import common - - memory_bindings = self._memory.get_common_tool_bindings() - for spec in common.TOOLS_SPEC: - name = spec["name"] - binding = memory_bindings.get(name) - if binding is None: - binding = (spec, common.TOOL_HANDLERS[name]) - tool_spec, handler = binding - self.add_tool(name, tool_spec, handler) - - # ------------------------------------------------------------------ - # Planner-facing API - # ------------------------------------------------------------------ - - @property - def memory(self) -> "MemoryManager": - """Return the toolkit's memory manager.""" - return self._memory + self._finish_result: dict[str, Any] | None = None + self._frames: list[np.ndarray] = [] @property def state(self) -> EnvState: - """Return the run's artifact and step store.""" - if self._state is None: - raise RuntimeError("toolkit has no environment state") return self._state - 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()]) - - 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 + @property + def memory(self) -> MemoryManager: + return self._memory + @property + def finish_result(self) -> dict[str, Any] | None: + """Return the accepted finish payload, including robot-specific metadata.""" + result = self._finish_result + return dict(result) if result is not None else None + + def list_tools(self) -> tuple[Tool, ...]: + return tuple(self._tools.values()) + + def record_frame(self, rgb: np.ndarray) -> None: + """Collect one environment-step image for action and episode videos.""" + self._frames.append(np.ascontiguousarray(np.asarray(rgb))) + + def execute_tool(self, name: str, arguments: dict) -> ToolResult: + """Validate and execute one call, then capture its observation.""" + tool = self._tools.get(name) + if tool is None: + return ToolResult(error=f"Unknown tool: {name}") + try: + args = tool.args_schema.model_validate(arguments) + except ValidationError as exc: + errors = exc.errors( + include_url=False, include_context=False, include_input=False + ) + details = json.dumps({"errors": errors}) + return ToolResult(error=f"Invalid arguments for {name}.\n{details}") with self._operation_lock: if self._active_operation is not None: - return ToolResult( - name=name, - result={"error": "another tool operation is still active"}, - ) + return ToolResult(error="another tool operation is still active") operation = _ToolOperation() self._active_operation = operation try: + if operation.cancel_event.is_set(): + return ToolResult(error="Tool call cancelled.") + ctx = ToolContext( + state=self._state, + memory=self._memory, + robot=self._robot, + output_dir=self._task_output_dir, + record_frame=self.record_frame, + _cancel_event=operation.cancel_event, + ) + capture = ( + not tool.readonly and tool not in COMMON_TOOLS and name != "finish" + ) started = time.perf_counter() - failed = False + # Read fields directly so nested models reach the handler intact. + kwargs = {name: getattr(args, name) for name in type(args).model_fields} try: - result = handler(**input_dict) - except TypeError as e: - result = { - "error": f"bad arguments for {name}: {e}", - "got": input_dict, - } - failed = True - 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()} - failed = True - - if not _is_readonly(handler): - 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 + result = tool.handler(**kwargs, ctx=ctx) + except Exception as exc: + logger.exception("Tool %s failed", name) + result = ToolResult(error=str(exc)[:500]) + if capture: + elapsed_s = time.perf_counter() - started + previous = self._state.latest_record() try: - captured = self.get_env_state( - command=command, - result=result_dict, + observation_data, observation_images = self._capture_observation( + command={"action": tool.name, **args.model_dump()}, + result=result, 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: - for key, value in result_dict.items(): - result.setdefault(key, value) - if record is not None: - self._publish_step(record) - - return ToolResult(name=name, result=result) + # The observation replaces action data; retain the action's error. + result.data = observation_data + result.images = result.images + observation_images + except Exception as exc: + logger.exception("State capture failed after %s", tool.name) + error = f"State capture failed: {str(exc)[:500]}" + if result.is_error: + error = f"{result.error}\n{error}" + result.error = error + record = self._state.latest_record() + if record is not None and record is not previous: + try: + self._dashboard_events.emit( + StepRecordEvent(record=record, env_state=self._state) + ) + except Exception: + logger.exception( + "Dashboard failed to publish step %s", record.step_idx + ) + # Images are logged by their owning artifact paths, not their bytes. + try: + result_text = json.dumps( + result.to_dict(), + ensure_ascii=False, + allow_nan=False, + default=str, + ) + except (TypeError, ValueError) as exc: + logger.exception("Tool %s result serialization failed", name) + result = ToolResult( + error=f"Tool result serialization failed: {str(exc)[:500]}", + images=result.images, + ) + result_text = result.to_text() + if name == "finish" and not result.is_error: + self._finish_result = { + key: value for key, value in result.data.items() if key != "_finish" + } + logger.info("Tool %s result: %s", name, result_text) + return 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, - ) - ) - - 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 + def _capture_observation( + self, *, command: dict[str, Any], result: ToolResult, elapsed_s: float + ) -> tuple[dict[str, Any], list[bytes]]: + """Save a step with the action log and return its observation and images. - # ------------------------------------------------------------------ - # Server lifecycle hooks (overridden by robot toolkits) - # ------------------------------------------------------------------ + Returned data replaces the action's data and must include the recorded + step and its artifact names. Include action details in that data where + needed (e.g. log.result). Robot toolkits save action video artifacts. + The executor appends the images and retains the action's error. Raise if + capture fails; an already saved step is still published to the Dashboard. + """ + raise NotImplementedError("This toolkit does not capture robot observations.") def cancel_active_and_wait(self) -> None: """Request cancellation and wait for the active tool to return.""" @@ -357,23 +219,10 @@ def cancel_active_and_wait(self) -> None: operation.cancel_event.set() operation.done_event.wait() - def raise_if_cancelled(self) -> None: - """Raise at an environment-defined safe cancellation boundary.""" - with self._operation_lock: - operation = self._active_operation - if operation is not None and operation.cancel_event.is_set(): - raise ToolCancelled("tool operation interrupted") - def close(self) -> None: - """Release the robot-side primitives / servers at end of run. Default: no-op.""" + """Release robot resources at the end of a run. Default: no-op.""" def solved(self) -> bool: - """Whether the env has reported the task complete. - - Ground truth for the session loop: an agent may call ``finish`` with - ``status="success"`` on a cell it did not actually finish, so the - handoff decision reads the environment, not the agent. - """ raise NotImplementedError def write_recipe(self, recipe_tag: str) -> str | None: diff --git a/tests/README.md b/tests/README.md index 03a2ca9d5..e76b61380 100644 --- a/tests/README.md +++ b/tests/README.md @@ -46,6 +46,10 @@ them when the first test for that module lands. - Place cross-layer tests with the primary contract owner. Registry and config contracts belong to `rpent/robots/`; extension toolkit and schema contracts belong to `robots/`. +- Keep historical tool schema comparisons for all robots and common tools in + `unit_tests/robots/test_tool_schema_contracts.py`. Store each robot's baseline + in `/fixtures/pre_native_tool_contracts.json`; keep runtime behavior + tests in the corresponding robot directory. - Keep one-off fakes in the test module that uses them. Put shared fixtures in the nearest `conftest.py`: use `tests/conftest.py` only for suite-wide fixtures and a module directory's `conftest.py` for local fixtures. diff --git a/tests/e2e_tests/common.py b/tests/e2e_tests/common.py index 10dd43fbf..27173a653 100644 --- a/tests/e2e_tests/common.py +++ b/tests/e2e_tests/common.py @@ -123,7 +123,7 @@ def run_scripted_policy_chain( ) transcript = json.loads(transcript_paths[0].read_text(encoding="utf-8")) finish = transcript.get("finish") - if not isinstance(finish, dict) or finish.get("_finish") is not True: + if not isinstance(finish, dict) or finish.get("status") != "stuck": raise RuntimeError(f"scripted {robot} session did not call finish") state_path = output_dir / "states.json" diff --git a/tests/e2e_tests/robocasa/scenario.py b/tests/e2e_tests/robocasa/scenario.py index c6c7e9ea2..a6c6e4060 100644 --- a/tests/e2e_tests/robocasa/scenario.py +++ b/tests/e2e_tests/robocasa/scenario.py @@ -93,7 +93,7 @@ def _validate_raw_observation(observation: Any) -> dict[str, Any]: def _capture_environment(output_dir: Path, args: Namespace) -> dict[str, Any]: spec = get_robot_spec() with runtime_phase(spec, args, output_dir / "env-capture", {"env"}) as runtime: - env = runtime["env_client"] + env = runtime["env"] raw_observation = env.reset() raw_check = _validate_raw_observation(raw_observation) cameras = {} @@ -127,7 +127,7 @@ def _capture_environment(output_dir: Path, args: Namespace) -> dict[str, Any]: def _rldx_checks(output_dir: Path, args: Namespace) -> dict[str, Any]: spec = get_robot_spec() with runtime_phase(spec, args, output_dir / "rldx", {"vla"}) as runtime: - client = runtime["vla_client"] + client = runtime["model"] modality = client.get_modality_config() frame_count = len(modality["video_delta_indices"]) image = synthetic_rgb() diff --git a/tests/unit_tests/robots/conftest.py b/tests/unit_tests/robots/conftest.py index d3874e244..5d85df9c3 100644 --- a/tests/unit_tests/robots/conftest.py +++ b/tests/unit_tests/robots/conftest.py @@ -21,99 +21,6 @@ import pytest -from rpent.tools.toolkit import readonly - - -class FakeSingleArmPrimitives: - """No-runtime primitive surface shared by LIBERO and RoboCasa tests.""" - - instances: list[FakeSingleArmPrimitives] = [] - - def __init__(self, **kwargs: Any) -> None: - self.kwargs = kwargs - self.reset_calls = 0 - self.recording_started = False - type(self).instances.append(self) - - def reset(self) -> dict[str, Any]: - self.reset_calls += 1 - return {"success": True} - - def reset_episode(self, reason: str) -> dict[str, Any]: - return {"success": True, "reason": reason} - - def start_recording(self) -> None: - self.recording_started = True - - def recorded_frame_count(self) -> int: - return 0 - - def frame_slice(self, start: int) -> list[Any]: - del start - return [] - - def stop_recording(self) -> list[Any]: - return [] - - def dump_success_criteria(self) -> str: - return "offline success criteria" - - @readonly - def segment(self, **kwargs: Any) -> dict[str, Any]: - return {"segment": kwargs} - - @staticmethod - def _operation(name: str, **kwargs: Any) -> dict[str, Any]: - return {"operation": name, "arguments": kwargs} - - def move_to(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("move_to", **kwargs) - - def pi0_pick(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("pi0_pick", **kwargs) - - def pi0_doubled(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("pi0_doubled", **kwargs) - - def release(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("release", **kwargs) - - def set_gripper(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("set_gripper", **kwargs) - - def rotate_wrist(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("rotate_wrist", **kwargs) - - def rotate_pitch(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("rotate_pitch", **kwargs) - - def move_pose(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("move_pose", **kwargs) - - def move_delta(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("move_delta", **kwargs) - - def scripted_grasp(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("scripted_grasp", **kwargs) - - def rldx_skill(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("rldx_skill", **kwargs) - - def rldx_arm(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("rldx_arm", **kwargs) - - def navigate_to(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("navigate_to", **kwargs) - - def move_base(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("move_base", **kwargs) - - -@pytest.fixture -def fake_single_arm_primitives() -> type[FakeSingleArmPrimitives]: - FakeSingleArmPrimitives.instances.clear() - return FakeSingleArmPrimitives - def _fake_module(name: str, **attrs: Any) -> types.ModuleType: module = types.ModuleType(name) diff --git a/tests/unit_tests/robots/dual_franka/fixtures/pre_native_tool_contracts.json b/tests/unit_tests/robots/dual_franka/fixtures/pre_native_tool_contracts.json new file mode 100644 index 000000000..e302b853e --- /dev/null +++ b/tests/unit_tests/robots/dual_franka/fixtures/pre_native_tool_contracts.json @@ -0,0 +1,527 @@ +{ + "source_commit": "bd9040aa99b8fc75433b424905d832fecbef270f", + "schemas": { + "describe_dual_franka_setup": { + "type": "object", + "properties": {} + }, + "view_env_state": { + "type": "object", + "properties": { + "step": { + "type": "integer", + "default": -1 + } + } + }, + "view_camera_meta": { + "type": "object", + "properties": { + "step": { + "type": "integer", + "default": -1 + } + } + }, + "back_project": { + "type": "object", + "properties": { + "camera": { + "type": "string", + "default": "d455", + "description": "Registered projection view name, e.g. d455 or base. Valid names come from perception.projection_views and the current state's saved artifacts." + }, + "row": { + "type": "integer", + "minimum": 0 + }, + "col": { + "type": "integer", + "minimum": 0 + }, + "target_name": { + "type": "string", + "default": "target" + }, + "step": { + "type": "integer" + }, + "window_radius": { + "type": "integer", + "minimum": 0, + "default": 2 + } + }, + "required": [ + "row", + "col" + ] + }, + "segment": { + "type": "object", + "properties": { + "camera": { + "type": "string", + "default": "d455", + "description": "Registered projection view name, e.g. d455 or base. Valid names come from perception.projection_views and the current state's saved artifacts." + }, + "prompt": { + "type": "string", + "default": "", + "description": "Text prompt for SAM3. Prefer short object/relation phrases; for the clean-desk box use 'white interior of the black cardboard box' or 'cardboard box'. Avoid over-specific surface words such as 'floor' when grounding is weak. Provide exactly one of prompt or point." + }, + "point": { + "type": "array", + "items": { + "type": "integer" + }, + "minItems": 2, + "maxItems": 2, + "description": "Positive SAM3 point in camera image coordinates [row, col]. Provide exactly one of prompt or point." + }, + "target_name": { + "type": "string", + "default": "target" + }, + "step": { + "type": "integer" + }, + "min_score": { + "type": "number", + "minimum": 0.0, + "maximum": 1.0, + "default": 0.2 + }, + "min_valid_depth_pixels": { + "type": "integer", + "minimum": 1, + "default": 25 + } + } + }, + "move_delta": { + "type": "object", + "properties": { + "arm": { + "type": "string", + "enum": [ + "left", + "right" + ], + "description": "Which arm to command; the other arm is left uncommanded." + }, + "delta_xyz": { + "type": "array", + "items": { + "type": "number" + }, + "minItems": 3, + "maxItems": 3 + } + }, + "required": [ + "arm", + "delta_xyz" + ] + }, + "rotate_delta": { + "type": "object", + "properties": { + "arm": { + "type": "string", + "enum": [ + "left", + "right" + ], + "description": "Which arm to command; the other arm is left uncommanded." + }, + "delta_rpy": { + "type": "array", + "items": { + "type": "number" + }, + "minItems": 3, + "maxItems": 3 + } + }, + "required": [ + "arm", + "delta_rpy" + ] + }, + "open_gripper": { + "type": "object", + "properties": { + "arm": { + "type": "string", + "enum": [ + "left", + "right" + ], + "description": "Which arm to command; the other arm is left uncommanded." + } + }, + "required": [ + "arm" + ] + }, + "close_gripper": { + "type": "object", + "properties": { + "arm": { + "type": "string", + "enum": [ + "left", + "right" + ], + "description": "Which arm to command; the other arm is left uncommanded." + } + }, + "required": [ + "arm" + ] + }, + "recover_joint_posture": { + "type": "object", + "properties": { + "reason": { + "type": "string", + "default": "" + }, + "return_to_start": { + "type": "boolean", + "default": true + } + } + }, + "request_scene_reset": { + "type": "object", + "properties": { + "reason": { + "type": "string", + "description": "Why the scene needs to be restored." + }, + "expected_scene_state": { + "type": "string", + "default": "", + "description": "Short instruction for the operator describing the desired restored layout." + } + }, + "required": [ + "reason" + ] + }, + "request_operator_verdict": { + "type": "object", + "properties": { + "question": { + "type": "string", + "default": "Does the current real-robot scene satisfy the task?" + } + } + }, + "vla_right_grasp": { + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "Planner-facing segment intent. This is recorded in the tool result; the current live clean-desk checkpoint still receives its fixed training instruction during policy inference." + }, + "max_chunks": { + "type": "integer", + "minimum": 1, + "maximum": 20, + "default": 20 + } + }, + "required": [ + "prompt" + ] + }, + "vla_handoff": { + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "Planner-facing segment intent. This is recorded in the tool result; the current live clean-desk checkpoint still receives its fixed training instruction during policy inference." + }, + "max_chunks": { + "type": "integer", + "minimum": 1, + "maximum": 20, + "default": 20 + } + }, + "required": [ + "prompt" + ] + }, + "vla_left_place": { + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "Planner-facing segment intent. This is recorded in the tool result; the current live clean-desk checkpoint still receives its fixed training instruction during policy inference." + }, + "max_chunks": { + "type": "integer", + "minimum": 1, + "maximum": 20, + "default": 20 + } + }, + "required": [ + "prompt" + ] + }, + "finish": { + "type": "object", + "properties": { + "status": { + "type": "string", + "description": "Outcome, e.g. 'success', 'failure', or 'stuck'." + }, + "summary": { + "type": "string", + "description": "Short natural-language summary of the run." + } + }, + "required": [ + "status", + "summary" + ] + } + }, + "cases": [ + { + "name": "view_env_state", + "arguments": {}, + "fields": [ + "agent_observation", + "artifact_images", + "artifacts", + "available_camera_views", + "image_base_path", + "image_block_order", + "image_d455_path", + "image_left_wrist_path", + "image_right_wrist_path", + "images", + "state", + "step_idx", + "terminated", + "truncated" + ], + "image_count": 1 + }, + { + "name": "view_camera_meta", + "arguments": {}, + "fields": [ + "camera_meta", + "step" + ], + "image_count": 0 + }, + { + "name": "move_delta", + "arguments": { + "arm": "left", + "delta_xyz": [ + 0.01, + 0, + 0 + ] + }, + "fields": [ + "agent_elapsed_s", + "agent_observation", + "artifact_images", + "artifacts", + "available_camera_views", + "command", + "elapsed_s", + "image_base_path", + "image_block_order", + "image_d455_path", + "image_left_wrist_path", + "image_right_wrist_path", + "images", + "result", + "state", + "step_idx", + "terminated", + "truncated" + ], + "image_count": 1, + "result_fields": [ + "arm", + "ok" + ] + }, + { + "name": "rotate_delta", + "arguments": { + "arm": "right", + "delta_rpy": [ + 0, + 0, + 0.1 + ] + }, + "fields": [ + "agent_elapsed_s", + "agent_observation", + "artifact_images", + "artifacts", + "available_camera_views", + "command", + "elapsed_s", + "image_base_path", + "image_block_order", + "image_d455_path", + "image_left_wrist_path", + "image_right_wrist_path", + "images", + "result", + "state", + "step_idx", + "terminated", + "truncated" + ], + "image_count": 1, + "result_fields": [ + "arm", + "ok" + ] + }, + { + "name": "open_gripper", + "arguments": { + "arm": "left" + }, + "fields": [ + "agent_elapsed_s", + "agent_observation", + "artifact_images", + "artifacts", + "available_camera_views", + "command", + "elapsed_s", + "image_base_path", + "image_block_order", + "image_d455_path", + "image_left_wrist_path", + "image_right_wrist_path", + "images", + "result", + "state", + "step_idx", + "terminated", + "truncated" + ], + "image_count": 1, + "result_fields": [ + "arm", + "ok", + "open" + ] + }, + { + "name": "close_gripper", + "arguments": { + "arm": "right" + }, + "fields": [ + "agent_elapsed_s", + "agent_observation", + "artifact_images", + "artifacts", + "available_camera_views", + "command", + "elapsed_s", + "image_base_path", + "image_block_order", + "image_d455_path", + "image_left_wrist_path", + "image_right_wrist_path", + "images", + "result", + "state", + "step_idx", + "terminated", + "truncated" + ], + "image_count": 1, + "result_fields": [ + "arm", + "ok", + "open" + ] + }, + { + "name": "recover_joint_posture", + "arguments": { + "reason": "joint drift" + }, + "fields": [ + "agent_elapsed_s", + "agent_observation", + "artifact_images", + "artifacts", + "available_camera_views", + "command", + "elapsed_s", + "image_base_path", + "image_block_order", + "image_d455_path", + "image_left_wrist_path", + "image_right_wrist_path", + "images", + "result", + "state", + "step_idx", + "terminated", + "truncated" + ], + "image_count": 1, + "result_fields": [ + "ok", + "reason", + "return_to_start" + ] + }, + { + "name": "finish", + "arguments": { + "status": "success", + "summary": "offline baseline" + }, + "fields": [ + "_finish", + "operator_verdict", + "status", + "summary" + ], + "image_count": 0 + } + ], + "descriptions": { + "describe_dual_franka_setup": "Read the dual-Franka runtime conventions, camera aliases, VLA policy conditioning text, semantic stop rules, and available primitive names before acting. This is read-only.", + "view_env_state": "Read a dual-Franka state snapshot. The D455 image is returned inline; left_wrist, base, and right_wrist are returned as artifact paths for targeted read_image inspection.", + "view_camera_meta": "Read camera intrinsics, serials, and projection metadata for the dual-Franka rig.", + "back_project": "Back-project one pixel from a registered RGBD camera view into shared right-base coordinates. Use a camera listed by view_env_state/view_camera_meta; default is the configured primary metric localization camera.", + "segment": "Use SAM3 on a registered RGB image with either a text prompt or one positive [row, col] point, return a mask overlay for verification, and estimate the mask median point in shared right-base coordinates.", + "move_delta": "Move one Franka TCP by a bounded world-frame xyz delta in meters.", + "rotate_delta": "Rotate one Franka TCP by a bounded world-frame rpy delta in radians.", + "open_gripper": "Open one Franka gripper and wait for the command to settle.", + "close_gripper": "Close one Franka gripper and wait for the command to settle.", + "recover_joint_posture": "Reset both arms to their healthy configured joint posture while preserving each gripper's open/closed state. Closed grippers are re-commanded before/after the joint reset so held objects stay clamped, then both TCPs return near their prior poses.", + "request_scene_reset": "Exploration-only real-robot reset gate. Ask the human operator to remove/secure held objects and restore the tabletop scene for another attempt, wait for terminal confirmation, then reset the robot posture. This does not automatically restore physical objects like a simulator.", + "request_operator_verdict": "Exploration-only human feedback gate. Ask the operator to mark the current physical task state as success, failure, or continue before the planner finishes or starts another attempt.", + "vla_right_grasp": "Run the learned right-grasp VLA segment. The active task prompt decides which object is currently allowed; this tool only defines the capability boundary: right gripper closes and the right TCP lifts.", + "vla_handoff": "Run the learned bimanual handoff VLA segment. The capability boundary is right-gripper release followed by the configured settle delay; do not rule-base pre-position either arm for it.", + "vla_left_place": "Run the learned left-placement VLA segment. The active task decides the destination; this tool only defines the capability boundary: left gripper opens and the left TCP lifts.", + "finish": "Call when the task is complete or unrecoverable. Halts the agent loop. Save any artifacts (recipe, audit) BEFORE calling finish." + } +} diff --git a/tests/unit_tests/robots/dual_franka/test_dual_franka_tools.py b/tests/unit_tests/robots/dual_franka/test_dual_franka_tools.py index 89ddbf414..5b2b3144e 100644 --- a/tests/unit_tests/robots/dual_franka/test_dual_franka_tools.py +++ b/tests/unit_tests/robots/dual_franka/test_dual_franka_tools.py @@ -16,19 +16,18 @@ from __future__ import annotations +import json from pathlib import Path +from threading import Event import numpy as np import pytest -from robots.dual_franka import get_robot_spec +from robots.dual_franka import get_robot_spec, tools from robots.dual_franka.perception import back_project, segment from robots.dual_franka.tasks import CLEAN_DESK_VLA_PROMPT -from robots.dual_franka.toolkit import DualFrankaToolkit +from robots.dual_franka.toolkit import DualFrankaRuntime, DualFrankaToolkit from robots.dual_franka.tools import ( - DualFrankaPrimitives, - coerce_arm, - coerce_vec3, dump_state, view_env_state, ) @@ -38,7 +37,7 @@ from rpent.dashboard.state import DashboardState from rpent.memory import MemoryManager from rpent.session import EnvState -from rpent.tools.toolkit import ToolResult +from rpent.tools import ToolContext class FakeEnv: @@ -179,13 +178,23 @@ def chunk_step(self, actions): return {"terminated": False, "truncated": False} -def _primitives(env: FakeEnv, *, model=None, check_cancelled=lambda: None): - return DualFrankaPrimitives( +def _runtime(env: FakeEnv, *, model=None): + return DualFrankaRuntime( env=env, model=model, task_description="default task", vla_instruction=CLEAN_DESK_VLA_PROMPT, - check_cancelled=check_cancelled, + ) + + +def _context(runtime=None, *, state=None): + return ToolContext( + robot=runtime, + state=state, + memory=None, + output_dir=Path("."), + record_frame=lambda frame: None, + _cancel_event=Event(), ) @@ -224,14 +233,14 @@ def test_toolkit_exploration_tools_are_opt_in(tmp_path: Path): refused_eval = evaluation.execute_tool( "finish", {"status": "success", "summary": "not confirmed"} ) - assert refused_eval.result["error"] == "finish refused" - assert refused_eval.is_finish is False + assert refused_eval.to_dict()["error"] == "finish refused" + assert refused_eval.is_error evaluation._read_operator_line = lambda prompt: "success checked by operator" evaluation.execute_tool("request_operator_verdict", {}) accepted_eval = evaluation.execute_tool( "finish", {"status": "success", "summary": "confirmed"} ) - assert accepted_eval.is_finish is True + assert not accepted_eval.is_error assert evaluation.solved() assert "request_scene_reset" in _tool_names(exploration) assert "request_operator_verdict" in _tool_names(exploration) @@ -240,8 +249,8 @@ def test_toolkit_exploration_tools_are_opt_in(tmp_path: Path): "finish", {"status": "success", "summary": "agent thinks done"}, ) - assert refused.result["error"].startswith("finish refused") - assert refused.is_finish is False + assert refused.to_dict()["error"].startswith("finish refused") + assert refused.is_error exploration._read_operator_line = lambda prompt: "done" exploration.execute_tool("request_scene_reset", {"reason": "prepare"}) @@ -251,8 +260,8 @@ def test_toolkit_exploration_tools_are_opt_in(tmp_path: Path): "finish", {"status": "success", "summary": "operator accepted"}, ) - assert accepted.is_finish is True - assert accepted.result["operator_verdict"] == "success" + assert not accepted.is_error + assert accepted.to_dict()["operator_verdict"] == "success" def test_scene_reset_waits_for_operator_then_resets_robot(tmp_path: Path): @@ -281,7 +290,7 @@ def test_scene_reset_waits_for_operator_then_resets_robot(tmp_path: Path): "reason": "retry with restored layout", "expected_scene_state": "objects back at the starting positions", }, - ).result + ).to_dict() assert env.resets == 1 assert result["result"]["robot_reset"] == {"ok": True} @@ -292,32 +301,37 @@ def test_scene_reset_waits_for_operator_then_resets_robot(tmp_path: Path): def test_arm_and_vec3_validation_and_motion_forwarding(): env = FakeEnv() - primitives = _primitives(env) + runtime = _runtime(env) - primitives.move_delta("left", [0.01, 0.0, -0.02]) - primitives.rotate_delta("right", [0.0, 0.0, 0.1]) - primitives.open_gripper("left") - primitives.close_gripper("right") + tools.move_delta.handler("left", [0.01, 0.0, -0.02], ctx=_context(runtime)) + tools.rotate_delta.handler("right", [0.0, 0.0, 0.1], ctx=_context(runtime)) + tools.open_gripper.handler("left", ctx=_context(runtime)) + tools.close_gripper.handler("right", ctx=_context(runtime)) assert env.moves[0][0] == "left" np.testing.assert_allclose(env.moves[0][1], [0.01, 0.0, -0.02]) assert env.rotations[0][0] == "right" assert env.grippers == [("left", True), ("right", False)] - assert coerce_arm("LEFT") == "left" + assert tools.open_gripper.args_schema.model_validate({"arm": "LEFT"}).arm == "left" with pytest.raises(ValueError, match="left.*right"): - coerce_arm("both") - with pytest.raises(ValueError, match="exactly 3"): - coerce_vec3([1.0, 2.0], name="delta") + tools.open_gripper.args_schema.model_validate({"arm": "both"}) + with pytest.raises(ValueError, match="at least 3"): + tools.move_delta.args_schema.model_validate( + { + "arm": "left", + "delta_xyz": [1.0, 2.0], + } + ) def test_dump_state_saves_three_camera_artifacts(tmp_path: Path): env = FakeEnv() - primitives = _primitives(env) + runtime = _runtime(env) state = EnvState(tmp_path) record = dump_state( - primitives, + runtime, state, command={"action": "move_delta"}, result={"ok": True}, @@ -335,8 +349,9 @@ def test_dump_state_saves_three_camera_artifacts(tmp_path: Path): "d455_depth.npy", "camera_meta.json", } - output = view_env_state(state=state) - assert output["_image_bytes"] + output = view_env_state.handler(ctx=_context(state=state)) + assert output.images + output = output.to_dict() assert "_image_nav_bytes" not in output assert "_image_cam_bytes" not in output assert "_image_wrist_bytes" not in output @@ -344,14 +359,16 @@ def test_dump_state_saves_three_camera_artifacts(tmp_path: Path): assert output["image_block_order"] == ["d455"] np.testing.assert_array_equal(state.load("base.png"), 7) np.testing.assert_array_equal(state.load("base_depth.npy"), 9) - camera_meta = view_camera_meta(state=state)["camera_meta"] + camera_meta = view_camera_meta.handler(ctx=_context(state=state)).data[ + "camera_meta" + ] assert camera_meta["observation_camera_map"]["main"] == "left_wrist_0_rgb" def test_dashboard_discovers_dual_franka_camera_artifacts(tmp_path: Path): state = EnvState(tmp_path / "state") record = dump_state( - _primitives(FakeEnv()), state, command=None, result=None, elapsed_s=None + _runtime(FakeEnv()), state, command=None, result=None, elapsed_s=None ) dashboard = DashboardState( output_dir=tmp_path, dashboard_spec=get_robot_spec().dashboard @@ -365,23 +382,24 @@ def test_dashboard_discovers_dual_franka_camera_artifacts(tmp_path: Path): def test_view_env_state_emits_multimodal_image_blocks(tmp_path: Path): env = FakeEnv() - primitives = _primitives(env) + runtime = _runtime(env) state = EnvState(tmp_path) - dump_state(primitives, state, command=None, result=None, elapsed_s=None) - output = view_env_state(state=state) + dump_state(runtime, state, command=None, result=None, elapsed_s=None) + output = view_env_state.handler(ctx=_context(state=state)) # Routine planner snapshots inline only D455, while auxiliary camera # artifacts stay available through returned paths/read_image. + output = output.to_dict() assert output["images"] == ["d455"] assert output["image_block_order"] == output["images"] assert output["artifact_images"] == ["base", "d455", "left_wrist", "right_wrist"] - result = ToolResult(name="view_env_state", result=output) - image_blocks = [b for b in result.content_blocks if b.get("type") == "image"] + result = view_env_state.handler(ctx=_context(state=state)) + image_blocks = result.images assert len(image_blocks) == 1 - text_block = next(b for b in result.content_blocks if b.get("type") == "text") + text = result.to_text() # Image bytes must be lifted out of the text block, not serialized into it. - assert "_image_" not in text_block["text"] + assert "_image_" not in text def test_back_project_reads_rpent_state_artifacts(tmp_path: Path): @@ -471,11 +489,11 @@ def test_back_project_returns_annotated_image_block(tmp_path: Path): assert result["_image_cam_bytes"] assert result["image_block_order"] == ["d455_selection_diagnostic"] - tool_result = ToolResult(name="back_project", result=result) - image_blocks = [b for b in tool_result.content_blocks if b.get("type") == "image"] + tool_result = tools._perception_result(result) + image_blocks = tool_result.images assert len(image_blocks) == 1 - text_block = next(b for b in tool_result.content_blocks if b.get("type") == "text") - assert "_image_" not in text_block["text"] + text = tool_result.to_text() + assert "_image_" not in text def test_segment_returns_mask_overlay_and_world_point(tmp_path: Path): @@ -529,8 +547,8 @@ def test_segment_returns_mask_overlay_and_world_point(tmp_path: Path): assert result["_image_cam_bytes"] assert result["image_block_order"] == ["d455_segment_overlay"] - tool_result = ToolResult(name="segment", result=result) - image_blocks = [b for b in tool_result.content_blocks if b.get("type") == "image"] + tool_result = tools._perception_result(result) + image_blocks = tool_result.images assert len(image_blocks) == 1 @@ -548,19 +566,23 @@ def test_segment_without_sam3_client_falls_back(tmp_path: Path): def test_vla_grasp_runs_bounded_chunks(): env = FakeEnv() - primitives = _primitives(env, model=FakeModel()) + runtime = _runtime(env, model=FakeModel(expected_prompt=CLEAN_DESK_VLA_PROMPT)) - result = primitives.vla_grasp("hand over the cube", max_chunks=3) + result = tools.vla_right_grasp.handler( + "hand over the cube", max_chunks=3, ctx=_context(runtime) + ).to_dict() assert result["chunks_executed"] == 3 - assert len(env.chunks) == 3 + assert len(env.chunks) == 6 def test_recover_joint_posture_forwards_to_env(): env = FakeEnv() - primitives = _primitives(env) + runtime = _runtime(env) - result = primitives.recover_joint_posture(reason="joint drift") + result = tools.recover_joint_posture.handler( + reason="joint drift", ctx=_context(runtime) + ).to_dict() assert result["ok"] assert result["reason"] == "joint drift" @@ -568,14 +590,14 @@ def test_recover_joint_posture_forwards_to_env(): def test_named_clean_desk_vla_uses_fixed_prompt_and_semantic_boundary(): env = BoundaryEnv() - primitives = _primitives( + runtime = _runtime( env, model=FakeModel(expected_prompt=CLEAN_DESK_VLA_PROMPT), ) - result = primitives.vla_right_grasp( - prompt="grasp the next task-allowed object", max_chunks=2 - ) + result = tools.vla_right_grasp.handler( + prompt="grasp the next task-allowed object", max_chunks=2, ctx=_context(runtime) + ).to_dict() assert result["ok"] assert result["skill_name"] == "vla_right_grasp" @@ -587,13 +609,54 @@ def test_named_clean_desk_vla_uses_fixed_prompt_and_semantic_boundary(): def test_named_vla_uses_task_configured_policy_instruction(): instruction = "custom checkpoint instruction" - primitives = DualFrankaPrimitives( + runtime = DualFrankaRuntime( env=BoundaryEnv(), model=FakeModel(expected_prompt=instruction), task_description="planner task description", vla_instruction=instruction, - check_cancelled=lambda: None, ) - result = primitives.vla_right_grasp(prompt="planner segment intent", max_chunks=2) + result = tools.vla_right_grasp.handler( + prompt="planner segment intent", max_chunks=2, ctx=_context(runtime) + ).to_dict() assert result["effective_policy_prompt"] == instruction assert result["prompt_overridden"] + + +# Target-branch schemas and returns captured at bd9040a, before native migration. +_CONTRACT = json.loads( + (Path(__file__).parent / "fixtures/pre_native_tool_contracts.json").read_text() +) + + +@pytest.mark.parametrize("case", _CONTRACT["cases"], ids=lambda case: case["name"]) +def test_normal_return_contract(case, tmp_path): + toolkit = DualFrankaToolkit( + runtime_kwargs={ + "env": FakeEnv(), + "model": None, + "task_description": "default task", + }, + dashboard_events=NullDashboardEventSink(), + memory=MemoryManager(tmp_path / "memory"), + state_output_dir=tmp_path, + ) + if case["name"] == "finish": + toolkit._read_operator_line = lambda _: "success" + assert not toolkit.execute_tool("request_operator_verdict", {}).is_error + result = toolkit.execute_tool(case["name"], case["arguments"]) + assert not result.is_error + data = result.to_dict() + assert sorted(data) == case["fields"] + assert len(result.images) == case["image_count"] + if "result_fields" in case: + assert sorted(data["result"]) == case["result_fields"] + if case["name"] == "finish": + assert data == { + "_finish": True, + "operator_verdict": "success", + **case["arguments"], + } + assert toolkit.finish_result == { + "operator_verdict": "success", + **case["arguments"], + } diff --git a/tests/unit_tests/robots/dual_franka/test_exploration.py b/tests/unit_tests/robots/dual_franka/test_exploration.py index d6e2a8978..c5e9658db 100644 --- a/tests/unit_tests/robots/dual_franka/test_exploration.py +++ b/tests/unit_tests/robots/dual_franka/test_exploration.py @@ -16,12 +16,13 @@ import argparse import json +from dataclasses import replace from pathlib import Path import numpy as np import pytest -from robots.dual_franka import robot_spec +from robots.dual_franka import robot_spec, tools from robots.dual_franka.toolkit import DualFrankaToolkit from rpent.dashboard.events import NullDashboardEventSink from rpent.memory import MemoryManager @@ -90,7 +91,7 @@ def setup(tmp_path, monkeypatch): def call(t, name, **kwargs): - return t.execute_tool(name, kwargs).result + return t.execute_tool(name, kwargs).to_dict() def reset(t, replies): @@ -174,6 +175,8 @@ def test_operator_abort_allows_finish_even_with_budget_and_never_succeeds(setup) call(t, "request_scene_reset", reason="initial") result = call(t, "finish", status="success", summary="stop") assert result["_finish"] and result["operator_aborted"] + assert t.finish_result["operator_aborted"] is True + assert "_finish" not in t.finish_result assert result["status"] == "failure" and not t.solved() and env.resets == 0 @@ -275,7 +278,10 @@ def test_successful_memory_pair_uses_existing_merge_and_index(setup, tmp_path): assert (t.memory.root / "task_only/dual_franka_t0_recipe.jsonl").exists() -def test_cli_two_sessions_operator_feedback_and_memory_pipeline(tmp_path, monkeypatch): +@pytest.mark.parametrize("first_verdict", ["failure dropped", "abort"]) +def test_cli_operator_feedback_controls_sessions_and_memory( + tmp_path, monkeypatch, first_verdict +): import sys from dataclasses import replace from types import SimpleNamespace @@ -285,7 +291,7 @@ def test_cli_two_sessions_operator_feedback_and_memory_pipeline(tmp_path, monkey env = FakeEnv() runtimes = [] planners = [] - replies = iter(["done", "failure dropped", "done", "success lifted"]) + replies = iter(["done", first_verdict, "done", "success lifted"]) class Operator: def __init__(self, **kwargs): @@ -341,7 +347,7 @@ def solve(self, *, toolkit, system_prompt, user_message, **kwargs): ) assert finish["_finish"] return SimpleNamespace( - finish_result=finish, messages=[], stats={}, error=None + finish_result=toolkit.finish_result, messages=[], stats={}, error=None ) def init_runtime(*args): @@ -378,7 +384,19 @@ def init_runtime(*args): ], ) assert cli.main() == 0 - assert len(planners) == 2 and len(runtimes) == 1 and env.resets == 2 + session_count = 1 if first_verdict == "abort" else 2 + assert len(planners) == session_count + assert len(runtimes) == 1 and env.resets == session_count + if first_verdict == "abort": + transcript = json.loads( + (tmp_path / "run/transcript_dual_franka_t0.json").read_text() + ) + assert transcript["finish"]["operator_aborted"] is True + assert transcript["finish"]["status"] == "failure" + assert "_finish" not in transcript["finish"] + assert not (tmp_path / "run/sessions/session_002").exists() + assert not (tmp_path / "memory/global/cli.md").exists() + return traces = [ json.loads((tmp_path / f"run/sessions/session_{i:03d}/states.json").read_text()) for i in (1, 2) @@ -460,11 +478,15 @@ def call(self, name, **kwargs): ("recover_joint_posture", {"reason": "joint warning"}), ], ) -def test_pr176_added_motion_tools_share_explore_guards(setup, tool, arguments): +def test_pr176_added_motion_tools_share_explore_guards( + setup, tool, arguments, monkeypatch +): t, env, replies = setup executed = [] - t._primitives._run_named_vla_skill = lambda **kwargs: ( - executed.append(kwargs) or {"ok": True} + monkeypatch.setattr( + tools, + "_run_named_vla_skill", + lambda ctx, **kwargs: executed.append(kwargs) or {"ok": True}, ) env.recover_joint_posture = lambda **kwargs: executed.append(kwargs) or {"ok": True} assert tool in t._tools @@ -489,7 +511,7 @@ def test_setup_describes_operator_reset_in_exploration(setup): def test_direct_success_stops_active_tool_and_records_memory(setup, tmp_path): import threading - from rpent.tools.toolkit import ToolCancelled + from rpent.tools import ToolCancelled t, env, replies = setup assert not t.request_direct_verdict("success") @@ -504,8 +526,12 @@ def active_motion(**kwargs): t.raise_if_cancelled() pytest.fail("motion must not continue after success") - t.add_tool("move_delta", t._tools["move_delta"][0], active_motion) - worker = threading.Thread(target=lambda: results.append(call(t, "move_delta"))) + t._tools["move_delta"] = replace(t._tools["move_delta"], handler=active_motion) + worker = threading.Thread( + target=lambda: results.append( + call(t, "move_delta", arm="right", delta_xyz=[0.01, 0, 0]) + ) + ) worker.start() assert entered.wait(2) assert t.request_direct_verdict("success") @@ -514,7 +540,7 @@ def active_motion(**kwargs): release.set() worker.join(2) assert not worker.is_alive() - assert results[0]["code"] == "tool_cancelled" + assert "operator submitted a terminal verdict" in results[0]["error"] with pytest.raises(ToolCancelled): t.raise_if_cancelled() result = t.finalize_direct_verdict() diff --git a/tests/unit_tests/robots/fixtures/common_tool_contracts.json b/tests/unit_tests/robots/fixtures/common_tool_contracts.json new file mode 100644 index 000000000..a4bfc685e --- /dev/null +++ b/tests/unit_tests/robots/fixtures/common_tool_contracts.json @@ -0,0 +1,60 @@ +{ + "source_commit": "949f61ecea1c571168ca6a29a59f6d3207f6f947", + "schemas": { + "read_text_file": { + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Absolute or repo-relative path" + }, + "max_chars": { + "type": "integer", + "description": "Max chars (default 40000)" + } + }, + "required": [ + "path" + ] + }, + "write_text_file": { + "type": "object", + "properties": { + "path": { + "type": "string" + }, + "content": { + "type": "string" + } + }, + "required": [ + "path", + "content" + ] + }, + "list_dir": { + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Default: {{output_dir}}" + } + } + }, + "read_image": { + "properties": { + "name": { + "type": "string" + }, + "step": { + "type": "integer", + "minimum": -1 + } + }, + "required": [ + "name" + ], + "type": "object" + } + } +} diff --git a/tests/unit_tests/robots/franka/_fakes.py b/tests/unit_tests/robots/franka/_fakes.py new file mode 100644 index 000000000..7eea67df3 --- /dev/null +++ b/tests/unit_tests/robots/franka/_fakes.py @@ -0,0 +1,74 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""CPU clients shared by behavior and pre-native compatibility tests.""" + +from __future__ import annotations + +import numpy as np + + +class FakeEnv: + def __init__(self) -> None: + self.moves: list[np.ndarray] = [] + self.rotations: list[np.ndarray] = [] + self.gripper_open = True + self.chunks: list[np.ndarray] = [] + self.observation_calls = 0 + + def reset(self): + return {"ok": True} + + def move_delta(self, value): + self.moves.append(np.asarray(value)) + return {"ok": True} + + def rotate_delta(self, value): + self.rotations.append(np.asarray(value)) + return {"ok": True} + + def set_gripper(self, *, open: bool): + self.gripper_open = open + return {"ok": True, "open": open} + + def _obs(self): + return { + "main_images": np.zeros((8, 8, 3), dtype=np.uint8), + "extra_view_images": np.ones((1, 8, 8, 3), dtype=np.uint8), + "main_depths": np.ones((8, 8), dtype=np.float32), + "extra_view_depths": np.ones((1, 8, 8), dtype=np.float32) * 2, + "states": np.zeros(8, dtype=np.float32), + } + + def get_observation(self): + self.observation_calls += 1 + return self._obs() + + def get_robot_state(self): + return {"tcp_pose": [0.5, 0.0, 0.2, 0.0, 0.0, 0.0, 1.0]} + + def get_camera_meta(self): + return {"depth_unit": "m", "cameras": {"wrist_1": {"fx": 100.0}}} + + def chunk_step(self, actions): + self.chunks.append(np.asarray(actions)) + return {"terminated": False, "truncated": False, "observation": self._obs()} + + +class FakeModel: + def predict(self, observation, options=None): + assert observation["task_descriptions"] == "pick up the cube" + assert options == {"mode": "eval"} + return np.zeros((2, 7), dtype=np.float32) diff --git a/tests/unit_tests/robots/franka/fixtures/pre_native_tool_contracts.json b/tests/unit_tests/robots/franka/fixtures/pre_native_tool_contracts.json new file mode 100644 index 000000000..3d81b25f7 --- /dev/null +++ b/tests/unit_tests/robots/franka/fixtures/pre_native_tool_contracts.json @@ -0,0 +1,408 @@ +{ + "source_commit": "8cad76c65ebf40f716e65eb5e8b9c50697d247d0", + "schemas": { + "finish": { + "type": "object", + "properties": { + "status": { + "type": "string", + "description": "Outcome, e.g. 'success', 'failure', or 'stuck'." + }, + "summary": { + "type": "string", + "description": "Short natural-language summary of the run." + } + }, + "required": [ + "status", + "summary" + ] + }, + "view_env_state": { + "type": "object", + "properties": { + "step": { + "type": "integer", + "default": -1 + } + } + }, + "view_camera_meta": { + "type": "object", + "properties": { + "step": { + "type": "integer", + "default": -1 + } + } + }, + "view_perception_setup": { + "type": "object", + "properties": { + "step": { + "type": "integer", + "default": -1 + } + } + }, + "back_project": { + "type": "object", + "properties": { + "row": { + "type": "integer", + "minimum": 0 + }, + "col": { + "type": "integer", + "minimum": 0 + }, + "step": { + "type": "integer" + }, + "camera": { + "type": "string", + "enum": [ + "wrist", + "third_person" + ] + }, + "debug": { + "type": "boolean", + "default": false + } + }, + "required": [ + "row", + "col" + ] + }, + "back_project_correspondence": { + "type": "object", + "properties": { + "third_person_row": { + "type": "integer", + "minimum": 0 + }, + "third_person_col": { + "type": "integer", + "minimum": 0 + }, + "wrist_row": { + "type": "integer", + "minimum": 0 + }, + "wrist_col": { + "type": "integer", + "minimum": 0 + }, + "pixels": { + "type": "array", + "items": { + "type": "object" + } + }, + "step": { + "type": "integer" + }, + "debug": { + "type": "boolean", + "default": false + } + } + }, + "move_delta": { + "type": "object", + "properties": { + "delta_xyz": { + "type": "array", + "items": { + "type": "number" + }, + "minItems": 3, + "maxItems": 3 + } + }, + "required": [ + "delta_xyz" + ] + }, + "rotate_delta": { + "type": "object", + "properties": { + "delta_rpy": { + "type": "array", + "items": { + "type": "number" + }, + "minItems": 3, + "maxItems": 3 + } + }, + "required": [ + "delta_rpy" + ] + }, + "open_gripper": { + "type": "object", + "properties": {} + }, + "close_gripper": { + "type": "object", + "properties": {} + }, + "vla_grasp": { + "type": "object", + "properties": { + "prompt": { + "type": "string" + }, + "max_chunks": { + "type": "integer", + "minimum": 1, + "maximum": 20 + } + }, + "required": [ + "prompt" + ] + } + }, + "cases": [ + { + "name": "view_env_state", + "arguments": {}, + "fields": [ + "artifacts", + "image_cam_path", + "image_wrist_path", + "state", + "step_idx", + "terminated", + "truncated" + ], + "image_count": 2 + }, + { + "name": "view_camera_meta", + "arguments": {}, + "fields": [ + "camera_meta", + "step" + ], + "image_count": 0 + }, + { + "name": "move_delta", + "arguments": { + "delta_xyz": [ + 0.01, + 0, + 0 + ] + }, + "fields": [ + "agent_elapsed_s", + "artifacts", + "command", + "elapsed_s", + "image_cam_path", + "image_wrist_path", + "result", + "state", + "step_idx", + "terminated", + "truncated" + ], + "result_fields": [ + "ok" + ], + "image_count": 2 + }, + { + "name": "rotate_delta", + "arguments": { + "delta_rpy": [ + 0, + 0, + 0.1 + ] + }, + "fields": [ + "agent_elapsed_s", + "artifacts", + "command", + "elapsed_s", + "image_cam_path", + "image_wrist_path", + "result", + "state", + "step_idx", + "terminated", + "truncated" + ], + "result_fields": [ + "ok" + ], + "image_count": 2 + }, + { + "name": "open_gripper", + "arguments": {}, + "fields": [ + "agent_elapsed_s", + "artifacts", + "command", + "elapsed_s", + "image_cam_path", + "image_wrist_path", + "result", + "state", + "step_idx", + "terminated", + "truncated" + ], + "result_fields": [ + "ok", + "open" + ], + "image_count": 2 + }, + { + "name": "close_gripper", + "arguments": {}, + "fields": [ + "agent_elapsed_s", + "artifacts", + "command", + "elapsed_s", + "image_cam_path", + "image_wrist_path", + "result", + "state", + "step_idx", + "terminated", + "truncated" + ], + "result_fields": [ + "ok", + "open" + ], + "image_count": 2 + }, + { + "name": "vla_grasp", + "arguments": { + "prompt": "pick up the cube", + "max_chunks": 1 + }, + "fields": [ + "agent_elapsed_s", + "artifacts", + "command", + "elapsed_s", + "image_cam_path", + "image_wrist_path", + "result", + "state", + "step_idx", + "terminated", + "truncated" + ], + "result_fields": [ + "chunks_executed", + "last_chunk", + "ok", + "robot_state" + ], + "last_chunk_fields": [ + "observation", + "terminated", + "truncated" + ], + "image_count": 2 + }, + { + "name": "finish", + "arguments": { + "status": "success", + "summary": "offline baseline" + }, + "fields": [ + "_finish", + "status", + "summary" + ], + "image_count": 0 + }, + { + "name": "view_perception_setup", + "arguments": {}, + "fields": [ + "calibration", + "camera_meta", + "convention", + "current_policy" + ], + "image_count": 0 + }, + { + "name": "back_project", + "arguments": { + "row": 2, + "col": 2, + "debug": true + }, + "fields": [ + "camera", + "camera_key", + "camera_name", + "coordinate_frame", + "debug", + "depth_m", + "pixel", + "point_base", + "selected_pixel_overlay", + "source", + "source_artifact", + "step", + "world_xyz" + ], + "image_count": 0 + }, + { + "name": "back_project_correspondence", + "arguments": { + "wrist_row": 2, + "wrist_col": 2, + "third_person_row": 2, + "third_person_col": 2 + }, + "fields": [ + "base_point_delta_m", + "confidence", + "pixel_correspondence", + "point_base", + "point_base_third_person", + "point_base_wrist", + "source", + "step", + "tcp_pose_source", + "warnings" + ], + "image_count": 0 + } + ], + "descriptions": { + "finish": "Call when the task is complete or unrecoverable. Halts the agent loop. Save any artifacts (recipe, audit) BEFORE calling finish.", + "view_env_state": "Read a Franka state snapshot and its synchronized RGB images.", + "view_camera_meta": "Read camera intrinsics, crop, depth, and calibration metadata.", + "view_perception_setup": "Read calibrated camera geometry and projection conventions.", + "back_project": "Back-project one wrist or external-camera pixel into Franka base coordinates.", + "back_project_correspondence": "Fuse matched wrist and external-camera pixels into a Franka base point.", + "move_delta": "Move the Franka TCP by a bounded base-frame xyz delta in meters.", + "rotate_delta": "Rotate the Franka TCP by a bounded base-frame rpy delta in radians.", + "open_gripper": "Open the Franka gripper and wait for the command to settle.", + "close_gripper": "Close the Franka gripper and wait for the command to settle.", + "vla_grasp": "Run bounded real-world VLA action chunks for a local grasp attempt." + } +} diff --git a/tests/unit_tests/robots/franka/test_tools.py b/tests/unit_tests/robots/franka/test_tools.py index f5283fa8c..59a2ec426 100644 --- a/tests/unit_tests/robots/franka/test_tools.py +++ b/tests/unit_tests/robots/franka/test_tools.py @@ -16,106 +16,60 @@ from __future__ import annotations +import json from pathlib import Path +from threading import Event +from types import SimpleNamespace import numpy as np import pytest +from robots.franka import runtime_config, tools +from robots.franka import toolkit as franka_toolkit from robots.franka.perception import back_project from robots.franka.runtime_config import set_calibration_path +from robots.franka.toolkit import FrankaRuntime, FrankaToolkit from robots.franka.tools import ( - FrankaPrimitives, - coerce_vec3, dump_state, view_camera_meta, view_env_state, ) +from rpent.dashboard.events import NullDashboardEventSink, StepRecordEvent +from rpent.memory import MemoryManager from rpent.session import EnvState - - -class FakeEnv: - def __init__(self) -> None: - self.moves: list[np.ndarray] = [] - self.rotations: list[np.ndarray] = [] - self.gripper_open = True - self.chunks: list[np.ndarray] = [] - self.observation_calls = 0 - - def reset(self): - return {"ok": True} - - def move_delta(self, value): - self.moves.append(np.asarray(value)) - return {"ok": True} - - def rotate_delta(self, value): - self.rotations.append(np.asarray(value)) - return {"ok": True} - - def set_gripper(self, *, open: bool): - self.gripper_open = open - return {"ok": True, "open": open} - - def _obs(self): - return { - "main_images": np.zeros((8, 8, 3), dtype=np.uint8), - "extra_view_images": np.ones((1, 8, 8, 3), dtype=np.uint8), - "main_depths": np.ones((8, 8), dtype=np.float32), - "extra_view_depths": np.ones((1, 8, 8), dtype=np.float32) * 2, - "states": np.zeros(8, dtype=np.float32), - } - - def get_observation(self): - self.observation_calls += 1 - return self._obs() - - def get_robot_state(self): - return {"tcp_pose": [0.5, 0.0, 0.2, 0.0, 0.0, 0.0, 1.0]} - - def get_camera_meta(self): - return {"depth_unit": "m", "cameras": {"wrist_1": {"fx": 100.0}}} - - def chunk_step(self, actions): - self.chunks.append(np.asarray(actions)) - return {"terminated": False, "truncated": False, "observation": self._obs()} - - -class FakeModel: - def predict(self, observation, options=None): - assert observation["task_descriptions"] == "pick up the cube" - assert options == {"mode": "eval"} - return np.zeros((2, 7), dtype=np.float32) - - -def _primitives(env: FakeEnv, *, model=None, check_cancelled=lambda: None): - return FrankaPrimitives( - env=env, - model=model, - task_description="default task", - check_cancelled=check_cancelled, +from rpent.tools import ToolContext +from tests.unit_tests.robots.franka._fakes import FakeEnv, FakeModel + + +def _context(env: FakeEnv, *, model=None, state=None): + return ToolContext( + robot=FrankaRuntime(env=env, model=model, task_description="default task"), + state=state, + memory=None, + output_dir=Path("."), + record_frame=lambda frame: None, + _cancel_event=Event(), ) def test_vec3_validation_and_motion_forwarding(): env = FakeEnv() - primitives = _primitives(env) + ctx = _context(env) - primitives.move_delta([0.01, 0.0, -0.02]) - primitives.rotate_delta([0.0, 0.0, 0.1]) + tools.move_delta.handler([0.01, 0.0, -0.02], ctx=ctx) + tools.rotate_delta.handler([0.0, 0.0, 0.1], ctx=ctx) np.testing.assert_allclose(env.moves[0], [0.01, 0.0, -0.02]) np.testing.assert_allclose(env.rotations[0], [0.0, 0.0, 0.1]) - with pytest.raises(ValueError, match="exactly 3"): - coerce_vec3([1.0, 2.0], name="delta") def test_dump_state_saves_canonical_rgbd_artifacts(tmp_path: Path): env = FakeEnv() - primitives = _primitives(env) + ctx = _context(env) state = EnvState(tmp_path) record = dump_state( - primitives, + ctx.robot, state, command={"action": "move_delta"}, result={"ok": True}, @@ -129,19 +83,29 @@ def test_dump_state_saves_canonical_rgbd_artifacts(tmp_path: Path): "wrist.png", "wrist_depth.npy", } - output = view_env_state(state=state) - assert output["_image_wrist_bytes"] - assert output["_image_cam_bytes"] - assert view_camera_meta(state=state)["camera_meta"]["depth_unit"] == "m" + output = view_env_state.handler(ctx=_context(env, state=state)) + assert output.images == [ + state.load_bytes("camera.png"), + state.load_bytes("wrist.png"), + ] + assert "images" not in output.data + assert output.data["image_cam_path"] == str(state.artifact_path("camera.png")) + assert output.data["image_wrist_path"] == str(state.artifact_path("wrist.png")) + assert ( + view_camera_meta.handler(ctx=_context(env, state=state)).data["camera_meta"][ + "depth_unit" + ] + == "m" + ) def test_vla_grasp_runs_bounded_chunks(): env = FakeEnv() - primitives = _primitives(env, model=FakeModel()) + ctx = _context(env, model=FakeModel()) - result = primitives.vla_grasp("pick up the cube", max_chunks=3) + result = tools.vla_grasp.handler("pick up the cube", max_chunks=3, ctx=ctx) - assert result["chunks_executed"] == 3 + assert result.data["chunks_executed"] == 3 assert len(env.chunks) == 3 # Obs is fetched once, then threaded from each chunk_step result. assert env.observation_calls == 1 @@ -172,3 +136,106 @@ def test_back_project_reads_rpent_state_artifacts(tmp_path: Path): assert result["coordinate_frame"] == "franka_base" assert result["depth_m"] == 0.5 assert len(result["point_base"]) == 3 + + +def test_toolkit_factory_validation_and_capture(tmp_path, monkeypatch): + monkeypatch.setattr(franka_toolkit, "get_output_dir", lambda: tmp_path) + env = FakeEnv() + # RPC state contains NumPy arrays; native result text must still be JSON. + get_robot_state = env.get_robot_state + env.get_robot_state = lambda: {**get_robot_state(), "states": np.zeros(8)} + events = [] + kwargs = {"env": env, "model": None, "task_description": "test task"} + toolkit = FrankaToolkit( + runtime_kwargs=kwargs, + dashboard_events=SimpleNamespace(enabled=True, emit=events.append), + memory=MemoryManager(root=tmp_path / "memory"), + ) + assert kwargs == {"env": env, "model": None, "task_description": "test task"} + assert env.observation_calls == 1 + assert len(events) == 1 and isinstance(events[0], StepRecordEvent) + assert all( + "ctx" not in tool.input_schema["properties"] for tool in toolkit.list_tools() + ) + for delta in ([1, 2], [0, 0, float("inf")]): + result = toolkit.execute_tool("move_delta", {"delta_xyz": delta}) + assert result.is_error + assert env.moves == [] + assert env.observation_calls == 1 + assert len(events) == 1 + result = toolkit.execute_tool("move_delta", {"delta_xyz": [0.01, 0, 0]}) + assert not result.is_error + assert "images" not in result.data + assert "image_cam_path" in result.data and "image_wrist_path" in result.data + assert len(result.images) == 2 + assert "states" in result.to_text() + assert len(events) == 2 + assert events[-1].record.command["action"] == "move_delta" + read = toolkit.execute_tool("view_env_state", {}) + np.testing.assert_equal( + read.data, + {key: value for key, value in result.data.items() if key != "agent_elapsed_s"}, + ) + assert read.images == result.images + assert len(events) == 2 + assert "finish" in {tool.name for tool in toolkit.list_tools()} + assert toolkit.finish_result is None + toolkit.close() + + +# Historical PR #172 schemas and return fields; do not regenerate from native tools. +_CONTRACT = json.loads( + (Path(__file__).parent / "fixtures/pre_native_tool_contracts.json").read_text() +) + + +def _prepare_contract_perception(state): + metadata = { + "observation_camera_map": {"main": "wrist_cam", "extra_0": "external_cam"}, + "cameras": { + name: {"intrinsic_K": [[100, 0, 2], [0, 100, 2], [0, 0, 1]]} + for name in ("wrist_cam", "external_cam") + }, + } + images = {name: state.load(f"{name}.png") for name in ("wrist", "camera")} + with state.record_step(state={"raw_base_state": state.get().state}) as step: + for name, image in images.items(): + state.save(f"{name}.png", image, step=step) + state.save( + f"{name}_depth.npy", np.full((4, 4), 0.5, dtype=np.float32), step=step + ) + state.save("camera_meta.json", metadata, step=step) + + +@pytest.mark.parametrize("case", _CONTRACT["cases"], ids=lambda case: case["name"]) +def test_normal_return_contract(case, tmp_path, monkeypatch): + monkeypatch.setattr(franka_toolkit, "get_output_dir", lambda: tmp_path) + toolkit = FrankaToolkit( + runtime_kwargs={ + "env": FakeEnv(), + "model": FakeModel(), + "task_description": "default task", + }, + dashboard_events=NullDashboardEventSink(), + memory=MemoryManager(tmp_path / "memory"), + ) + name = case["name"] + if name.startswith("back_project") or name == "view_perception_setup": + _prepare_contract_perception(toolkit.state) + monkeypatch.setattr( + runtime_config, + "_calibration_path", + Path(__file__).parent / "fixtures/hand_eye_calibration.json", + ) + result = toolkit.execute_tool(name, case["arguments"]) + assert not result.is_error + data = result.to_dict() + assert sorted(data) == case["fields"] + assert len(result.images) == case["image_count"] + if "result_fields" in case: + assert sorted(data["result"]) == case["result_fields"] + if "last_chunk_fields" in case: + assert sorted(data["result"]["last_chunk"]) == case["last_chunk_fields"] + if name == "finish": + assert data == {"_finish": True, **case["arguments"]} + assert toolkit.finish_result == case["arguments"] diff --git a/tests/unit_tests/robots/libero/conftest.py b/tests/unit_tests/robots/libero/conftest.py new file mode 100644 index 000000000..6a71f056e --- /dev/null +++ b/tests/unit_tests/robots/libero/conftest.py @@ -0,0 +1,138 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import numpy as np +import pytest + +from robots.libero.toolkit import LiberoToolkit +from rpent.dashboard.events import NullDashboardEventSink +from rpent.memory import MemoryManager +from rpent.robots.components.sam3_client import Sam3Result + + +class FakeEnv: + return_all_frames = False + + def __init__(self): + self.actions = [] + self.reset_calls = 0 + self.terminated = self.truncated = False + self.after_step = lambda: None + self.image = np.zeros((8, 8, 3), dtype=np.uint8) + self.pos = np.array([0.0, 0.0, 0.3]) + + def obs(self): + return { + "main_images": self.image, + "states": [*self.pos, 0, 0, 0, 0.04, -0.04], + "task_descriptions": "original task", + } + + def reset(self): + self.reset_calls += 1 + self.terminated = self.truncated = False + return self.obs(), {} + + def step(self, action): + self.actions.append(action.copy()) + self.after_step() + return self.obs(), 0, self.terminated, self.truncated, {} + + def chunk_step(self, actions, *, return_all_frames=False): + observations = [self.step(action)[0] for action in actions] + return ( + observations if return_all_frames else observations[-1], + 0, + np.zeros(len(actions), dtype=bool), + np.zeros(len(actions), dtype=bool), + {}, + ) + + def raw_obs(self): + return { + "robot0_eef_pos": self.pos, + "robot0_eef_quat": [1, 0, 0, 0], + "robot0_gripper_qpos": [0.04, -0.04], + "agentview_image": self.image, + "agentview_depth": np.ones((8, 8)), + "robot0_eye_in_hand_image": self.image, + "robot0_eye_in_hand_depth": np.ones((8, 8)), + } + + def get_task_language(self): + return "original task" + + def get_camera_meta(self, camera_name, height, width): + return { + "intrinsic_K": [[4, 0, 4], [0, 4, 4], [0, 0, 1]], + "extrinsic_cam2world": np.eye(4).tolist(), + } + + def render_camera(self, **kwargs): + return self.image, np.ones((8, 8)) + + +class FakeModel: + def __init__(self): + self.instructions = [] + + def predict(self, obs, *, options): + self.instructions.append(obs["task_descriptions"]) + assert options == {"mode": "eval"} + return np.zeros((3, 7), dtype=np.float32) + + +class FakeSam: + def segment(self, image, **kwargs): + return Sam3Result( + found=True, + mask=np.ones((8, 8), dtype=bool), + score=0.8, + box=[0, 0, 8, 8], + mask_shape=(8, 8), + ) + + +@pytest.fixture +def make_toolkit(tmp_path): + instances = [] + + def make(*, mode="evaluation", attempts=0, molmo_client=None, flywheel_config=None): + output = tmp_path / str(len(instances)) + env, model = FakeEnv(), FakeModel() + toolkit = LiberoToolkit( + runtime_kwargs={ + "env": env, + "model": model, + "sam3_client": FakeSam(), + "molmo_client": molmo_client, + "flywheel_config": flywheel_config, + }, + output_dir=output, + state_output_dir=output / "sessions" / "session_001", + memory=MemoryManager(output / "memory"), + dashboard_events=NullDashboardEventSink(), + mode=mode, + attempts_per_session=attempts, + ) + instances.append(toolkit) + return toolkit, env, model + + yield make + for toolkit in instances: + # Avoid video encoding in CPU contract tests. + toolkit._frames.clear() + toolkit.close() diff --git a/tests/unit_tests/robots/libero/fixtures/pre_native_tool_contracts.json b/tests/unit_tests/robots/libero/fixtures/pre_native_tool_contracts.json new file mode 100644 index 000000000..7a6a7b6f4 --- /dev/null +++ b/tests/unit_tests/robots/libero/fixtures/pre_native_tool_contracts.json @@ -0,0 +1,417 @@ +{ + "source_commit": "949f61ecea1c571168ca6a29a59f6d3207f6f947", + "schemas": { + "finish": { + "type": "object", + "properties": { + "status": { + "type": "string", + "description": "Outcome, e.g. 'success', 'failure', or 'stuck'." + }, + "summary": { + "type": "string", + "description": "Short natural-language summary of the run." + } + }, + "required": [ + "status", + "summary" + ] + }, + "reset": { + "type": "object", + "properties": { + "reason": { + "type": "string", + "description": "Why this episode is unrecoverable and what will change." + } + }, + "required": [ + "reason" + ] + }, + "view_env_state": { + "type": "object", + "properties": { + "step": { + "type": "integer", + "default": -1, + "description": "Step number; 0 = initial, -1 = latest." + } + } + }, + "move_to": { + "type": "object", + "properties": { + "xyz": { + "type": "array", + "description": "World-frame target [x, y, z] in meters", + "items": { + "type": "number" + }, + "minItems": 3, + "maxItems": 3 + }, + "gripper": { + "type": "number", + "description": "Gripper command: -1 open, +1 close (default -1)" + }, + "tol": { + "type": "number", + "description": "Position tolerance, m (default 0.012)" + }, + "step_clip": { + "type": "number", + "description": "Per-step \u0394xyz cap before action_scale, m (default 0.025)" + }, + "max_steps": { + "type": "integer", + "description": "Step budget (default 80)" + }, + "action_scale": { + "type": "number", + "description": "OSC action scale (default 0.05)" + }, + "target_yaw": { + "type": [ + "number", + "null" + ], + "description": "Optional world-frame yaw target in radians" + }, + "yaw_step_clip": { + "type": "number", + "description": "Per-step yaw clip, rad (default 0.10)" + } + }, + "required": [ + "xyz" + ] + }, + "pi0_pick": { + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "Pi0 prompt (e.g. 'pick up the akita black bowl')." + }, + "max_chunks": { + "type": "integer", + "description": "Action-chunk budget (default 24)" + }, + "lift_thresh": { + "type": "number", + "description": "EEF post-descent ascent threshold for success, m (default 0.05)" + }, + "gripper_closed_thresh": { + "type": "number", + "description": "Finger-separation closed threshold (default 0.06)" + } + }, + "required": [ + "prompt" + ] + }, + "pi0_doubled": { + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "Contact-skill prompt, e.g. 'turn on the stove'." + }, + "max_chunks": { + "type": "integer", + "description": "Action-chunk budget (default 20)" + } + }, + "required": [ + "prompt" + ] + }, + "release": { + "type": "object", + "properties": { + "max_steps": { + "type": "integer", + "description": "Step budget (default 20)" + } + } + }, + "set_gripper": { + "type": "object", + "properties": { + "gripper": { + "type": "number", + "description": "Gripper command: -1 open, +1 close (default -1)" + }, + "steps": { + "type": "integer", + "description": "Number of env steps (default 5)" + } + } + }, + "rotate_wrist": { + "type": "object", + "properties": { + "target_yaw": { + "type": [ + "number", + "null" + ], + "description": "Absolute world-frame yaw target, rad" + }, + "delta_yaw": { + "type": [ + "number", + "null" + ], + "description": "Relative yaw delta, rad" + }, + "gripper": { + "type": "number", + "description": "Gripper command held during rotation (default +1)" + }, + "max_steps": { + "type": "integer", + "description": "Step budget (default 40)" + }, + "tol": { + "type": "number", + "description": "Yaw tolerance, rad (default 0.02)" + }, + "step_clip": { + "type": "number", + "description": "Per-step yaw clip, rad (default 0.10)" + } + } + }, + "rotate_pitch": { + "type": "object", + "properties": { + "target_pitch": { + "type": [ + "number", + "null" + ], + "description": "Absolute world-frame pitch target, rad" + }, + "delta_pitch": { + "type": [ + "number", + "null" + ], + "description": "Relative pitch delta, rad" + }, + "gripper": { + "type": "number", + "description": "Gripper command held during rotation (default +1)" + }, + "max_steps": { + "type": "integer", + "description": "Step budget (default 40)" + }, + "tol": { + "type": "number", + "description": "Pitch tolerance, rad (default 0.02)" + }, + "step_clip": { + "type": "number", + "description": "Per-step pitch clip, rad (default 0.10)" + } + } + }, + "move_pose": { + "type": "object", + "properties": { + "xyz": { + "type": "array", + "description": "World-frame target [x, y, z] in meters", + "items": { + "type": "number" + }, + "minItems": 3, + "maxItems": 3 + }, + "target_pitch": { + "type": [ + "number", + "null" + ], + "description": "Absolute pitch target, rad" + }, + "target_yaw": { + "type": [ + "number", + "null" + ], + "description": "Absolute yaw target, rad" + }, + "gripper": { + "type": "number", + "description": "Gripper command held during the move (default -1)" + }, + "step_clip": { + "type": "number", + "description": "Per-step \u0394xyz cap, m (default 0.02)" + }, + "pitch_step": { + "type": "number", + "description": "Per-step pitch clip, rad (default 0.08)" + }, + "yaw_step": { + "type": "number", + "description": "Per-step yaw clip, rad (default 0.08)" + }, + "tol": { + "type": "number", + "description": "Position tolerance, m (default 0.012)" + }, + "ori_tol": { + "type": "number", + "description": "Orientation tolerance, rad (default 0.05)" + }, + "action_scale": { + "type": "number", + "description": "OSC action scale (default 0.05)" + }, + "max_steps": { + "type": "integer", + "description": "Step budget (default 150)" + } + }, + "required": [ + "xyz" + ] + }, + "view_camera_meta": { + "type": "object", + "properties": { + "camera": { + "type": "string", + "enum": [ + "agentview", + "wrist" + ], + "description": "Camera metadata to read (default agentview)." + }, + "step": { + "type": "integer", + "default": -1, + "description": "Metadata step to use; -1 = latest." + } + } + }, + "segment": { + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "Object/text prompt to segment." + }, + "camera": { + "type": "string", + "enum": [ + "agentview", + "wrist" + ], + "description": "Artifact camera to use (default agentview)." + }, + "step": { + "type": "integer", + "default": -1, + "description": "Step to segment; -1 = latest." + }, + "point": { + "type": [ + "array", + "null" + ], + "description": "Optional single positive point as [row, col]. Mutually exclusive with prompt.", + "items": { + "type": "integer" + }, + "minItems": 2, + "maxItems": 2 + }, + "min_score": { + "type": "number", + "description": "Minimum accepted mask score (default 0.2)." + } + } + }, + "back_project": { + "type": "object", + "properties": { + "row": { + "type": [ + "integer", + "null" + ], + "description": "Pixel row (0=top) in the selected resolution image." + }, + "col": { + "type": [ + "integer", + "null" + ], + "description": "Pixel column (0=left) in the selected resolution image." + }, + "step": { + "type": "integer", + "default": -1, + "description": "Depth/world-map step; 0 = initial, -1 = latest." + }, + "camera": { + "type": "string", + "enum": [ + "agentview", + "wrist" + ], + "description": "Camera to back-project from (default agentview)." + }, + "resolution": { + "type": "string", + "enum": [ + "high", + "low" + ], + "description": "Coordinate system for row/col (default high). Use low only when row/col came from the embedded/standard 256 image." + }, + "row_range": { + "type": [ + "array", + "null" + ], + "items": { + "type": "integer" + }, + "description": "Region mode: [r0, r1] pixel row window. Requires col_range." + }, + "col_range": { + "type": [ + "array", + "null" + ], + "items": { + "type": "integer" + }, + "description": "Region mode: [c0, c1] pixel col window. Requires row_range." + }, + "z_min": { + "type": [ + "number", + "null" + ], + "description": "Region mode: keep only pixels with world z >= z_min." + }, + "z_max": { + "type": [ + "number", + "null" + ], + "description": "Region mode: keep only pixels with world z <= z_max." + } + } + } + } +} diff --git a/tests/unit_tests/robots/libero/test_flash.py b/tests/unit_tests/robots/libero/test_flash.py new file mode 100644 index 000000000..934fb0e43 --- /dev/null +++ b/tests/unit_tests/robots/libero/test_flash.py @@ -0,0 +1,370 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +from robots.libero.flash.replay import ( + execute, + load, + locate, + pick_succeeded, + plans, + replay, +) +from robots.libero.robot_spec import FLASH_SUITES, _parse_config +from rpent.planner.flash import FlashPlanner +from rpent.robots.components.molmo_client import MolmoResult +from rpent.tools import ToolResult + + +def test_locate_uses_native_back_projection(make_toolkit): + grounded = [] + + def ground(image, query): + grounded.append((image, query)) + return MolmoResult(found=True, point_xy=(4, 4)) + + molmo = SimpleNamespace(ground=ground) + toolkit, env, _ = make_toolkit(molmo_client=molmo) + opening = toolkit.state.latest_step + + found = locate(toolkit.molmo_client, toolkit, opening, "agentview", "bowl") + + assert found is not None + assert found["xy"] == pytest.approx([0.0, 0.0]) + assert found["z_top"] == pytest.approx(1.0) + assert grounded[0][0].startswith(b"\x89PNG") + assert grounded[0][1] == "bowl" + assert toolkit.state.latest_step == opening + assert env.actions == [] + + +def test_flash_planner_executes_native_toolkit(make_toolkit): + toolkit, env, model = make_toolkit(molmo_client=SimpleNamespace()) + plan_dir = toolkit.memory.root / "flash" + plan_dir.mkdir(parents=True) + (plan_dir / "object_task_t0_plan.json").write_text( + json.dumps({"plan": [{"action": "pi0_pick", "arguments": {"prompt": "bowl"}}]}) + ) + (plan_dir / "object_task_t0_anchors.json").write_text('{"anchors": []}') + env.after_step = lambda: setattr(env, "terminated", True) + + result = FlashPlanner(recipe_tag="object_task_t0_s1", robot_name="libero").solve( + system_prompt="", user_message="", toolkit=toolkit, max_turns=1 + ) + + assert result.error is None + assert result.finish_result["status"] == "success" + assert result.stats["total_input_tokens"] == 0 + assert model.instructions == ["bowl"] + assert env.reset_calls == 1 + assert toolkit.solved() + + +def test_replay_relocates_segment_anchor_with_native_toolkit(make_toolkit): + toolkit, env, _ = make_toolkit() + env.after_step = lambda: setattr(env, "terminated", True) + + result = replay( + toolkit, + molmo=SimpleNamespace(), + program={ + "plan": [ + { + "action": "move_to", + "arguments": {"xyz": [0.0, 0.0, 0.7]}, + "anchor": "bowl", + "anchor_distance": 0.0, + "offset": [0.01, -0.02], + } + ], + "reference": {"bowl": [0.0, 0.0]}, + "locator_of": {"bowl": "segment"}, + }, + ) + + assert result == {"done": True, "anchors": 1, "plan": 1} + assert len(env.actions) == 1 + command = toolkit.state.latest_record().command + assert command["action"] == "move_to" + assert command["xyz"] == pytest.approx([-0.115, -0.145, 0.7]) + + +@pytest.mark.parametrize( + ("gripper_open_thresh", "descent_thresh", "success"), + [(0.003, 0.0, True), (0.05, 0.0, False), (0.003, 0.10, False)], +) +def test_native_pick_applies_flash_thresholds( + make_toolkit, monkeypatch, gripper_open_thresh, descent_thresh, success +): + toolkit, env, _ = make_toolkit() + + def lift(ctx, prompt): + env.pos[2] = 0.36 + obs = env.obs() + obs["states"][-2:] = [0.02, -0.02] + ctx.robot.set_obs(obs) + + monkeypatch.setattr("robots.libero.tools._vlm_chunk", lift) + result = toolkit.execute_tool( + "pi0_pick", + { + "prompt": "bowl", + "max_chunks": 1, + "gripper_open_thresh": gripper_open_thresh, + "descent_thresh": descent_thresh, + }, + ) + + assert not result.is_error + assert result.data["log"]["result"]["success"] is success + + +class _Toolkit: + def __init__(self, result: dict | None = None) -> None: + self.result = result or {} + self.state = SimpleNamespace(latest_step=0) + self.calls = [] + + def execute_tool(self, name: str, arguments: dict): + self.calls.append((name, arguments)) + return ToolResult(data=self.result, error=self.result.get("error")) + + def solved(self) -> bool: + return False + + +def test_execute_rejects_tool_error() -> None: + with pytest.raises(RuntimeError, match="move_to failed: unreachable"): + execute(_Toolkit({"error": "unreachable"}), "move_to", {}) + + +def test_pick_succeeded_uses_pi0_pick_success_contract() -> None: + wrapped = {"log": {"result": {"success": True, "peak_lift_m": 0.05}}} + + assert pick_succeeded(wrapped) is True + assert ( + pick_succeeded( + { + "log": { + "result": { + "success": False, + "peak_lift_m": 0.10, + "final_gripper_opening": 0.01, + } + } + } + ) + is False + ) + + +def test_replay_reuses_toolkit_opening_observation() -> None: + toolkit = _Toolkit() + toolkit.state.latest_step = 7 + + result = replay( + toolkit, + molmo=SimpleNamespace(), + program={"plan": [], "reference": {}, "locator_of": {}}, + ) + + assert result == {"done": False, "anchors": 0, "plan": 0} + + +def test_replay_passes_legacy_pick_thresholds_to_pi0_pick() -> None: + toolkit = _Toolkit({"success": True}) + + replay( + toolkit, + molmo=SimpleNamespace(), + program={ + "plan": [ + { + "action": "pi0_pick", + "arguments": { + "prompt": "pick up the bowl", + "lift_thresh": 0.08, + }, + } + ], + "reference": {}, + "locator_of": {}, + }, + ) + + assert toolkit.calls == [ + ( + "pi0_pick", + { + "prompt": "pick up the bowl", + "lift_thresh": 0.04, + "gripper_closed_thresh": 0.07, + "gripper_open_thresh": 0.003, + "descent_thresh": 0.0, + }, + ) + ] + + +def test_pick_retry_reuses_relocated_move_arguments() -> None: + class RetryToolkit(_Toolkit): + def execute_tool(self, name: str, arguments: dict): + self.calls.append((name, dict(arguments))) + if name == "segment": + result = {"world_xyz": [0.2, 0.1, 0.0]} + elif name == "pi0_pick": + result = {"success": False} + else: + result = {} + return ToolResult(data=result) + + toolkit = RetryToolkit() + replay( + toolkit, + molmo=SimpleNamespace(), + program={ + "plan": [ + { + "action": "move_to", + "arguments": {"xyz": [0.0, 0.0, 0.7], "gripper": -1}, + "anchor": "bowl", + "anchor_distance": 0.0, + "offset": [0.01, -0.02], + }, + {"action": "pi0_pick", "arguments": {"prompt": "pick up the bowl"}}, + ], + "reference": {"bowl": [0.0, 0.0]}, + "locator_of": {"bowl": "segment"}, + }, + ) + + retried_moves = [args for name, args in toolkit.calls if name == "move_to"] + assert len(retried_moves) == 3 + assert all(args["xyz"] == [0.21, 0.08, 0.7] for args in retried_moves) + + +def test_replay_stops_when_attached_anchor_is_not_located() -> None: + toolkit = _Toolkit() + notes = [] + + result = replay( + toolkit, + molmo=SimpleNamespace(), + program={ + "plan": [ + { + "action": "move_to", + "arguments": {"xyz": [0.3, 0.2, 0.7], "gripper": -1}, + "anchor": "bowl", + "anchor_distance": 0.0, + "offset": [0.01, -0.02], + }, + {"action": "release", "arguments": {}}, + ], + "reference": {"bowl": [0.0, 0.0]}, + "locator_of": {"bowl": "segment"}, + }, + note=notes.append, + ) + + assert result["done"] is False + assert [name for name, _ in toolkit.calls] == ["segment"] + assert any("unavailable; stopping replay" in note for note in notes) + + +def test_replay_propagates_toolkit_exceptions() -> None: + class FailingToolkit(_Toolkit): + def execute_tool(self, name: str, arguments: dict): + raise ConnectionError("RPC disconnected") + + with pytest.raises(ConnectionError, match="RPC disconnected"): + replay( + FailingToolkit(), + molmo=SimpleNamespace(), + program={ + "plan": [ + { + "action": "move_to", + "arguments": {"xyz": [0.1, 0.1, 0.7], "gripper": -1}, + } + ], + "reference": {}, + "locator_of": {}, + }, + ) + + +def test_flash_supports_all_libero_pro_task_and_swap_suites() -> None: + assert FLASH_SUITES == { + f"libero_{family}_{regime}" + for family in ("spatial", "object", "goal", "10") + for regime in ("task", "swap") + } + + +def test_non_flash_planner_rejects_molmo_endpoint() -> None: + args = SimpleNamespace( + suite="libero_object_swap", + task=0, + planner="api", + molmo_endpoint="http://127.0.0.1:8115", + ) + + with pytest.raises(ValueError, match="requires --planner flash"): + _parse_config(args) + + +def test_missing_plans_reports_selected_memory_without_downloading( + monkeypatch, tmp_path +) -> None: + sync_calls = [] + monkeypatch.setattr( + "rpent.memory.MemoryManager.sync", + lambda *args, **kwargs: sync_calls.append(kwargs), + ) + + with pytest.raises(FileNotFoundError, match="RLinf/RPent-memory") as error: + plans(tmp_path / "flash") + + assert "--cards" not in str(error.value) + assert str(tmp_path / "flash") in str(error.value) + assert sync_calls == [] + + +def test_plans_do_not_require_an_index(tmp_path) -> None: + root = tmp_path / "flash" + root.mkdir(parents=True) + (root / "object_swap_t0_plan.json").write_text('{"plan": []}') + + assert plans(root) == root + + +def test_load_reads_only_runtime_plan_fields(tmp_path) -> None: + (tmp_path / "object_swap_t0_plan.json").write_text('{"plan": []}') + (tmp_path / "object_swap_t0_anchors.json").write_text( + '{"anchors": [{"phrase": "bowl", "locator": "segment", ' + '"median_xy": [0.1, 0.2]}]}' + ) + + card = load(tmp_path, "object_swap_t0") + + assert card["plan"] == [] + assert card["reference"]["bowl"].tolist() == [0.1, 0.2] + assert card["locator_of"] == {"bowl": "segment"} diff --git a/tests/unit_tests/robots/libero/test_libero_integration.py b/tests/unit_tests/robots/libero/test_libero_integration.py index 801c877c6..cff521ec1 100644 --- a/tests/unit_tests/robots/libero/test_libero_integration.py +++ b/tests/unit_tests/robots/libero/test_libero_integration.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import threading from types import SimpleNamespace from unittest.mock import Mock @@ -21,8 +22,10 @@ from robots.libero import robot_spec from robots.libero import toolkit as libero_toolkit from robots.libero.flywheel import LIBERO_SPEC -from robots.libero.tools import LiberoPrimitives +from robots.libero.toolkit import LiberoRuntime +from robots.libero.tools import _step_env, _vlm_chunk from rpent.flywheel.episode import validate_episode +from rpent.tools import ToolContext def _obs(value: int) -> dict: @@ -65,22 +68,32 @@ def predict(self, observation, *, options): return np.ones((2, 7), np.float32) -def _primitives(env, config=None, molmo_client=None): - return LiberoPrimitives( +def _runtime(env, config=None, molmo_client=None): + return LiberoRuntime( env=env, model=_Model(), sam3_client=SimpleNamespace(), - check_cancelled=lambda: None, flywheel_config=config, molmo_client=molmo_client, ) +def _context(runtime): + return ToolContext( + robot=runtime, + state=None, + memory=None, + output_dir=None, + record_frame=lambda frame: None, + _cancel_event=threading.Event(), + ) + + @pytest.mark.parametrize("with_molmo", [False, True]) def test_collection_records_scripted_and_vla_actions(tmp_path, with_molmo): env = _Env() molmo = SimpleNamespace() if with_molmo else None - primitives = _primitives( + runtime = _runtime( env, { "root": tmp_path, @@ -90,16 +103,15 @@ def test_collection_records_scripted_and_vla_actions(tmp_path, with_molmo): }, molmo_client=molmo, ) - assert primitives.molmo_client is molmo - primitives.reset() - primitives.begin_primitive("move_to") - primitives._step_env(np.zeros(7)) - primitives.end_primitive() - primitives.begin_primitive("pi0_pick") - primitives._vlm_chunk("pick up the bowl") - primitives.end_primitive() - - path = primitives.finalize_flywheel() + assert runtime.molmo_client is molmo + runtime.reset() + ctx = _context(runtime) + runtime.execute_primitive("move_to", _step_env, ctx=ctx, action=np.zeros(7)) + runtime.execute_primitive( + "pi0_pick", _vlm_chunk, ctx=ctx, instruction="pick up the bowl" + ) + + path = runtime.finalize_flywheel() metadata = validate_episode(path, spec=LIBERO_SPEC) assert metadata["step_count"] == 3 assert metadata["training_step_count"] == 3 @@ -109,23 +121,21 @@ def test_collection_records_scripted_and_vla_actions(tmp_path, with_molmo): np.testing.assert_array_equal(data["primitive_id"], [0, 1, 1]) -def test_molmo_positional_argument_keeps_collection_disabled(): +def test_optional_molmo_keeps_collection_disabled(): molmo = SimpleNamespace() - primitives = LiberoPrimitives( - _Env(), _Model(), SimpleNamespace(), lambda: None, molmo - ) - primitives.reset() - assert primitives.molmo_client is molmo - assert primitives.finalize_flywheel() is None + runtime = LiberoRuntime(_Env(), _Model(), SimpleNamespace(), molmo) + runtime.reset() + assert runtime.molmo_client is molmo + assert runtime.finalize_flywheel() is None -def test_collection_disabled_keeps_fast_chunk_path(): +def test_collection_disabled_keeps_native_frame_recording(): env = _Env() - primitives = _primitives(env) - primitives.reset() - primitives._vlm_chunk("pick up the bowl") - assert env.chunk_return_all_frames is None - assert primitives.finalize_flywheel() is None + runtime = _runtime(env) + runtime.reset() + _vlm_chunk(_context(runtime), "pick up the bowl") + assert env.chunk_return_all_frames is True + assert runtime.finalize_flywheel() is None def test_dashboard_flywheel_config_belongs_to_unique_env(tmp_path, monkeypatch): @@ -159,22 +169,20 @@ def test_dashboard_flywheel_config_belongs_to_unique_env(tmp_path, monkeypatch): } -@pytest.mark.parametrize("failure", [None, "finalize", "stop", "save"]) +@pytest.mark.parametrize("failure", [None, "finalize", "save"]) def test_close_handles_collection_and_video_independently( tmp_path, monkeypatch, failure ): toolkit = libero_toolkit.LiberoToolkit.__new__(libero_toolkit.LiberoToolkit) frames = [_obs(0)["main_images"]] finalize = Mock(return_value=tmp_path / "episode") - stop = Mock(return_value=frames) save = Mock() if failure is not None: - {"finalize": finalize, "stop": stop, "save": save}[ + {"finalize": finalize, "save": save}[failure].side_effect = RuntimeError( failure - ].side_effect = RuntimeError(failure) - toolkit._primitives = SimpleNamespace( - finalize_flywheel=finalize, stop_recording=stop - ) + ) + toolkit._robot = SimpleNamespace(finalize_flywheel=finalize) + toolkit._frames = frames toolkit._state = SimpleNamespace(save=save) logger = Mock() monkeypatch.setattr(libero_toolkit, "logger", logger) @@ -182,11 +190,7 @@ def test_close_handles_collection_and_video_independently( toolkit.close() finalize.assert_called_once_with() - stop.assert_called_once_with() - if failure == "stop": - save.assert_not_called() - else: - save.assert_called_once_with("episode.mp4", frames, step=None, fps=20) + save.assert_called_once_with("episode.mp4", frames, step=None, fps=20) if failure == "finalize": logger.info.assert_not_called() else: @@ -194,3 +198,24 @@ def test_close_handles_collection_and_video_independently( "flywheel episode finalized: %s", tmp_path / "episode" ) assert logger.warning.call_count == (failure is not None) + + +@pytest.mark.parametrize("fails", [False, True]) +def test_native_action_brackets_flywheel_recording(make_toolkit, monkeypatch, fails): + from robots.libero import flywheel + + writer = Mock() + monkeypatch.setattr(flywheel, "create_episode_writer", lambda config, obs: writer) + toolkit, env, _ = make_toolkit(flywheel_config={"enabled": True}) + if fails: + monkeypatch.setattr(env, "step", Mock(side_effect=RuntimeError("step failed"))) + + result = toolkit.execute_tool("release", {}) + + assert result.is_error is fails + writer.begin_primitive.assert_called_once_with("release") + writer.end_primitive.assert_called_once_with() + if fails: + writer.add_transition.assert_not_called() + else: + assert writer.add_transition.call_count > 0 diff --git a/tests/unit_tests/robots/libero/test_libero_toolkit_contracts.py b/tests/unit_tests/robots/libero/test_libero_toolkit_contracts.py index 087bed807..dbfb92905 100644 --- a/tests/unit_tests/robots/libero/test_libero_toolkit_contracts.py +++ b/tests/unit_tests/robots/libero/test_libero_toolkit_contracts.py @@ -12,190 +12,362 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Offline contracts for the LIBERO toolkit.""" - from __future__ import annotations -from pathlib import Path -from types import SimpleNamespace -from typing import Any +import json +import threading +from concurrent.futures import ThreadPoolExecutor import pytest -from robots.libero import robot_spec, toolkit +from robots.libero import robot_spec from rpent.dashboard.events import NullDashboardEventSink -from rpent.memory import MemoryManager from rpent.robots import RunConfig -from rpent.tools.toolkit import Toolkit, _is_readonly -from rpent.utils import templates +from rpent.robots.components.sam3_client import Sam3Result -COMMON_TOOLS = {"read_text_file", "write_text_file", "list_dir", "finish"} -EVALUATION_TOOLS = COMMON_TOOLS | { - "view_env_state", - "move_to", - "pi0_pick", - "pi0_doubled", - "release", - "set_gripper", - "rotate_wrist", - "rotate_pitch", - "move_pose", - "view_camera_meta", - "segment", - "back_project", -} +@pytest.mark.parametrize("reset_after_success", [False, True]) +def test_robot_finish_uses_cumulative_solved_state(make_toolkit, reset_after_success): + toolkit, env, _ = make_toolkit(mode="exploration", attempts=4) + assert toolkit._robot.mode == "exploration" + assert not toolkit.solved() + env.after_step = lambda: setattr(env, "terminated", True) + assert not toolkit.execute_tool("set_gripper", {"steps": 1}).is_error + assert toolkit.solved() and toolkit._robot.solved + if reset_after_success: + assert not toolkit.execute_tool( + "reset", {"reason": "record another attempt"} + ).is_error + assert not env.terminated + assert toolkit.solved() and toolkit._robot.solved + record_count = len(toolkit.state.records()) + result = toolkit.execute_tool("finish", {"status": "success", "summary": "完成"}) + assert not result.is_error + assert toolkit.finish_result == {"status": "success", "summary": "完成"} + assert len(toolkit.state.records()) == record_count -def _record(step_idx: int = 0) -> SimpleNamespace: - return SimpleNamespace(step_idx=step_idx, terminated=False) +@pytest.mark.parametrize("mode, attempts", [("evaluation", 4), ("exploration", 0)]) +def test_robot_finish_preserves_unrestricted_modes(make_toolkit, mode, attempts): + toolkit, _, _ = make_toolkit(mode=mode, attempts=attempts) + assert not toolkit.solved() + result = toolkit.execute_tool("finish", {"status": "stuck", "summary": "无法完成"}) + assert result.data == {"_finish": True, "status": "stuck", "summary": "无法完成"} + assert not result.is_error + assert not toolkit.solved() -def _tool_names(robot_toolkit: Toolkit) -> set[str]: - return {spec["name"] for spec in robot_toolkit.get_tools_spec()} +def test_modes_filters_directories_and_exploration_guards(make_toolkit): + evaluation, env, _ = make_toolkit() + assert env.reset_calls == 1 + assert "reset" not in {tool.name for tool in evaluation.list_tools()} + assert evaluation.execute_tool("reset", {"reason": "again"}).error.startswith( + "Unknown tool: " + ) + exploration, env, _ = make_toolkit(mode="exploration", attempts=3) + assert "reset" in {t.name for t in exploration.list_tools()} + assert exploration.execute_tool( + "finish", {"status": "success", "summary": "early"} + ).is_error + assert exploration.finish_result is None + for attempt in [2, 3]: + reset = exploration.execute_tool("reset", {"reason": "new strategy"}) + assert not reset.is_error + assert reset.data["log"]["result"]["attempt"] == attempt + assert reset.data["step"] == attempt - 1 + assert exploration.execute_tool("reset", {"reason": "too many"}).is_error + assert env.reset_calls == 3 + accepted = exploration.execute_tool( + "finish", {"status": "failure", "summary": "spent"} + ) + assert accepted.data == {"_finish": True, "status": "failure", "summary": "spent"} + assert exploration.finish_result == {"status": "failure", "summary": "spent"} -def _readonly_names(robot_toolkit: Toolkit) -> set[str]: - return { - name - for name, (_, handler) in robot_toolkit._tools.items() - if _is_readonly(handler) - } +@pytest.mark.parametrize( + ("name", "args"), + [ + ("move_to", {"xyz": [0, 0, 0.3]}), + ("move_pose", {"xyz": [0, 0, 0.3]}), + ("rotate_wrist", {"target_yaw": 0}), + ("rotate_pitch", {"target_pitch": 0}), + ("move_pose", {"xyz": [1, 1, 1], "max_steps": 0}), + ("move_to", {"xyz": [1, 1, 1], "max_steps": 0}), + ], +) +def test_already_reached_and_zero_budget_report_zero_actions(make_toolkit, name, args): + toolkit, env, _ = make_toolkit() + result = toolkit.execute_tool(name, args) + assert not result.is_error + assert result.data["log"]["result"]["steps_used"] == 0 + assert result.data["step"] == 1 + assert env.actions == [] + assert len(result.images) == 3 + assert all(image.startswith(b"\x89PNG\r\n\x1a\n") for image in result.images) -def _run_config(memory_dir: Path, *, recipe_tag: str = "cell-s0") -> RunConfig: - return RunConfig( - recipe_tag=recipe_tag, - output_dir=memory_dir.parent / "run", - prompt_vars={"memory_dir": str(memory_dir)}, - task_desc={}, +@pytest.mark.parametrize( + "value", [float("nan"), float("inf"), -float("inf"), "NaN", "Infinity", "-Infinity"] +) +def test_nonfinite_motion_arguments_fail_before_execution(make_toolkit, value): + toolkit, env, _ = make_toolkit() + result = toolkit.execute_tool( + "move_to", {"xyz": [0, 0, 0.3], "tol": value, "max_steps": 0} ) + assert result.is_error + assert "finite_number" in result.error + assert env.actions == [] + assert toolkit.state.latest_step == 0 + assert not toolkit.execute_tool("view_env_state", {}).is_error -def test_toolkit_factory_configures_memory_access_by_mode( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - captured: list[dict[str, Any]] = [] - def fake_toolkit(**kwargs: Any) -> SimpleNamespace: - captured.append(kwargs) - return SimpleNamespace(**kwargs) +@pytest.mark.parametrize( + ("name", "args", "field"), + [ + ("move_to", {"xyz": [1, 0, 0.3], "max_steps": 4}, "steps_used"), + ("move_pose", {"xyz": [1, 0, 0.3], "max_steps": 4}, "steps_used"), + ("rotate_wrist", {"target_yaw": 1, "max_steps": 4}, "steps_used"), + ("rotate_pitch", {"target_pitch": 1, "max_steps": 4}, "steps_used"), + ("set_gripper", {"steps": 4}, "steps"), + ("release", {"max_steps": 4}, "steps_used"), + ], +) +def test_early_termination_counts_only_sent_actions(make_toolkit, name, args, field): + toolkit, env, _ = make_toolkit() + env.after_step = lambda: setattr(env, "terminated", True) + result = toolkit.execute_tool(name, args) + assert not result.is_error + assert result.data["log"]["result"][field] == len(env.actions) == 1 + assert toolkit.solved() + assert result.data["terminated"] - monkeypatch.setattr(toolkit, "LiberoToolkit", fake_toolkit) - memory_dir = tmp_path / "libero-memory" - config = _run_config(memory_dir) - evaluation = robot_spec.get_toolkit( - runtime_kwargs={"env": "evaluation"}, - dashboard_events=NullDashboardEventSink(), - config=config, - ) - exploration = robot_spec.get_toolkit( - runtime_kwargs={"env": "exploration"}, - dashboard_events=NullDashboardEventSink(), - config=config, - mode="exploration", - attempts_per_session=2, - state_output_dir=tmp_path / "state", - ) +ACTION_CASES = [ + ("move_to", {"xyz": [1, 0, 0.3], "max_steps": 4}, 1), + ("move_pose", {"xyz": [1, 0, 0.3], "max_steps": 4}, 1), + ("rotate_wrist", {"target_yaw": 1, "max_steps": 4}, 1), + ("rotate_pitch", {"target_pitch": 1, "max_steps": 4}, 1), + ("set_gripper", {"steps": 4}, 1), + ("release", {"max_steps": 4}, 1), + ("pi0_pick", {"prompt": "pick bowl", "max_chunks": 2}, 3), + ("pi0_doubled", {"prompt": "touch bowl", "max_chunks": 2}, 3), +] - assert evaluation.memory.root == memory_dir.resolve() - assert exploration.memory.root == memory_dir.resolve() - evaluation_write = evaluation.memory.get_common_tool_bindings()["write_text_file"][ - 1 - ] - exploration_write = exploration.memory.get_common_tool_bindings()[ - "write_text_file" - ][1] - own_draft = memory_dir / "_internal" / "inbox" / config.recipe_tag / "draft.md" - with pytest.raises(PermissionError, match="writing to memory is denied"): - evaluation_write(str(own_draft), "draft") - assert exploration_write(str(own_draft), "draft")["bytes_written"] == 5 - assert captured[0]["mode"] == "evaluation" - assert captured[1]["mode"] == "exploration" - assert captured[1]["attempts_per_session"] == 2 - - -def test_toolkit_modes_construct_with_fake_primitives( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, - fake_single_arm_primitives: type[Any], -) -> None: - dumped: list[Any] = [] - monkeypatch.setattr( - templates, "default_variables", lambda: {"output_dir": "/offline/output"} - ) - monkeypatch.setattr( - toolkit.libero_tools, - "LiberoPrimitives", - fake_single_arm_primitives, - ) - monkeypatch.setattr( - toolkit.libero_tools, - "dump_state", - lambda primitives, state, log: dumped.append(primitives) or _record(), - ) - evaluation = toolkit.LiberoToolkit( - runtime_kwargs={"env_client": object()}, - dashboard_events=NullDashboardEventSink(), - memory=MemoryManager(tmp_path / "evaluation-memory"), - mode="evaluation", - state_output_dir=tmp_path / "evaluation", - ) - exploration = toolkit.LiberoToolkit( - runtime_kwargs={"env_client": object()}, - dashboard_events=NullDashboardEventSink(), - memory=MemoryManager( - tmp_path / "exploration-memory", - memory_access="inbox_write", - inbox_cell_tag="offline-cell", - ), - mode="exploration", - attempts_per_session=3, - state_output_dir=tmp_path / "exploration", - ) +@pytest.mark.parametrize(("name", "args", "completed"), ACTION_CASES) +def test_cancellation_preserves_partial_execution_capture_and_resume( + make_toolkit, name, args, completed +): + toolkit, env, _ = make_toolkit() + stepped, release_step = threading.Event(), threading.Event() - assert _tool_names(evaluation) == EVALUATION_TOOLS - assert _tool_names(exploration) == EVALUATION_TOOLS | {"reset"} - assert _readonly_names(evaluation) == COMMON_TOOLS | { - "view_env_state", - "view_camera_meta", - "segment", - "back_project", + def after_step(): + stepped.set() + assert release_step.wait(3) + + env.after_step = after_step + with ThreadPoolExecutor(max_workers=2) as pool: + action = pool.submit(toolkit.execute_tool, name, args) + assert stepped.wait(3) + cancellation = pool.submit(toolkit.cancel_active_and_wait) + # Confirm cancellation was delivered before allowing another action boundary. + assert toolkit._active_operation.cancel_event.wait(3) + release_step.set() + result = action.result(3) + cancellation.result(3) + assert result.is_error and result.error == "Tool call cancelled." + assert len(env.actions) == toolkit._robot.executed_steps == completed + assert toolkit.state.latest_record().result == {"error": result.error} + assert result.data["step"] == 1 + env.after_step = lambda: None + assert not toolkit.execute_tool("set_gripper", {"steps": 1}).is_error + assert len(env.actions) == completed + 1 + + +@pytest.mark.parametrize(("name", "args", "completed"), ACTION_CASES) +def test_action_failures_keep_completed_steps_and_capture( + make_toolkit, monkeypatch, name, args, completed +): + toolkit, env, _ = make_toolkit() + step = env.step + + def fail_after_completed(action): + if toolkit._robot.executed_steps >= completed: + raise TypeError("driver defect") + return step(action) + + monkeypatch.setattr(env, "step", fail_after_completed) + result = toolkit.execute_tool(name, args) + assert result.is_error + assert result.data["log"]["result"] == {} + assert toolkit._robot.executed_steps == completed + assert toolkit.state.latest_record().result == {"error": result.error} + assert len(env.actions) == completed + assert result.data["step"] == 1 + assert len(result.images) == 3 + + +def test_model_results_keep_original_observation_shape_and_full_disk_history( + make_toolkit, +): + toolkit, _, _ = make_toolkit() + result = toolkit.execute_tool("rotate_wrist", {}) + model_text = result.to_text() + assert model_text.count(result.error) == 1 + assert "error" not in result.data["log"]["result"] + payload = json.loads(model_text) + assert set(payload) == { + "step", + "terminated", + "truncated", + "state", + "artifacts", + "task_language", + "log", + "agent_elapsed_s", + "error", } - assert _readonly_names(exploration) == _readonly_names(evaluation) - assert len(dumped) == 2 - assert all( - instance.reset_calls == 1 for instance in fake_single_arm_primitives.instances - ) - assert all( - instance.recording_started for instance in fake_single_arm_primitives.instances + assert payload["error"] == result.error + assert payload["log"]["result"] == {"name": "rotate_wrist"} + # Rendering must not strip the error from the stored observation/history. + assert toolkit.state.latest_record().result["error"] == result.error + observed = toolkit.execute_tool("view_env_state", {}) + historical = json.loads(observed.to_text()) + assert "error" not in historical + assert historical["log"]["result"] == { + "name": "rotate_wrist", + "error": result.error, + } + + manifest = json.loads((toolkit.state._output_dir / "states.json").read_text()) + record = manifest["steps"][-1] + assert record["command"] == { + "action": "rotate_wrist", + **toolkit._tools.get("rotate_wrist").args_schema().model_dump(), + } + assert record["result"] == {"name": "rotate_wrist", "error": result.error} + assert record["elapsed_s"] >= 0 + + +def test_validation_precedes_execution_and_uses_model_defaults(make_toolkit): + toolkit, env, _ = make_toolkit(mode="exploration") + assert toolkit.execute_tool("move_to", {"xyz": [0, 1]}).error.startswith( + "Invalid arguments for " ) - assert all( - callable(instance.kwargs["check_cancelled"]) - for instance in fake_single_arm_primitives.instances + assert toolkit.execute_tool("reset", {}).error.startswith("Invalid arguments for ") + assert len(toolkit.state.records()) == 1 + rejected = toolkit.execute_tool( + "set_gripper", {"steps": "2", "ctx": "ignored", "extra": 1} ) + assert rejected.is_error + assert env.actions == [] + assert len(toolkit.state.records()) == 1 + result = toolkit.execute_tool("set_gripper", {"steps": "2"}) + assert not result.is_error + assert result.data["log"]["result"]["steps"] == 2 + assert len(env.actions) == 2 - refused = exploration.execute_tool( - "finish", {"status": "failure", "summary": "first attempt"} - ) - assert refused.result["error"] == "finish refused" - assert refused.is_finish is False - exploration.get_env_state = lambda *, command, result, elapsed_s: dict(result) - assert ( - exploration.execute_tool("reset", {"reason": "new approach"}).result["attempt"] - == 2 +def test_segment_saves_artifacts_without_capturing_an_observation(make_toolkit): + toolkit, env, _ = make_toolkit() + definition = next(t for t in toolkit.list_tools() if t.name == "segment") + assert definition.readonly + for index in [0, 1]: + result = toolkit.execute_tool("segment", {"prompt": "bowl", "step": 0}) + assert not result.is_error + data = result.data + assert data["step"] == 0 + assert data["world_xyz"] == pytest.approx([-0.125, -0.125, 1.0]) + assert data["segment_artifact"] == f"segment_{index:02d}.json" + assert toolkit.state.load(data["segment_artifact"], step=0)["prompt"] == "bowl" + assert len(result.images) == 1 + assert result.images[0].startswith(b"\x89PNG\r\n\x1a\n") + assert len(toolkit.state.records()) == 1 + assert env.actions == [] + assert not toolkit.execute_tool("back_project", {"row": 4, "col": 4}).is_error + assert not toolkit.execute_tool("view_camera_meta", {}).is_error + + +@pytest.mark.parametrize("save_fails", [False, True]) +def test_segment_errors_appear_once_and_keep_diagnostics( + make_toolkit, monkeypatch, save_fails +): + toolkit, env, _ = make_toolkit() + reason = "SAM3 found no matching mask" + monkeypatch.setattr( + toolkit._robot._sam3_client, + "segment", + lambda *args, **kwargs: Sam3Result(found=False, reason=reason), ) - assert ( - exploration.execute_tool("reset", {"reason": "third approach"}).result[ - "attempt" - ] - == 3 + if save_fails: + save = toolkit.state.save + + def fail_segment_save(name, *args, **kwargs): + if name.startswith("segment_"): + return None + return save(name, *args, **kwargs) + + monkeypatch.setattr(toolkit.state, "save", fail_segment_save) + + result = toolkit.execute_tool("segment", {"prompt": "bowl"}) + data = result.data + assert not ({"code", "segmentation_error", "world_error"} & data.keys()) + assert data["found"] is False + assert data["world_xyz"] is None + assert result.to_text().count(reason) == 1 + assert len(toolkit.state.records()) == 1 + assert env.actions == [] + if save_fails: + assert result.is_error + assert json.loads(result.error.split("\n", 1)[1]) == { + "segmentation_error": reason + } + assert "segment_artifact" not in data + else: + assert result.is_error + assert result.error == reason + assert toolkit.state.load(data["segment_artifact"], step=0)["error"] == reason + + +@pytest.mark.parametrize("name", ["pi0_pick", "pi0_doubled"]) +def test_vla_prompt_chunks_and_unsolved_result_are_preserved(make_toolkit, name): + toolkit, env, model = make_toolkit() + result = toolkit.execute_tool(name, {"prompt": "touch bowl", "max_chunks": 2}) + assert not result.is_error + assert result.data["log"]["result"]["success"] is False + assert result.data["log"]["result"]["chunks_used"] == 2 + assert model.instructions == ["touch bowl", "touch bowl"] + assert toolkit._robot._last_obs["task_descriptions"] == "original task" + assert len(env.actions) == toolkit._robot.executed_steps == 6 + + +def test_factory_binds_memory_permissions_and_task_root(monkeypatch, tmp_path): + from robots.libero import toolkit as module + + captured = [] + monkeypatch.setattr( + module, "LiberoToolkit", lambda **kwargs: captured.append(kwargs) or kwargs ) - allowed = exploration.execute_tool( - "finish", {"status": "failure", "summary": "budget spent"} + config = RunConfig( + recipe_tag="cell", + output_dir=tmp_path / "run", + prompt_vars={"memory_dir": str(tmp_path / "memory")}, + task_desc={}, ) - assert allowed.is_finish is True + for mode in ["evaluation", "exploration"]: + result = robot_spec.get_toolkit( + runtime_kwargs={"env": "offline"}, + dashboard_events=NullDashboardEventSink(), + config=config, + mode=mode, + state_output_dir=tmp_path / "state", + ) + assert result["runtime_kwargs"] == {"env": "offline"} + assert result["output_dir"] == config.output_dir + path = tmp_path / "memory" / "_internal" / "inbox" / "cell" / "draft.md" + if mode == "evaluation": + with pytest.raises(PermissionError): + result["memory"].authorize_write(path) + else: + assert result["memory"].authorize_write(path) == path diff --git a/tests/unit_tests/robots/robocasa/conftest.py b/tests/unit_tests/robots/robocasa/conftest.py new file mode 100644 index 000000000..f5cbb3760 --- /dev/null +++ b/tests/unit_tests/robots/robocasa/conftest.py @@ -0,0 +1,199 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Small CPU fakes for the RoboCasa client boundaries.""" + +from __future__ import annotations + +from copy import deepcopy +from types import SimpleNamespace + +import numpy as np +import pytest + +from robots.robocasa.robot_spec import get_toolkit +from rpent.dashboard.events import NullDashboardEventSink +from rpent.robots import RunConfig + + +class FakeEnv: + def __init__(self): + self.actions = [] + self.reset_calls = 0 + self.success = False + self.language = "Open the left drawer." + self.on_step = lambda: None + self.eef_pos = np.array([0.0, 0.0, 1.0]) + self.eef_quat = np.array([0.0, 0.0, 0.0, 1.0]) + self.gripper_qpos = np.array([0.02, -0.02]) + self.base_pos = np.zeros(3) + self.base_yaw = 0.0 + + @property + def current_raw_obs(self): + return { + "language": self.language, + "robot0_base_pos": self.base_pos.copy(), + "robot0_base_quat": np.array( + [ + 0, + 0, + np.sin(self.base_yaw / 2), + np.cos(self.base_yaw / 2), + ] + ), + "robot0_gripper_qpos": self.gripper_qpos.copy(), + "robot0_base_to_eef_pos": self.eef_pos - self.base_pos, + "robot0_base_to_eef_quat": self.eef_quat.copy(), + } + + @property + def terminated(self): + return self.success + + def reset(self): + self.reset_calls += 1 + self.eef_pos = np.array([0.0, 0.0, 1.0]) + self.base_pos = np.zeros(3) + self.base_yaw = 0.0 + self.success = False + + def step(self, action): + self.actions.append(action.copy()) + self.eef_pos += action[:3] * 0.05 + self.base_pos[:2] += ( + np.array( + [ + np.cos(self.base_yaw) * action[7] + - np.sin(self.base_yaw) * action[8], + np.sin(self.base_yaw) * action[7] + + np.cos(self.base_yaw) * action[8], + ] + ) + * 0.01 + ) + self.base_yaw += action[9] * 0.1 + self.on_step() + + def render_camera(self, camera_name, height=4, width=4, depth=False): + colors = {"agentview": 1, "navview": 2, "wrist": 3} + rgb = np.full((4, 4, 3), colors.get(camera_name, 4), dtype=np.uint8) + rgb[..., 0] = len(self.actions) % 256 + return (rgb, np.ones((4, 4))) if depth else rgb + + def world_map(self, *args): + world = np.ones((4, 4, 3)) + world[..., 2] = 0.9 + return world + + def get_camera_meta(self, camera): + return {"camera": camera} + + def get_task_language(self): + return self.language + + def get_task_progress(self): + return {"steps": len(self.actions)} + + def get_success_criteria_text(self): + return "offline success criteria" + + def check_success(self): + return self.success + + def grasp_contact(self): + return False, None + + def reassemble_env_action(self, action): + return np.concatenate( + [ + action[name] + for name in ( + "action.end_effector_position", + "action.end_effector_rotation", + "action.gripper_close", + "action.base_motion", + "action.control_mode", + ) + ] + ) + + +class FakeModel: + def __init__(self): + self.calls = [] + self.resets = 0 + self.on_predict = lambda: None + + def get_modality_config(self): + return {"video_delta_indices": [-2, 0], "hist_maxlen": 3} + + def predict(self, obs, options): + self.calls.append((deepcopy(obs), deepcopy(options))) + self.on_predict() + return { + "action.end_effector_position": np.zeros((1, 2, 3)), + "action.end_effector_rotation": np.zeros((1, 2, 3)), + "action.gripper_close": np.ones((1, 2, 1)), + "action.base_motion": np.ones((1, 2, 4)), + "action.control_mode": np.ones((1, 2, 1)), + } + + def reset_session(self): + self.resets += 1 + + +@pytest.fixture +def make_toolkit(monkeypatch, tmp_path): + for name in ( + "RLDX_MAX_CHUNKS", + "RLDX_ACTION_STEPS_PER_CHUNK", + "RLDX_SETTLE_PATIENCE", + "RLDX_ALLOW_RESET", + "RLDX_KEEP_HEAVY_NPY", + "RLDX_VIDEO_DIR", + ): + monkeypatch.delenv(name, raising=False) + runs = [] + + def make(**settings): + for name, value in settings.items(): + monkeypatch.setenv(name, str(value)) + env = FakeEnv() + model = FakeModel() + output_dir = tmp_path / f"run-{len(runs)}" + config = RunConfig( + recipe_tag="OpenDrawer_s1", + output_dir=output_dir, + prompt_vars={"memory_dir": str(tmp_path / "memory")}, + task_desc={}, + ) + robot_toolkit = get_toolkit( + runtime_kwargs={"env": env, "model": model, "hi_res": 8}, + dashboard_events=NullDashboardEventSink(), + config=config, + ) + # The simulator package owns action conversion; these tests exercise the + # rollout, frame history, and RPC-independent toolkit execution. + robot_toolkit._robot._rldx._unmap = lambda action: action + run = SimpleNamespace( + toolkit=robot_toolkit, env=env, model=model, config=config + ) + runs.append(run) + return run + + yield make + for run in runs: + run.toolkit._frames.clear() + run.toolkit.close() diff --git a/tests/unit_tests/robots/robocasa/fixtures/pre_native_tool_contracts.json b/tests/unit_tests/robots/robocasa/fixtures/pre_native_tool_contracts.json new file mode 100644 index 000000000..81933125f --- /dev/null +++ b/tests/unit_tests/robots/robocasa/fixtures/pre_native_tool_contracts.json @@ -0,0 +1,482 @@ +{ + "source_commit": "014a0fa97f69c991ee5e0f62f14e1d6f89c3dcd7", + "schemas": { + "move_to": { + "type": "object", + "properties": { + "xyz": { + "type": "array", + "description": "World-frame target [x, y, z] in meters", + "items": { + "type": "number" + }, + "minItems": 3, + "maxItems": 3 + }, + "gripper": { + "type": [ + "number", + "string" + ], + "description": "Gripper: +1 close, -1 open, or 'hold' to maintain current finger width (default 'hold')" + }, + "step_clip": { + "type": "number", + "description": "Per-step dxyz cap, m (default 0.02)" + }, + "max_steps": { + "type": "integer", + "description": "Step budget (default 200)" + }, + "tol": { + "type": "number", + "description": "Position tolerance, m (default 0.012)" + } + }, + "required": [ + "xyz" + ] + }, + "move_delta": { + "type": "object", + "properties": { + "dxyz": { + "type": "array", + "description": "Relative displacement [dx, dy, dz] in meters", + "items": { + "type": "number" + }, + "minItems": 3, + "maxItems": 3 + }, + "gripper": { + "type": [ + "number", + "string" + ], + "description": "Gripper: +1 close, -1 open, or 'hold' (default 'hold')" + }, + "step_clip": { + "type": "number", + "description": "Per-step dxyz cap, m (default 0.02)" + }, + "max_steps": { + "type": "integer", + "description": "Step budget (default 80)" + } + }, + "required": [ + "dxyz" + ] + }, + "rotate_pitch": { + "type": "object", + "properties": { + "target_pitch": { + "type": "number", + "description": "Absolute pitch target, radians (clamped +/-1.5; default 0.6)" + }, + "gripper": { + "type": "number", + "description": "Gripper command held during rotation (default +1)" + }, + "n": { + "type": "integer", + "description": "Number of env steps for the rotation (default 12)" + } + } + }, + "set_gripper": { + "type": "object", + "properties": { + "gripper": { + "type": "number", + "description": "Gripper command: +1 close, -1 open (default +1)" + }, + "steps": { + "type": "integer", + "description": "Number of env steps to hold (default 10)" + } + } + }, + "release": { + "type": "object", + "properties": { + "steps": { + "type": "integer", + "description": "Number of env steps (default 10)" + } + } + }, + "scripted_grasp": { + "type": "object", + "properties": { + "xyz": { + "type": "array", + "description": "World-frame grasp target [x, y, z] in meters", + "items": { + "type": "number" + }, + "minItems": 3, + "maxItems": 3 + }, + "approach_z": { + "type": "number", + "description": "Z offset above target before descent, m (default 0.10)" + }, + "grasp_z_offset": { + "type": "number", + "description": "Z offset at grasp point (default 0.0; negative = below target)" + }, + "step_clip": { + "type": "number", + "description": "Per-step dxyz cap during descent, m (default 0.02)" + } + }, + "required": [ + "xyz" + ] + }, + "rldx_skill": { + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "Complete live task_language, copied verbatim" + }, + "base_clip": { + "type": [ + "number", + "null" + ], + "description": "Base motion magnitude cap (default null = no clamp)" + }, + "max_chunks": { + "type": "integer", + "description": "Action-chunk budget (default 70)" + }, + "force_reset": { + "type": "boolean", + "description": "Force VLA frame history reset (default False)" + }, + "n_action_steps": { + "type": "integer", + "description": "Actions per VLA chunk (default 8)" + }, + "settle_patience": { + "type": "integer", + "description": "Settle step budget before declaring done (default 999; do NOT set small)" + }, + "settle_eps": { + "type": "number", + "description": "Settle position tolerance, m (default 0.012)" + } + }, + "required": [ + "prompt" + ] + }, + "rldx_arm": { + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "Complete live task_language, copied verbatim" + }, + "base_clip": { + "type": [ + "number", + "null" + ], + "description": "Base motion magnitude cap (default 0.1 = small)" + }, + "max_chunks": { + "type": "integer", + "description": "Action-chunk budget (default 70)" + }, + "force_reset": { + "type": "boolean", + "description": "Force VLA frame history reset (default False)" + }, + "n_action_steps": { + "type": "integer", + "description": "Actions per VLA chunk (default 8)" + }, + "settle_patience": { + "type": "integer", + "description": "Settle step budget before declaring done (default 999; do NOT set small)" + }, + "settle_eps": { + "type": "number", + "description": "Settle position tolerance, m (default 0.012)" + } + }, + "required": [ + "prompt" + ] + }, + "navigate_to": { + "type": "object", + "properties": { + "xy": { + "type": "array", + "description": "World-frame target [x, y] in meters (z ignored if provided)", + "items": { + "type": "number" + }, + "minItems": 2, + "maxItems": 2 + }, + "tol": { + "type": "number", + "description": "Distance threshold to stop, m (default 0.20)" + }, + "max_steps": { + "type": "integer", + "description": "Step budget (default 300)" + }, + "gripper": { + "type": [ + "number", + "string" + ], + "description": "Gripper while driving: +1 close, -1 open, or 'hold' (default 'hold')" + } + }, + "required": [ + "xy" + ] + }, + "move_base": { + "type": "object", + "properties": { + "forward": { + "type": "number", + "description": "Forward velocity, [-1, 1] (default 0)" + }, + "lateral": { + "type": "number", + "description": "Lateral / strafe velocity, [-1, 1] (default 0)" + }, + "turn": { + "type": "number", + "description": "Yaw rotation velocity, [-1, 1] (default 0)" + }, + "steps": { + "type": "integer", + "description": "Number of env steps (default 10)" + }, + "gripper": { + "type": [ + "number", + "string" + ], + "description": "Gripper while driving: +1 close, -1 open, or 'hold' (default 'hold')" + } + } + }, + "reset": { + "type": "object", + "properties": {} + }, + "view_env_state": { + "type": "object", + "properties": { + "step": { + "type": [ + "integer", + "null" + ], + "description": "Step number; 0 = initial. Null = latest." + } + } + }, + "view_camera_meta": { + "type": "object", + "properties": { + "camera": { + "type": "string", + "enum": [ + "agentview", + "navview" + ], + "description": "Camera metadata to read (default agentview)." + }, + "step": { + "type": [ + "integer", + "null" + ], + "description": "Step number (default latest)." + } + } + }, + "back_project": { + "type": "object", + "properties": { + "row": { + "type": "integer", + "description": "Pixel row (0=top) in the selected resolution image." + }, + "col": { + "type": "integer", + "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." + }, + "camera": { + "type": "string", + "enum": [ + "agentview", + "navview", + "wrist" + ], + "description": "Camera to back-project from (default agentview)." + }, + "resolution": { + "type": "string", + "enum": [ + "high", + "low" + ], + "description": "Coordinate system for row/col (default high). Use 'low' when row/col came from the standard 256x256 embedded image." + } + }, + "required": [ + "row", + "col" + ] + }, + "back_project_batch": { + "type": "object", + "properties": { + "pixels": { + "type": "array", + "description": "List of [row, col] pixel coordinates (max 50)", + "items": { + "type": "array", + "items": { + "type": "integer" + }, + "minItems": 2, + "maxItems": 2 + }, + "minItems": 1, + "maxItems": 50 + }, + "step": { + "type": [ + "integer", + "null" + ], + "description": "Depth / world-map step to use (default latest)." + }, + "camera": { + "type": "string", + "enum": [ + "agentview", + "navview", + "wrist" + ], + "description": "Camera to back-project from (default agentview)." + }, + "resolution": { + "type": "string", + "enum": [ + "high", + "low" + ], + "description": "Coordinate system for pixels (default low). Use 'low' for the standard 256x256 world map." + } + }, + "required": [ + "pixels" + ] + }, + "query_world_map": { + "type": "object", + "properties": { + "z_min": { + "type": "number", + "description": "Minimum Z in meters (default 0.85 for counter height)." + }, + "z_max": { + "type": "number", + "description": "Maximum Z in meters (default 0.95 for counter height)." + }, + "x_range": { + "type": [ + "array", + "null" + ], + "description": "Optional X range [min, max] in meters; null = no filter.", + "items": { + "type": "number" + }, + "minItems": 2, + "maxItems": 2 + }, + "y_range": { + "type": [ + "array", + "null" + ], + "description": "Optional Y range [min, max] in meters; null = no filter.", + "items": { + "type": "number" + }, + "minItems": 2, + "maxItems": 2 + }, + "camera": { + "type": "string", + "enum": [ + "agentview", + "navview", + "wrist" + ], + "description": "Camera world map to query (default agentview)." + }, + "resolution": { + "type": "string", + "enum": [ + "high", + "low" + ], + "description": "World map resolution (default low)." + }, + "min_cluster_size": { + "type": "integer", + "description": "Minimum pixels per cluster to report (default 10)." + } + } + }, + "finish": { + "type": "object", + "properties": { + "status": { + "type": "string", + "enum": [ + "success", + "failure", + "stuck" + ], + "description": "Task outcome classification." + }, + "summary": { + "type": "string", + "description": "1-3 sentence summary of what worked / what failed." + } + }, + "required": [ + "status", + "summary" + ] + } + } +} diff --git a/tests/unit_tests/robots/robocasa/test_memory_contracts.py b/tests/unit_tests/robots/robocasa/test_memory_contracts.py index f96d61e8a..37b2f4516 100644 --- a/tests/unit_tests/robots/robocasa/test_memory_contracts.py +++ b/tests/unit_tests/robots/robocasa/test_memory_contracts.py @@ -21,12 +21,13 @@ from pathlib import Path from types import SimpleNamespace -import pytest - from robots.robocasa.prompt_bundle import system_prompt from robots.robocasa.robot_spec import _parse_config +from robots.robocasa.tools import finish from rpent.memory import MemoryManager from rpent.prompt.utils import format_prompt +from rpent.session import EnvState +from rpent.tools import Toolkit def _args( @@ -90,13 +91,20 @@ def test_results_corpus_is_readable_through_memory_tool(monkeypatch, tmp_path): audit.write_text('{"success": true}\n') manager = MemoryManager(root=memory_root) - bindings = manager.get_common_tool_bindings() - read_text_file = bindings["read_text_file"][1] - write_text_file = bindings["write_text_file"][1] - - assert read_text_file(path=str(audit))["content"] == '{"success": true}\n' - with pytest.raises(PermissionError, match="writing to memory is denied"): - write_text_file(path=str(audit), content="{}\n") + toolkit = Toolkit( + state=EnvState(tmp_path / "run"), + memory=manager, + robot=None, + output_dir=tmp_path / "run", + tools=(finish,), + ) + read = toolkit.execute_tool("read_text_file", {"path": str(audit)}) + assert not read.is_error + assert read.data["content"] == '{"success": true}\n' + write = toolkit.execute_tool( + "write_text_file", {"path": str(audit), "content": "{}\n"} + ) + assert write.is_error and "writing to memory is denied" in write.error def test_parse_config_resolves_local_memory_dir(tmp_path): diff --git a/tests/unit_tests/robots/robocasa/test_robocasa_toolkit_contracts.py b/tests/unit_tests/robots/robocasa/test_robocasa_toolkit_contracts.py index 642ce45dd..5b2599ac4 100644 --- a/tests/unit_tests/robots/robocasa/test_robocasa_toolkit_contracts.py +++ b/tests/unit_tests/robots/robocasa/test_robocasa_toolkit_contracts.py @@ -12,137 +12,236 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Offline contracts for the RoboCasa toolkit.""" +"""Native RoboCasa execution, observation, and lifecycle contracts.""" from __future__ import annotations -from pathlib import Path -from types import SimpleNamespace -from typing import Any +import json +import threading +from concurrent.futures import ThreadPoolExecutor +import numpy as np import pytest -from robots.robocasa import robot_spec, toolkit -from rpent.dashboard.events import NullDashboardEventSink -from rpent.memory import MemoryManager -from rpent.robots import RunConfig -from rpent.tools.toolkit import Toolkit, _is_readonly -from rpent.utils import templates - -COMMON_TOOLS = {"read_text_file", "write_text_file", "list_dir", "finish"} - -EXPECTED_TOOLS = COMMON_TOOLS | { - "move_to", - "move_delta", - "rotate_pitch", - "set_gripper", - "release", - "scripted_grasp", - "rldx_skill", - "rldx_arm", - "navigate_to", - "move_base", - "reset", +from rpent.tools import ToolContext + +PERCEPTION = { "view_env_state", "back_project_batch", "query_world_map", } - - -def _record(step_idx: int = 0) -> SimpleNamespace: - return SimpleNamespace(step_idx=step_idx, terminated=False, extras={}) - - -def _tool_names(robot_toolkit: Toolkit) -> set[str]: - return {spec["name"] for spec in robot_toolkit.get_tools_spec()} - - -def _readonly_names(robot_toolkit: Toolkit) -> set[str]: - return { - name - for name, (_, handler) in robot_toolkit._tools.items() - if _is_readonly(handler) - } - - -def test_toolkit_falls_back_to_memory_root( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - memory_dir = tmp_path / "robocasa" - monkeypatch.setattr(robot_spec, "get_memory_dir", lambda _: memory_dir) - monkeypatch.setattr( - toolkit, - "RoboCasaToolkit", - lambda **kwargs: SimpleNamespace(**kwargs), - ) - config = RunConfig( - recipe_tag="cell-s0", - output_dir=tmp_path / "run", - prompt_vars={}, - task_desc={}, +COMMON = {"read_text_file", "write_text_file", "list_dir", "read_image", "finish"} + + +def test_toolkit_initializes_environment_and_tool_policies(make_toolkit): + run = make_toolkit() + tk = run.toolkit + assert run.env.reset_calls == 1 + assert not run.env.actions + assert ( + run.config.output_dir / "success_criteria.md" + ).read_text() == "offline success criteria" + for definition in tk.list_tools(): + if definition.name not in COMMON: + assert definition.readonly is (definition.name in PERCEPTION) + + +def test_observation_images_and_errors_are_native_and_saved(make_toolkit): + run = make_toolkit() + tk = run.toolkit + initial = tk.execute_tool("view_env_state", {"step": 0}) + assert initial.data["step"] == 0 + assert initial.data["task_language"] == run.env.language + names = [image["artifact"] for image in initial.data["images"]] + assert names == ["agentview_high.png", "navview.png", "wrist.png"] + assert initial.images == [tk.state.load_bytes(name, step=0) for name in names] + assert all(image.startswith(b"\x89PNG") for image in initial.images) + json.dumps(initial.to_dict()) + result = tk.execute_tool("move_to", {"xyz": [2, 0, 1], "max_steps": 1}) + assert result.is_error + assert "did not reach" in result.error + assert result.data["step"] == 1 + assert result.data["log"]["result"]["ok"] is False + assert "error" not in result.data["log"]["result"] + assert tk.state.get(1).result["error"] == result.error + assert result.images and result.data["artifacts"] + historical = tk.execute_tool("view_env_state", {"step": 1}) + assert historical.data["log"]["result"]["error"] == result.error + assert not historical.is_error + assert tk.state.latest_record().step_idx == 1 + + +@pytest.mark.parametrize( + ("name", "arguments", "expected_steps"), + [ + ("set_gripper", {"steps": 3}, 3), + ("release", {"steps": 4}, 4), + ("rotate_pitch", {"n": 5}, 5), + ("move_base", {"forward": 0.2, "steps": 3}, 3), + ("move_to", {"xyz": [2, 0, 1], "max_steps": 2}, 11), + ("move_delta", {"dxyz": [2, 0, 0], "max_steps": 2}, 11), + ("navigate_to", {"xy": [2, 0], "max_steps": 2}, 8), + ], +) +def test_actions_preserve_step_counts_and_invalidate_vla_state( + make_toolkit, name, arguments, expected_steps +): + run = make_toolkit() + result = run.toolkit.execute_tool(name, arguments) + assert result.data["step"] == 1 + assert len(run.env.actions) == expected_steps + assert run.toolkit._robot._vla_desync is True + + +def test_scripted_grasp_executes_stages_and_stops_on_failure(make_toolkit): + run = make_toolkit() + result = run.toolkit.execute_tool("scripted_grasp", {"xyz": [0, 0, 1]}) + assert not result.is_error + assert len(run.env.actions) > 18 + assert np.all([action[6] == -1 for action in run.env.actions[:4]]) + failed = make_toolkit() + failed.toolkit._robot._pos_jac = np.zeros((3, 3)) + result = failed.toolkit.execute_tool("scripted_grasp", {"xyz": [2, 0, 1]}) + assert result.is_error + assert result.data["log"]["result"]["stage"] == "approach" + assert len(failed.env.actions) == 204 + + +def test_reset_guard_preserves_evaluation_and_clears_exploration_state(make_toolkit): + run = make_toolkit() + denied = run.toolkit.execute_tool("reset", {}) + assert denied.is_error and "DISABLED" in denied.error + assert run.env.reset_calls == 1 + assert not run.env.actions + explore = make_toolkit(RLDX_ALLOW_RESET=1) + runtime = explore.toolkit._robot + runtime._pos_jac = np.eye(3) + runtime._fwd_offset = 1.0 + runtime._vla_desync = False + assert not explore.toolkit.execute_tool("reset", {}).is_error + assert explore.env.reset_calls == 2 + assert runtime._pos_jac is None and runtime._fwd_offset is None + assert runtime._vla_desync is True + assert explore.model.resets == 1 + + +def test_cancellation_stops_remaining_steps_and_uses_fresh_call_signal( + make_toolkit, monkeypatch +): + run = make_toolkit() + entered = threading.Event() + resume = threading.Event() + cancelled = threading.Event() + check = ToolContext.check_cancelled + + def checkpoint(ctx): + try: + check(ctx) + finally: + if ctx._cancel_event.is_set(): + cancelled.set() + + def on_step(): + entered.set() + assert resume.wait(5) + + monkeypatch.setattr(ToolContext, "check_cancelled", checkpoint) + run.env.on_step = on_step + with ThreadPoolExecutor(max_workers=2) as pool: + action = pool.submit(run.toolkit.execute_tool, "move_base", {"steps": 20}) + assert entered.wait(5) + # Request cancellation while the first step is in progress. + stop = pool.submit(run.toolkit.cancel_active_and_wait) + assert run.toolkit._active_operation.cancel_event.wait(5) + resume.set() + result = action.result(timeout=5) + stop.result(timeout=5) + assert cancelled.is_set() + assert result.is_error and "cancelled" in result.error + assert len(run.env.actions) == 1 + assert result.data["task_progress"]["steps"] == 1 + run.env.on_step = lambda: None + assert not run.toolkit.execute_tool("release", {"steps": 1}).is_error + assert len(run.env.actions) == 2 + + +def test_partial_failure_captures_current_state_and_retains_both_errors(make_toolkit): + run = make_toolkit() + + def fail_step(): + raise RuntimeError("actuator failed after moving") + + run.env.on_step = fail_step + result = run.toolkit.execute_tool("move_base", {"forward": 1, "steps": 5}) + assert result.is_error and "actuator failed" in result.error + assert result.data["state"]["robot0_base_pos"][0] == pytest.approx(0.01) + assert len(run.env.actions) == 1 + + def fail_render(*args, **kwargs): + raise RuntimeError("camera unavailable") + + run.env.render_camera = fail_render + result = run.toolkit.execute_tool("move_base", {"forward": 1, "steps": 5}) + assert "actuator failed" in result.error and "camera unavailable" in result.error + assert "state" not in result.data + + +def test_finish_keeps_environment_success_source(make_toolkit): + run = make_toolkit() + tk = run.toolkit + assert not tk.execute_tool("release", {"steps": 2}).is_error + finish = tk.execute_tool( + "finish", {"status": "success", "summary": "planner claim"} ) - - robot_toolkit = robot_spec.get_toolkit( - runtime_kwargs={}, - dashboard_events=NullDashboardEventSink(), - config=config, + assert not finish.is_error + assert tk.finish_result == {"status": "success", "summary": "planner claim"} + assert tk.solved() is False + run.env.success = True + tk.execute_tool("release", {"steps": 1}) + assert tk.solved() is True + run.env.success = False + tk.execute_tool("release", {"steps": 1}) + assert tk.solved() is False + + +@pytest.mark.parametrize( + ("name", "arguments"), + [ + ("move_to", {"xyz": [1, 2]}), + ("navigate_to", {"xy": [1, 2, 3]}), + ("back_project_batch", {"pixels": [[1, 2, 3]]}), + ("back_project_batch", {"pixels": [[1, 2]] * 51}), + ("back_project_batch", {"pixels": [[1, 2]], "camera": "typo"}), + ("query_world_map", {"x_range": [1]}), + ("rldx_skill", {"prompt": "atomic pick", "use_prompt": True}), + ("finish", {"status": "invalid", "summary": "bad"}), + ], +) +def test_model_argument_validation_precedes_execution(make_toolkit, name, arguments): + run = make_toolkit() + result = run.toolkit.execute_tool(name, arguments) + assert result.is_error and "Invalid arguments" in result.error + assert not run.env.actions + assert not run.model.calls + assert run.toolkit.state.latest_record().step_idx == 0 + + +def test_perception_uses_saved_maps_and_reports_missing_artifacts(make_toolkit): + run = make_toolkit(RLDX_KEEP_HEAVY_NPY=1) + tk = run.toolkit + result = tk.execute_tool( + "back_project_batch", {"pixels": [[0, 0], [1, 1], [-1, 0]]} ) - - assert robot_toolkit.memory.root == memory_dir.resolve() - - -def test_toolkit_constructs_and_classifies_tools_with_a_fake( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, - fake_single_arm_primitives: type[Any], -) -> None: - monkeypatch.delenv("RLDX_MAX_CHUNKS", raising=False) - monkeypatch.delenv("RLDX_SETTLE_PATIENCE", raising=False) - import robots.robocasa.primitives as primitives_module - - dumped: list[Any] = [] - monkeypatch.setattr( - templates, "default_variables", lambda: {"output_dir": "/offline/output"} - ) - monkeypatch.setattr( - primitives_module, - "RoboCasaPrimitives", - fake_single_arm_primitives, - ) - monkeypatch.setattr(toolkit, "get_output_dir", lambda: tmp_path) - monkeypatch.setattr( - toolkit.robocasa_tools, - "dump_state", - lambda primitives, state, log: dumped.append(primitives) or _record(), - ) - - robot_toolkit = toolkit.RoboCasaToolkit( - runtime_kwargs={"env_client": object(), "vla_client": object()}, - dashboard_events=NullDashboardEventSink(), - memory=MemoryManager(tmp_path / "memory"), - ) - - assert _tool_names(robot_toolkit) == EXPECTED_TOOLS - assert _readonly_names(robot_toolkit) == COMMON_TOOLS | { - "view_env_state", + assert not result.is_error + assert result.data["summary"]["valid_count"] == 2 + assert result.data["summary"]["median_xyz"] == [1.0, 1.0, 0.9] + assert not result.data["results"][2]["valid"] + assert tk.execute_tool("query_world_map", {"min_cluster_size": 1}).data["clusters"] + assert tk.execute_tool( "back_project_batch", - "query_world_map", - } - assert dumped == fake_single_arm_primitives.instances - primitive = fake_single_arm_primitives.instances[0] - assert primitive.reset_calls == 1 - assert primitive.recording_started is True - assert callable(primitive.kwargs["check_cancelled"]) - assert (tmp_path / "success_criteria.md").read_text() == "offline success criteria" - - -def test_solved_reads_only_the_final_environment_record() -> None: - records = [SimpleNamespace(extras={"success": True})] - robot_toolkit = toolkit.RoboCasaToolkit.__new__(toolkit.RoboCasaToolkit) - robot_toolkit._state = SimpleNamespace(latest_record=lambda: records[-1]) - - assert robot_toolkit.solved() is True - - records.append(SimpleNamespace(extras={"success": False})) - assert robot_toolkit.solved() is False + {"pixels": [[0, 0]], "camera": "wrist", "resolution": "high"}, + ).is_error + tk.execute_tool("release", {"steps": 1}) + pruned = tk.execute_tool("back_project_batch", {"pixels": [[0, 0]], "step": 0}) + assert pruned.is_error and "not found" in pruned.error + assert tk.execute_tool("view_env_state", {"step": 0}).images diff --git a/tests/unit_tests/robots/robocasa/test_vla_protocol_contracts.py b/tests/unit_tests/robots/robocasa/test_vla_protocol_contracts.py index 83af7bfda..6f763bacf 100644 --- a/tests/unit_tests/robots/robocasa/test_vla_protocol_contracts.py +++ b/tests/unit_tests/robots/robocasa/test_vla_protocol_contracts.py @@ -18,27 +18,22 @@ from pathlib import Path from types import SimpleNamespace -from typing import Any -from robots.robocasa.primitives import RoboCasaPrimitives +import numpy as np +import pytest + +from robots.robocasa import tools from robots.robocasa.prompt_bundle import system_prompt -from robots.robocasa.tools import TOOLS_SPEC from robots.robocasa.vla_server import _normalize_legacy_processor_geometry from rpent.prompt.utils import format_prompt +from rpent.tools import ToolContext -class _RecordingRldx: - def __init__(self) -> None: - self.calls: list[dict[str, Any]] = [] +def _record_vla(run): + calls = [] - def run( - self, - prompt: str, - max_chunks: int, - n_action_steps: int, - **kwargs: Any, - ) -> dict[str, Any]: - self.calls.append( + def record(prompt, max_chunks, n_action_steps, **kwargs): + calls.append( { "prompt": prompt, "max_chunks": max_chunks, @@ -48,19 +43,8 @@ def run( ) return {"ok": True, "prompt": prompt, "status": "cap"} - -def _fake_primitives(task_language: str) -> tuple[RoboCasaPrimitives, _RecordingRldx]: - rldx = _RecordingRldx() - primitives = RoboCasaPrimitives.__new__(RoboCasaPrimitives) - primitives.env = SimpleNamespace( - current_raw_obs={"language": task_language}, - get_task_language=lambda: task_language, - ) - primitives._rldx = rldx - primitives._vla_desync = True - primitives._recording = False - primitives.record_frame = lambda: None - return primitives, rldx + run.toolkit._robot._rldx.run = record + return calls def test_prompt_requires_live_task_language_and_fresh_geometry(tmp_path: Path) -> None: @@ -83,136 +67,181 @@ def test_prompt_requires_live_task_language_and_fresh_geometry(tmp_path: Path) - assert "{{" not in rendered -def test_vla_tool_schema_hides_historical_prompt_override() -> None: - vla_specs = { - spec["name"]: spec for spec in TOOLS_SPEC if spec["name"].startswith("rldx_") - } - - assert set(vla_specs) == {"rldx_skill", "rldx_arm"} - for spec in vla_specs.values(): - schema = spec["input_schema"] +def test_vla_schema_exposes_budgets_and_hides_legacy_prompt_override(): + for t in (tools.rldx_skill, tools.rldx_arm): + schema = t.input_schema assert schema["required"] == ["prompt"] - assert "use_prompt" not in schema["properties"] - assert "complete live task_language" in spec["description"] - + assert set(schema["properties"]) == { + "prompt", + "base_clip", + "max_chunks", + "force_reset", + "n_action_steps", + "settle_patience", + "settle_eps", + } + assert "complete live task_language" in t.description + + +def test_vla_uses_live_language_and_run_defaults(make_toolkit): + run = make_toolkit() + calls = _record_vla(run) + first = run.toolkit.execute_tool("rldx_skill", {"prompt": "atomic pick"}) + second = run.toolkit.execute_tool("rldx_arm", {"prompt": "atomic place"}) + assert not first.is_error and not second.is_error + assert [call["prompt"] for call in calls] == [run.env.language] * 2 + assert [call["force_reset"] for call in calls] == [True, False] + assert [call["base_clip"] for call in calls] == [None, 0.1] + for result in (first, second): + data = result.data["log"]["result"] + assert data["effective_max_chunks"] == 70 + assert data["effective_n_action_steps"] == 8 + assert data["effective_settle_patience"] == 999 + assert data["effective_prompt"] == run.env.language + assert data["prompt_overridden"] is True + assert first.data["log"]["result"]["requested_prompt"] == "atomic pick" + assert run.toolkit._robot._vla_desync is False + for name in ("rldx_skill", "rldx_arm"): + result = run.toolkit.execute_tool( + name, + { + "prompt": run.env.language, + "max_chunks": 3, + "n_action_steps": 4, + "settle_patience": 5, + }, + ) + assert not result.is_error + data = result.data["log"]["result"] + assert data["effective_max_chunks"] == calls[-1]["max_chunks"] == 3 + assert data["effective_n_action_steps"] == calls[-1]["n_action_steps"] == 4 + assert data["effective_settle_patience"] == calls[-1]["settle_patience"] == 5 -def test_vla_always_uses_live_task_language_and_preserves_continuity() -> None: - task_language = "Pick the squash up and place it in the microwave." - primitives, rldx = _fake_primitives(task_language) - first = primitives.rldx_skill( - prompt="Pick the squash up.", - use_prompt=True, - max_chunks=3, +def test_environment_overrides_call_budgets_at_execution(make_toolkit, monkeypatch): + run = make_toolkit( + RLDX_MAX_CHUNKS=40, RLDX_ACTION_STEPS_PER_CHUNK=6, RLDX_SETTLE_PATIENCE=99 ) - second = primitives.rldx_arm( - prompt="Place it in the microwave.", - use_prompt=True, - max_chunks=4, + calls = _record_vla(run) + for name in ("rldx_skill", "rldx_arm"): + result = run.toolkit.execute_tool( + name, + { + "prompt": run.env.language, + "max_chunks": 2, + "n_action_steps": 1, + "settle_patience": 2, + }, + ) + assert not result.is_error + data = result.data["log"]["result"] + assert data["effective_max_chunks"] == 40 + assert data["effective_n_action_steps"] == 6 + assert data["effective_settle_patience"] == 99 + assert data["prompt_overridden"] is False + assert "requested_prompt" not in data + assert result.data["log"]["command"]["max_chunks"] == 2 + assert all( + (call["max_chunks"], call["n_action_steps"], call["settle_patience"]) + == (40, 6, 99) + for call in calls ) - - assert [call["prompt"] for call in rldx.calls] == [task_language, task_language] - assert rldx.calls[0]["force_reset"] is True - assert rldx.calls[1]["force_reset"] is False - assert primitives._vla_desync is False - assert first["effective_prompt"] == task_language - assert first["effective_max_chunks"] == 3 - assert first["prompt_overridden"] is True - assert first["requested_prompt"] == "Pick the squash up." - assert second["effective_prompt"] == task_language - assert second["effective_max_chunks"] == 4 - assert second["prompt_overridden"] is True - - -def test_environment_max_chunks_locks_the_formal_protocol(monkeypatch) -> None: - task_language = "Open the left drawer." - primitives, rldx = _fake_primitives(task_language) - monkeypatch.setenv("RLDX_MAX_CHUNKS", "40") - - result = primitives.rldx_skill(prompt=task_language, max_chunks=70) - - assert rldx.calls[0]["max_chunks"] == 40 - assert result["effective_max_chunks"] == 40 - - -def test_ordinary_robocasa_keeps_default_max_chunks_at_70(monkeypatch) -> None: - task_language = "Open the left drawer." - primitives, rldx = _fake_primitives(task_language) - monkeypatch.delenv("RLDX_MAX_CHUNKS", raising=False) - - result = primitives.rldx_skill(prompt=task_language) - - assert rldx.calls[0]["max_chunks"] == 70 - assert result["effective_max_chunks"] == 70 - - -def test_environment_locks_all_formal_rldx_runtime_values(monkeypatch) -> None: - task_language = "Open the left drawer." - primitives, rldx = _fake_primitives(task_language) - monkeypatch.setenv("RLDX_MAX_CHUNKS", "40") - monkeypatch.setenv("RLDX_ACTION_STEPS_PER_CHUNK", "8") - monkeypatch.setenv("RLDX_SETTLE_PATIENCE", "999") - - result = primitives.rldx_skill( - prompt=task_language, - max_chunks=2, - n_action_steps=1, - settle_patience=2, + monkeypatch.setenv("RLDX_MAX_CHUNKS", "5") + result = run.toolkit.execute_tool( + "rldx_skill", {"prompt": run.env.language, "max_chunks": 2} ) - - assert rldx.calls[0]["max_chunks"] == 40 - assert rldx.calls[0]["n_action_steps"] == 8 - assert rldx.calls[0]["settle_patience"] == 999 - assert result["effective_max_chunks"] == 40 - assert result["effective_n_action_steps"] == 8 - assert result["effective_settle_patience"] == 999 - - -def test_invalid_vla_budgets_return_errors_without_executing_rldx( - monkeypatch, -) -> None: - monkeypatch.delenv("RLDX_MAX_CHUNKS", raising=False) - monkeypatch.delenv("RLDX_ACTION_STEPS_PER_CHUNK", raising=False) - monkeypatch.delenv("RLDX_SETTLE_PATIENCE", raising=False) - - for arguments, parameter in ( - ({"max_chunks": 0}, "max_chunks"), - ({"n_action_steps": 0}, "n_action_steps"), - ({"settle_patience": 0}, "settle_patience"), - ): - primitives, rldx = _fake_primitives("Open the left drawer.") - - result = primitives.rldx_skill( - prompt="Open the left drawer.", - **arguments, - ) - - assert result == { - "error": f"{parameter} must be positive; VLA was not executed" - } - assert rldx.calls == [] - - -def test_matching_vla_prompt_is_reported_without_override() -> None: - task_language = "Open the left drawer." - primitives, rldx = _fake_primitives(task_language) - - result = primitives.rldx_skill(prompt=task_language, use_prompt=False) - - assert rldx.calls[0]["prompt"] == task_language - assert result["effective_prompt"] == task_language - assert result["prompt_overridden"] is False - assert "requested_prompt" not in result - - -def test_vla_does_not_run_without_environment_task_language() -> None: - primitives, rldx = _fake_primitives("") - - result = primitives.rldx_skill(prompt="atomic fallback", use_prompt=True) - - assert "task language is unavailable" in result["error"] - assert rldx.calls == [] - assert primitives._vla_desync is True + assert not result.is_error + assert result.data["log"]["result"]["effective_max_chunks"] == 5 + assert calls[-1]["max_chunks"] == 5 + + +@pytest.mark.parametrize( + ("variable", "parameter"), + [ + ("RLDX_MAX_CHUNKS", "max_chunks"), + ("RLDX_ACTION_STEPS_PER_CHUNK", "n_action_steps"), + ("RLDX_SETTLE_PATIENCE", "settle_patience"), + ], +) +@pytest.mark.parametrize("source", ["environment", "arguments"]) +def test_invalid_effective_budgets_fail_without_vla_execution( + make_toolkit, variable, parameter, source +): + run = make_toolkit(**({variable: 0} if source == "environment" else {})) + calls = _record_vla(run) + arguments = {"prompt": run.env.language} + if source == "arguments": + arguments[parameter] = 0 + result = run.toolkit.execute_tool("rldx_skill", arguments) + assert result.error == f"{parameter} must be positive; VLA was not executed" + assert not calls + assert not run.env.actions + assert run.toolkit._robot._vla_desync is True + + +def test_vla_rejects_missing_live_language(make_toolkit): + run = make_toolkit() + run.env.language = "" + calls = _record_vla(run) + result = run.toolkit.execute_tool("rldx_skill", {"prompt": "atomic fallback"}) + assert result.is_error and "task language is unavailable" in result.error + assert result.data["log"]["result"]["effective_prompt"] == "" + assert not calls + assert run.toolkit._robot._vla_desync is True + + +def test_rollout_preserves_or_reseeds_history(make_toolkit): + run = make_toolkit(RLDX_MAX_CHUNKS=1) + tk = run.toolkit + for name in ("rldx_skill", "rldx_arm"): + result = tk.execute_tool(name, {"prompt": run.env.language}) + assert not result.is_error + assert result.data["log"]["result"]["steps_applied"] == 2 + # Hitting the chunk cap with no task success is an ordinary VLA result. + assert result.data["log"]["result"]["status"] == "cap" + assert len(run.env.actions) == 4 + np.testing.assert_allclose(run.env.actions[0][7:11], 1) + np.testing.assert_allclose(run.env.actions[2][7:11], 0.1) + assert [options["reset_memory"] for _, options in run.model.calls] == [ + [True], + [False], + ] + before_manual = run.model.calls[1][0]["video.robot0_agentview_left"] + assert before_manual.shape == (1, 2, 4, 4, 3) + assert before_manual[0, :, 0, 0, 0].tolist() == [0, 2] + tk.execute_tool("release", {"steps": 1}) + assert tk._robot._vla_desync is True + tk.execute_tool("rldx_arm", {"prompt": run.env.language}) + assert run.model.calls[2][1]["reset_memory"] == [True] + assert run.model.calls[2][0]["video.robot0_agentview_left"][ + 0, :, 0, 0, 0 + ].tolist() == [5, 5] + tk.execute_tool("rldx_arm", {"prompt": run.env.language, "force_reset": True}) + assert run.model.calls[3][1]["reset_memory"] == [True] + assert len(run.env.actions) == 9 + + +def test_vla_checks_current_call_cancellation_after_inference( + make_toolkit, monkeypatch +): + run = make_toolkit(RLDX_MAX_CHUNKS=1) + contexts = [] + original = ToolContext.check_cancelled + + def checkpoint(ctx): + contexts.append(ctx) + original(ctx) + + monkeypatch.setattr(ToolContext, "check_cancelled", checkpoint) + run.model.on_predict = lambda: contexts[-1]._cancel_event.set() + result = run.toolkit.execute_tool("rldx_skill", {"prompt": run.env.language}) + assert result.is_error and "cancelled" in result.error + assert len(run.model.calls) == 1 + assert not run.env.actions + run.model.on_predict = lambda: None + result = run.toolkit.execute_tool("rldx_skill", {"prompt": run.env.language}) + assert not result.is_error + assert len(run.env.actions) == 2 def test_legacy_rldx_processor_null_geometry_uses_release_defaults(monkeypatch) -> None: diff --git a/tests/unit_tests/robots/robotwin/conftest.py b/tests/unit_tests/robots/robotwin/conftest.py new file mode 100644 index 000000000..217924abb --- /dev/null +++ b/tests/unit_tests/robots/robotwin/conftest.py @@ -0,0 +1,158 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Offline RoboTwin clients for testing the real native tools.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest + +from robots.robotwin.robot_spec import ROBOTWIN_CAMERA_NAMES +from robots.robotwin.toolkit import RoboTwinToolkit +from rpent.dashboard.events import NullDashboardEventSink +from rpent.memory import MemoryManager + + +class FakeRoboTwinEnv: + def __init__(self): + self.terminated = False + self.truncated = False + self.server_meta = {"task_name": "stack_blocks"} + self.execution_capabilities = {"chunk_step_all_frames": True} + self.last_info = { + "episode_status": { + "eval_success": False, + "take_action_cnt": 0, + "step_lim": 200, + "actual_seed": 7, + }, + "robot_state": { + "qpos_target14": [0.0] * 14, + "left_eef_pose": [0.1, 0.2, 0.3, 1.0, 0.0, 0.0, 0.0], + "right_eef_pose": [0.4, 0.5, 0.6, 1.0, 0.0, 0.0, 0.0], + "left_gripper": 0.0, + "right_gripper": 0.0, + }, + } + self.last_reset_info = { + "actual_seed": 7, + "instruction": self.get_task_language(), + } + self.renders = [] + self.plans = [] + self.steps = [] + self.chunks = [] + self.path = np.arange(36, dtype=float).reshape(6, 6) / 100 + self.plan_status = "Success" + self.on_step = lambda: None + self.on_chunk = lambda: None + self.success_at = None + + def get_task_language(self): + return "Stack the red block on the blue block." + + def render_camera(self, camera, *, depth=False): + self.renders.append((camera, depth)) + rgb = np.full( + (4, 5, 3), ROBOTWIN_CAMERA_NAMES.index(camera) * 50, dtype=np.uint8 + ) + return (rgb, np.ones((4, 5))) if depth else rgb + + def get_camera_meta(self, camera): + return { + "intrinsic_K": [[2.0, 0, 0], [0, 2.0, 0], [0, 0, 1]], + "cam2world_gl": np.eye(4).tolist(), + "height": 4, + "width": 5, + } + + def plan_arm_path(self, arm, target_pose): + self.plans.append((arm, target_pose.copy())) + return {"status": self.plan_status, "position": self.path.copy()} + + def _advance(self, requested): + status = self.last_info["episode_status"] + count = min(requested, status["step_lim"] - status["take_action_cnt"]) + if self.success_at is not None: + count = min(count, self.success_at - status["take_action_cnt"]) + status["take_action_cnt"] += count + status["eval_success"] = status["take_action_cnt"] == self.success_at + self.terminated = status["eval_success"] + self.truncated = status["take_action_cnt"] >= status["step_lim"] + self.last_info["executed_actions"] = count + return count + + def step(self, action, *, action_type): + assert not self.terminated and not self.truncated + assert action_type == "qpos" + self.steps.append(action.copy()) + self._advance(1) + state = self.last_info["robot_state"] + state["qpos_target14"] = action.tolist() + state["left_gripper"] = action[6] + state["right_gripper"] = action[13] + self.on_step() + return ( + {"main_images": np.zeros((4, 5, 3), dtype=np.uint8)}, + 0, + self.terminated, + self.truncated, + self.last_info, + ) + + def chunk_step(self, actions, *, action_type, return_all_frames): + assert not self.terminated and not self.truncated + assert action_type == "ee" + self.chunks.append((actions.copy(), return_all_frames)) + executed = self._advance(len(actions)) + rgb = np.zeros((4, 5, 3), dtype=np.uint8) + payload = {"main_images": rgb} + if return_all_frames: + payload = { + "frames": [rgb.copy() for _ in range(executed)], + "final": payload, + } + self.on_chunk() + return payload, 0, self.terminated, self.truncated, self.last_info + + +class FakeLingBot: + def __init__(self): + self.observations = [] + self.on_infer = lambda: None + + def infer(self, observation): + self.observations.append(observation) + self.on_infer() + return np.arange(60 * 16, dtype=float).reshape(60, 16) + + +@pytest.fixture +def robotwin(tmp_path): + env = FakeRoboTwinEnv() + model = FakeLingBot() + toolkit = RoboTwinToolkit( + runtime_kwargs={"env": env, "model": model, "seed": 7}, + output_dir=tmp_path / "run", + dashboard_events=NullDashboardEventSink(), + memory=MemoryManager(tmp_path / "memory"), + ) + yield SimpleNamespace( + env=env, model=model, toolkit=toolkit, output_dir=tmp_path / "run" + ) + toolkit._frames.clear() + toolkit.close() diff --git a/tests/unit_tests/robots/robotwin/fixtures/pre_native_tool_contracts.json b/tests/unit_tests/robots/robotwin/fixtures/pre_native_tool_contracts.json new file mode 100644 index 000000000..bb4171c34 --- /dev/null +++ b/tests/unit_tests/robots/robotwin/fixtures/pre_native_tool_contracts.json @@ -0,0 +1,252 @@ +{ + "source_commit": "014a0fa97f69c991ee5e0f62f14e1d6f89c3dcd7", + "schemas": { + "view_env_state": { + "type": "object", + "properties": { + "step": { + "type": "integer", + "default": -1, + "description": "Step number; 0 = initial, -1 = latest." + } + } + }, + "render": { + "type": "object", + "properties": {} + }, + "sample_world_xyz": { + "type": "object", + "properties": { + "view": { + "type": "string", + "description": "Artifact view and pixel coordinate space. It must match the RGB image used to choose pixels." + }, + "pixels": { + "type": "array", + "minItems": 1, + "maxItems": 256, + "items": { + "type": "array", + "items": { + "type": "integer" + }, + "minItems": 2, + "maxItems": 2 + } + }, + "step": { + "type": [ + "integer", + "null" + ] + }, + "neighborhood": { + "type": "integer", + "minimum": 0, + "maximum": 32, + "default": 1 + } + }, + "required": [ + "view", + "pixels" + ] + }, + "query_world_map": { + "type": "object", + "properties": { + "view": { + "type": "string", + "description": "Artifact view and bbox coordinate space. It must match the RGB image used to choose the bbox." + }, + "bbox": { + "type": "array", + "items": { + "type": "integer" + }, + "minItems": 4, + "maxItems": 4 + }, + "step": { + "type": [ + "integer", + "null" + ] + }, + "max_points": { + "type": "integer", + "minimum": 1, + "maximum": 4096, + "default": 256 + } + }, + "required": [ + "view", + "bbox" + ] + }, + "lingbot_act": { + "type": "object", + "properties": { + "chunks": { + "type": "integer", + "minimum": 1, + "default": 4 + }, + "use_length": { + "type": "integer", + "const": 50, + "default": 50 + }, + "prompt": { + "type": [ + "string", + "null" + ] + } + } + }, + "move_to": { + "type": "object", + "properties": { + "arm": { + "type": "string", + "enum": [ + "left", + "right" + ] + }, + "xyz": { + "type": "array", + "items": { + "type": "number" + }, + "minItems": 3, + "maxItems": 3 + }, + "quat": { + "type": [ + "array", + "null" + ], + "items": { + "type": "number" + }, + "minItems": 4, + "maxItems": 4 + }, + "gripper": { + "type": [ + "number", + "null" + ] + }, + "substeps": { + "type": "integer", + "minimum": 0, + "default": 25 + } + }, + "required": [ + "arm", + "xyz" + ] + }, + "rotate_wrist": { + "type": "object", + "properties": { + "arm": { + "type": "string", + "enum": [ + "left", + "right" + ] + }, + "delta_yaw_deg": { + "type": "number" + }, + "gripper": { + "type": [ + "number", + "null" + ] + }, + "substeps": { + "type": "integer", + "minimum": 0, + "default": 25 + } + }, + "required": [ + "arm", + "delta_yaw_deg" + ] + }, + "set_gripper": { + "type": "object", + "properties": { + "arm": { + "type": "string", + "enum": [ + "left", + "right" + ] + }, + "val": { + "type": "number", + "minimum": 0, + "maximum": 1 + }, + "steps": { + "type": "integer", + "minimum": 1, + "default": 10 + } + }, + "required": [ + "arm", + "val" + ] + }, + "release": { + "type": "object", + "properties": { + "arm": { + "type": "string", + "enum": [ + "left", + "right" + ] + }, + "val": { + "type": "number", + "default": 1.0 + }, + "steps": { + "type": "integer", + "minimum": 1, + "default": 10 + } + }, + "required": [ + "arm" + ] + }, + "finish": { + "type": "object", + "properties": { + "status": { + "type": "string" + }, + "summary": { + "type": "string" + } + }, + "required": [ + "status", + "summary" + ] + } + } +} diff --git a/tests/unit_tests/robots/robotwin/test_robotwin_tool_schema_contracts.py b/tests/unit_tests/robots/robotwin/test_robotwin_tool_schema_contracts.py index 9e30a41a3..e768368fd 100644 --- a/tests/unit_tests/robots/robotwin/test_robotwin_tool_schema_contracts.py +++ b/tests/unit_tests/robots/robotwin/test_robotwin_tool_schema_contracts.py @@ -14,16 +14,51 @@ """Robot-specific schema contracts for RoboTwin tools.""" +import pytest + from robots.robotwin import tools def test_perception_schemas_use_the_same_view_coordinate_space() -> None: - by_name = {spec["name"]: spec for spec in tools.TOOLS_SPEC} + by_name = {t.name: t for t in tools.ROBOTWIN_TOOLS} for tool_name, coordinate_name in ( ("sample_world_xyz", "pixels"), ("query_world_map", "bbox"), ): - schema = by_name[tool_name]["input_schema"] + schema = by_name[tool_name].input_schema assert schema["required"] == ["view", coordinate_name] assert schema["properties"]["view"]["type"] == "string" + + +def test_lingbot_keeps_fixed_chunk_length(): + use_length = tools.lingbot_act.input_schema["properties"]["use_length"] + assert use_length["const"] == use_length["default"] == 50 + + +@pytest.mark.parametrize( + "name,arguments", + [ + ("lingbot_act", {"chunks": 0}), + ("lingbot_act", {"use_length": 10}), + ("move_to", {"arm": "both", "xyz": [1, 2, 3]}), + ("move_to", {"arm": "left", "xyz": [1, 2]}), + ("move_to", {"arm": "left", "xyz": [1, 2, 3], "quat": [1, 0, 0]}), + ("move_to", {"arm": "left", "xyz": [1, 2, 3], "substeps": -1}), + ("set_gripper", {"arm": "left", "val": 1.1}), + ("release", {"arm": "right", "steps": 0}), + ("sample_world_xyz", {"view": "head", "pixels": []}), + ("sample_world_xyz", {"view": "head", "pixels": [[0]]}), + ("sample_world_xyz", {"view": "head", "pixels": [[0, 0]], "neighborhood": 33}), + ("query_world_map", {"view": "head", "bbox": [0, 0, 1]}), + ("query_world_map", {"view": "head", "bbox": [0, 0, 1, 1], "max_points": 4097}), + ], +) +def test_invalid_arguments_are_rejected_before_execution_or_capture( + robotwin, name, arguments +): + result = robotwin.toolkit.execute_tool(name, arguments) + assert result.is_error and result.error.startswith("Invalid arguments") + assert not robotwin.env.steps and not robotwin.env.chunks and not robotwin.env.plans + assert robotwin.toolkit.state.latest_step == 0 + assert len(robotwin.env.renders) == 3 diff --git a/tests/unit_tests/robots/robotwin/test_robotwin_toolkit_contracts.py b/tests/unit_tests/robots/robotwin/test_robotwin_toolkit_contracts.py index ff3bcc3ac..cc2aae57b 100644 --- a/tests/unit_tests/robots/robotwin/test_robotwin_toolkit_contracts.py +++ b/tests/unit_tests/robots/robotwin/test_robotwin_toolkit_contracts.py @@ -12,192 +12,403 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Offline contracts for the RoboTwin toolkit.""" +"""Native RoboTwin action, observation, cancellation, and lifecycle contracts.""" from __future__ import annotations -from pathlib import Path +import argparse +import threading from types import SimpleNamespace -from typing import Any +import numpy as np import pytest -from robots.robotwin import toolkit -from robots.robotwin.primitives import RoboTwinPrimitives +from robots.robotwin import robot_spec, tools +from robots.robotwin.env_client import RoboTwinEnvClient +from robots.robotwin.robot_spec import ROBOTWIN_CAMERA_NAMES +from robots.robotwin.vla_client import LingBotVLAClient from rpent.dashboard.events import NullDashboardEventSink -from rpent.memory import MemoryManager -from rpent.tools.toolkit import Toolkit, _is_readonly, readonly -from rpent.utils import templates - -COMMON_TOOLS = {"read_text_file", "write_text_file", "list_dir", "finish"} - -EXPECTED_TOOLS = COMMON_TOOLS | { - "view_env_state", - "render", - "sample_world_xyz", - "query_world_map", - "lingbot_act", - "move_to", - "rotate_wrist", - "set_gripper", - "release", -} - -PRIMITIVE_METHODS = { - "start_recording", - "recorded_frame_count", - "frame_slice", - "stop_recording", - "status", - "finish", - "lingbot_act", - "move_to", - "rotate_wrist", - "set_gripper", - "release", -} - - -class FakeRoboTwinPrimitives: - instances: list[FakeRoboTwinPrimitives] = [] - - def __init__(self, **kwargs: Any) -> None: - self.kwargs = kwargs - self.status_calls = 0 - self.recording_started = False - self.env = SimpleNamespace(last_reset_info={"actual_seed": 7}) - type(self).instances.append(self) - - def start_recording(self) -> None: - self.recording_started = True - - def recorded_frame_count(self) -> int: - return 0 - - def frame_slice(self, start: int) -> list[Any]: - del start - return [] - - def stop_recording(self) -> list[Any]: - return [] - - def status(self) -> dict[str, Any]: - self.status_calls += 1 - return { - "eval_success": False, - "take_action_cnt": 0, - "step_lim": 100, - "actual_seed": 7, - } - - def finish(self, *, status: str, summary: str) -> dict[str, Any]: - return {"_finish": True, "status": status, "summary": summary} - - @staticmethod - def _operation(name: str, **kwargs: Any) -> dict[str, Any]: - return {"operation": name, "arguments": kwargs} - - def lingbot_act(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("lingbot_act", **kwargs) - - def move_to(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("move_to", **kwargs) - - def rotate_wrist(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("rotate_wrist", **kwargs) - - def set_gripper(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("set_gripper", **kwargs) - - def release(self, **kwargs: Any) -> dict[str, Any]: - return self._operation("release", **kwargs) - - -def _record(step_idx: int = 0) -> SimpleNamespace: - return SimpleNamespace(step_idx=step_idx, terminated=False) - - -def _tool_names(robot_toolkit: Toolkit) -> set[str]: - return {spec["name"] for spec in robot_toolkit.get_tools_spec()} - - -def _readonly_names(robot_toolkit: Toolkit) -> set[str]: - return { - name - for name, (_, handler) in robot_toolkit._tools.items() - if _is_readonly(handler) +from rpent.tools import ToolCancelled, ToolContext +from rpent.tools.common_tools import COMMON_TOOLS + + +def test_initial_observation_and_readonly_tools(robotwin): + tk = robotwin.toolkit + assert {t.name for t in tk.list_tools()} == { + t.name for t in (*COMMON_TOOLS, *tools.ROBOTWIN_TOOLS) + } + assert { + t.name for t in tools.ROBOTWIN_TOOLS if t.readonly and t.name != "finish" + } == { + "view_env_state", + "sample_world_xyz", + "query_world_map", } + assert robotwin.env.renders == [(camera, True) for camera in ROBOTWIN_CAMERA_NAMES] + initial = tk.state.get(0) + assert initial.command == {"action": "reset"} + assert initial.result == {**robotwin.env.last_reset_info, "success": True} + assert len(initial.artifacts) == 12 + result = tk.execute_tool("view_env_state", {}) + assert not result.is_error + assert len(result.images) == 3 + assert result.images == [ + tk.state.load_bytes(f"{c}_rgb.png", step=0) for c in ROBOTWIN_CAMERA_NAMES + ] + assert "_image_bytes" not in result.data + assert "log" not in result.data["state"] + assert tk.state.latest_step == 0 + assert len(robotwin.env.renders) == 3 + assert not tk.solved() + assert tk.execute_tool("view_env_state", {"step": 99}).is_error + + +@pytest.mark.parametrize( + "substeps,indices", + [(0, [0, 1, 2, 3, 4, 5]), (1, [5]), (3, [0, 2, 5]), (25, [0, 1, 2, 3, 4, 5])], +) +def test_move_preserves_path_sampling_and_fresh_other_arm_state( + robotwin, substeps, indices +): + env = robotwin.env + + def change_other_arm(): + env.last_info["robot_state"]["qpos_target14"][7] += 1 + + env.on_step = change_other_arm + result = robotwin.toolkit.execute_tool( + "move_to", + { + "arm": "left", + "xyz": [0.2, 0.3, 0.4], + "gripper": 0.7, + "substeps": substeps, + }, + ) + assert not result.is_error + np.testing.assert_allclose(env.plans[0][1], [0.2, 0.3, 0.4, 1, 0, 0, 0]) + np.testing.assert_allclose(np.asarray(env.steps)[:, :6], env.path[indices]) + np.testing.assert_allclose(np.asarray(env.steps)[:, 6], 0.7) + np.testing.assert_allclose(np.asarray(env.steps)[:, 7], np.arange(len(indices))) + assert result.data["log"]["result"]["executed_steps"] == len(indices) + assert result.data["state"]["episode_status"]["native_actions"] == len(indices) + assert result.data["state"]["episode_status"]["policy_actions"] == 0 + assert robotwin.toolkit.state.latest_step == 1 + assert len(env.renders) == 6 + + +def test_rotate_uses_world_z_and_only_captures_once(robotwin): + result = robotwin.toolkit.execute_tool( + "rotate_wrist", + { + "arm": "right", + "delta_yaw_deg": 90, + "substeps": 1, + }, + ) + assert not result.is_error + arm, target = robotwin.env.plans[0] + assert arm == "right" + np.testing.assert_allclose( + target, [0.4, 0.5, 0.6, np.sqrt(0.5), 0, 0, np.sqrt(0.5)] + ) + assert result.data["log"]["result"]["requested_delta_yaw_deg"] == 90 + assert robotwin.toolkit.state.latest_step == 1 + assert len(robotwin.env.renders) == 6 + + +@pytest.mark.parametrize( + "name,args,target", + [ + ("set_gripper", {"val": 0.8}, 0.8), + ("release", {}, 1.0), + ], +) +def test_gripper_interpolation_and_release_composition(robotwin, name, args, target): + result = robotwin.toolkit.execute_tool(name, {"arm": "right", "steps": 4, **args}) + assert not result.is_error + np.testing.assert_allclose( + np.asarray(robotwin.env.steps)[:, 13], np.arange(1, 5) * target / 4 + ) + np.testing.assert_allclose(np.asarray(robotwin.env.steps)[:, :13], 0) + assert result.data["log"]["result"]["gripper_val"] == target + assert robotwin.toolkit.state.latest_step == 1 + + +@pytest.mark.parametrize("empty_path", [False, True]) +def test_plan_failure_reports_native_error_and_captures(robotwin, empty_path): + if empty_path: + robotwin.env.path = np.empty((0, 6)) + else: + robotwin.env.plan_status = "Failure" + result = robotwin.toolkit.execute_tool("move_to", {"arm": "left", "xyz": [1, 2, 3]}) + assert result.is_error + assert result.data["log"]["result"]["stop_reason"] == "plan_failed" + assert "error" not in result.data["log"]["result"] + assert "error" in robotwin.toolkit.state.get(1).result + assert len(result.images) == 3 + assert not robotwin.env.steps + + +def test_partial_execution_failure_keeps_counts_and_observation(robotwin): + original = robotwin.env.step + + def step(action, **kwargs): + if len(robotwin.env.steps) == 2: + raise RuntimeError("offline step failure") + return original(action, **kwargs) + + robotwin.env.step = step + result = robotwin.toolkit.execute_tool("set_gripper", {"arm": "left", "val": 1}) + assert result.is_error and "offline step failure" in result.error + assert len(robotwin.env.steps) == 2 + assert result.data["state"]["episode_status"]["native_actions"] == 2 + assert robotwin.toolkit.state.latest_step == 1 + + +def _context(robotwin): + cancel = threading.Event() + tk = robotwin.toolkit + return ToolContext( + state=tk.state, + memory=tk.memory, + robot=tk._robot, + output_dir=robotwin.output_dir, + record_frame=tk.record_frame, + _cancel_event=cancel, + ), cancel + + +def test_cancel_stops_before_next_waypoint_and_keeps_completed_count(robotwin): + ctx, cancel = _context(robotwin) + robotwin.env.on_step = cancel.set + with pytest.raises(ToolCancelled): + tools.move_to.handler(arm="left", xyz=[0.1, 0.2, 0.3], ctx=ctx) + assert len(robotwin.env.steps) == 1 + assert ctx.robot.native_actions == 1 + + +@pytest.mark.parametrize("cancel_at,executed", [("infer", 0), ("chunk", 50)]) +def test_lingbot_cancellation_boundaries_preserve_counters( + robotwin, cancel_at, executed +): + ctx, cancel = _context(robotwin) + if cancel_at == "infer": + robotwin.model.on_infer = cancel.set + else: + robotwin.env.on_chunk = cancel.set + with pytest.raises(ToolCancelled): + tools.lingbot_act.handler(chunks=2, ctx=ctx) + assert ctx.robot.policy_actions == ctx.robot.native_actions == executed + assert len(robotwin.env.chunks) == (1 if executed else 0) + + +def test_lingbot_native_instruction_and_rgb_only_inference(robotwin): + result = robotwin.toolkit.execute_tool( + "lingbot_act", {"chunks": 2, "prompt": "ignored instruction"} + ) + assert not result.is_error + assert len(robotwin.model.observations) == 2 + for observation in robotwin.model.observations: + assert observation["task_language"] == robotwin.env.get_task_language() + assert all(set(view) == {"rgb"} for view in observation["views"].values()) + assert robotwin.env.renders[3:9] == [ + (camera, False) for _ in range(2) for camera in ROBOTWIN_CAMERA_NAMES + ] + assert all(actions.shape == (50, 16) for actions, _ in robotwin.env.chunks) + action = result.data["log"]["result"] + assert action["prompt"] == robotwin.env.get_task_language() + assert action["agent_prompt_ignored"] is True + assert action["ignored_agent_prompt"] == "ignored instruction" + assert action["requested_steps"] == action["executed_steps"] == 100 + assert result.data["state"]["episode_status"]["native_actions"] == 100 + assert result.data["state"]["episode_status"]["policy_actions"] == 100 + assert robotwin.toolkit.state.latest_step == 1 + + +@pytest.mark.parametrize( + "tool_name,arguments", + [ + ("lingbot_act", {"chunks": 2}), + ("set_gripper", {"arm": "left", "val": 1}), + ], +) +@pytest.mark.parametrize("success", [True, False]) +def test_native_success_and_budget_exhaustion_stop_execution( + robotwin, tool_name, arguments, success +): + if success: + robotwin.env.success_at = 2 + else: + robotwin.env.last_info["episode_status"]["step_lim"] = 2 + result = robotwin.toolkit.execute_tool(tool_name, arguments) + assert not result.is_error + action = result.data["log"]["result"] + assert action["executed_steps"] == 2 + assert action["completed"] is False + assert action["stop_reason"] == ( + "native_success" if success else "budget_exhausted" + ) + assert result.data["terminated"] is success + assert result.data["truncated"] is (not success) + assert robotwin.toolkit.solved() is success + if tool_name == "lingbot_act": + robotwin.toolkit.execute_tool(tool_name, arguments) + assert len(robotwin.model.observations) == 1 + + +@pytest.mark.parametrize( + "native_success,requested,reported", + [ + (False, "success", "failure"), + (True, "failure", "success"), + (False, "stuck", "stuck"), + ], +) +def test_finish_verifies_native_status_and_does_not_capture( + robotwin, native_success, requested, reported +): + robotwin.env.last_info["episode_status"]["eval_success"] = native_success + result = robotwin.toolkit.execute_tool( + "finish", {"status": requested, "summary": "done"} + ) + assert not result.is_error + assert result.data["status"] == reported + assert result.data["success"] is native_success + assert robotwin.toolkit.finish_result == { + "status": reported, + "summary": "done", + "requested_success": requested == "success", + "success": native_success, + "episode_status": { + **robotwin.env.last_info["episode_status"], + "policy_actions": 0, + "native_actions": 0, + }, + } + assert robotwin.toolkit.state.latest_step == 0 + assert len(robotwin.env.renders) == 3 -def test_fake_and_real_implement_toolkit_primitive_protocol() -> None: - for primitive_type in (RoboTwinPrimitives, FakeRoboTwinPrimitives): - missing = { - name - for name in PRIMITIVE_METHODS - if not callable(getattr(primitive_type, name, None)) - } - assert missing == set(), ( - f"{primitive_type.__name__} is missing toolkit methods: {sorted(missing)}" - ) - - -def test_toolkit_constructs_and_captures_an_initial_observation( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - FakeRoboTwinPrimitives.instances.clear() - dumped: list[dict[str, Any]] = [] - monkeypatch.setattr( - templates, "default_variables", lambda: {"output_dir": "/offline/output"} +def test_finish_still_ends_planner_when_native_status_is_unavailable(robotwin): + def failed_status(): + raise RuntimeError("status unavailable") + + robotwin.toolkit._robot.status = failed_status + result = robotwin.toolkit.execute_tool( + "finish", {"status": "success", "summary": "done"} ) - monkeypatch.setattr(toolkit, "RoboTwinPrimitives", FakeRoboTwinPrimitives) - monkeypatch.setattr(toolkit, "get_output_dir", lambda: tmp_path) - monkeypatch.setattr( - toolkit.RoboTwinToolkit, - "_capture_full_observation", - lambda self: {"views": {}, "robot_state": {}, "task_language": "offline"}, + assert not result.is_error + assert result.data["runtime_error"] == "RuntimeError: status unavailable" + assert robotwin.toolkit.finish_result == { + "status": "error", + "summary": "done", + "requested_status": "success", + "requested_success": True, + "runtime_error": "RuntimeError: status unavailable", + } + + +def test_persisted_perception_uses_same_step_and_view_without_rendering(robotwin): + tk = robotwin.toolkit + sampled = tk.execute_tool( + "sample_world_xyz", {"view": "head", "pixels": [[1, 2]], "neighborhood": 0} ) - monkeypatch.setattr( - toolkit.tools, - "dump_observation", - lambda observation, env_state, status, log: ( - dumped.append({"observation": observation, "status": status, "log": log}) - or _record() - ), + assert not sampled.is_error + np.testing.assert_allclose(sampled.data["samples"][0]["xyz"], [1, -0.5, -1]) + queried = tk.execute_tool( + "query_world_map", {"view": "left_wrist", "bbox": [0, 0, 2, 3], "max_points": 2} ) - monkeypatch.setattr( - toolkit.tools, - "view_env_state", - readonly(lambda step=-1, *, state: {"step": step}), + assert not queried.is_error + assert queried.data["valid_points"] == 6 + assert queried.data["returned_points"] == 2 + assert [p["pixel"] for p in queried.data["points"]] == [[0, 0], [1, 2]] + assert tk.state.latest_step == 0 + assert len(robotwin.env.renders) == 3 + outside = tk.execute_tool("sample_world_xyz", {"view": "head", "pixels": [[99, 0]]}) + assert outside.is_error and outside.data["code"] == "pixel_out_of_bounds" + missing = tk.execute_tool( + "query_world_map", {"view": "missing", "bbox": [0, 0, 1, 1]} ) + assert missing.is_error and missing.data["code"] == "view_not_found" + + +def test_malformed_planner_waypoint_cannot_broadcast_into_joint_targets(robotwin): + robotwin.env.path = np.ones((2, 1)) + result = robotwin.toolkit.execute_tool("move_to", {"arm": "left", "xyz": [1, 2, 3]}) + assert result.is_error and "shape [N,6]" in result.error + assert not robotwin.env.steps + assert robotwin.toolkit.state.latest_step == 1 + - robot_toolkit = toolkit.RoboTwinToolkit( - runtime_kwargs={"env": object(), "model": object(), "seed": 7}, - dashboard_events=NullDashboardEventSink(), - memory=MemoryManager(tmp_path / "memory"), +def test_perception_reports_missing_initial_record(robotwin): + robotwin.toolkit.state.reset() + result = robotwin.toolkit.execute_tool( + "sample_world_xyz", {"view": "head", "pixels": [[0, 0]]} ) + assert result.is_error and result.data["code"] == "state_not_found" + assert robotwin.toolkit.execute_tool("view_env_state", {}).is_error + + +@pytest.mark.parametrize("components", [{"env"}, {"vla"}, None]) +def test_runtime_builds_native_clients_and_resets_exactly_once( + monkeypatch, tmp_path, components, robotwin +): + env = robotwin.env + metadata = robot_spec.env_runtime_contract( + task_name="stack_blocks", + task_config="demo_randomized", + seed=7, + max_episode_steps=200, + ) + calls = [] - assert _tool_names(robot_toolkit) == EXPECTED_TOOLS - assert _readonly_names(robot_toolkit) == COMMON_TOOLS | { - "view_env_state", - "sample_world_xyz", - "query_world_map", - } - assert len(dumped) == 1 - assert dumped[0]["log"] == { - "command": {"action": "reset"}, - "result": {"actual_seed": 7, "success": True}, - "elapsed_s": 0.0, - } - primitive = FakeRoboTwinPrimitives.instances[0] - assert primitive.status_calls == 1 - assert primitive.recording_started is True - assert callable(primitive.kwargs["check_cancelled"]) - - robot_toolkit.get_env_state = lambda *, command, result, elapsed_s: dict(result) - render = robot_toolkit.execute_tool("render", {}) - assert render.result == {"success": True} - finish = robot_toolkit.execute_tool( - "finish", {"status": "failure", "summary": "offline"} - ) - assert finish.is_finish is True + def call(name, **kwargs): + calls.append(name) + if name == "env.get_env_meta": + return metadata + assert name == "env.reset" + return {}, {**env.last_info, "instruction": env.get_task_language()} + + rpc = SimpleNamespace(call=call) + daemons = {name: object() for name in ("env", "vla")} + + def spawn(owned, events, name, starter): + owned[name] = daemons[name] + return daemons[name], rpc + + monkeypatch.setattr(robot_spec, "try_spawn_server", spawn) + monkeypatch.setattr(robot_spec, "try_wait_server", lambda *a, post_fn: post_fn()) + monkeypatch.setattr( + robot_spec, + "_spawn_vla_server", + lambda *a: (daemons["vla"], ("localhost", 9000)), + ) + monkeypatch.setattr(robot_spec, "_wait_for_tcp", lambda *a, **kw: None) + contracts = [] + monkeypatch.setattr( + LingBotVLAClient, + "validate_contract", + lambda self, contract: contracts.append(contract), + ) + args = argparse.Namespace( + task_name="stack_blocks", + task_config="demo_randomized", + seed=7, + max_episode_steps=200, + ) + owned, resources = robot_spec._init_runtime( + args, tmp_path, NullDashboardEventSink(), components + ) + selected = {"env", "vla"} if components is None else components + assert set(owned) == {daemons[name] for name in selected} + assert set(resources) == ( + {"env", "seed", "seed_mode"} if "env" in selected else set() + ) | ({"model"} if "vla" in selected else set()) + if "env" in selected: + assert isinstance(resources["env"], RoboTwinEnvClient) + assert resources["seed"] == 7 and resources["seed_mode"] == "exact" + assert calls == ["env.get_env_meta", "env.reset"] + if "vla" in selected: + assert isinstance(resources["model"], LingBotVLAClient) + assert contracts == [robot_spec.vla_runtime_contract()] diff --git a/tests/unit_tests/robots/test_tool_schema_contracts.py b/tests/unit_tests/robots/test_tool_schema_contracts.py index 61ce9889d..862794880 100644 --- a/tests/unit_tests/robots/test_tool_schema_contracts.py +++ b/tests/unit_tests/robots/test_tool_schema_contracts.py @@ -14,82 +14,119 @@ from __future__ import annotations +import inspect +import json +from importlib import import_module +from pathlib import Path + import pytest -from robots.libero import tools as libero_tools -from robots.robocasa import tools as robocasa_tools -from robots.robotwin import tools as robotwin_tools +from rpent.tools import Tool +from rpent.tools.common_tools import COMMON_TOOLS -ROBOT_SCHEMAS = { - "libero": libero_tools.TOOLS_SPEC, - "robocasa": robocasa_tools.TOOLS_SPEC, - "robotwin": robotwin_tools.TOOLS_SPEC, -} +ROBOT_NAMES = ("franka", "dual_franka", "libero", "robocasa", "robotwin") -EXPECTED_TOOL_NAMES = { - "libero": { - "reset", - "view_env_state", - "move_to", - "pi0_pick", - "pi0_doubled", - "release", - "set_gripper", - "rotate_wrist", - "rotate_pitch", - "move_pose", - "view_camera_meta", - "segment", - "back_project", - }, - "robocasa": { - "move_to", - "move_delta", - "rotate_pitch", - "set_gripper", - "release", - "scripted_grasp", - "rldx_skill", - "rldx_arm", - "navigate_to", - "move_base", - "reset", - "view_env_state", - "back_project_batch", - "query_world_map", - "finish", - }, - "robotwin": { - "view_env_state", - "render", - "sample_world_xyz", - "query_world_map", - "lingbot_act", - "move_to", - "rotate_wrist", - "set_gripper", - "release", - "finish", +# These existing optional inputs now explicitly advertise their None value. +NULLABLE_INPUTS = { + "franka": { + "back_project": ("step",), + "back_project_correspondence": ( + "third_person_row", + "third_person_col", + "wrist_row", + "wrist_col", + "pixels", + "step", + ), }, + "dual_franka": {"back_project": ("step",), "segment": ("point", "step")}, } -@pytest.mark.parametrize("robot_name", sorted(ROBOT_SCHEMAS)) -def test_robot_tool_names_are_an_explicit_unique_contract(robot_name: str) -> None: - specs = ROBOT_SCHEMAS[robot_name] - names = [spec["name"] for spec in specs] +def _tools(group: str) -> tuple[Tool, ...]: + if group == "common": + return COMMON_TOOLS + module = import_module(f"robots.{group}.tools") + return getattr(module, f"{group.upper()}_TOOLS") - assert set(names) == EXPECTED_TOOL_NAMES[robot_name] - assert len(names) == len(set(names)) +@pytest.mark.parametrize("group", ("common", *ROBOT_NAMES)) +def test_input_schemas_match_before_native_migration(group: str) -> None: + # Extracted from historical TOOLS_SPEC declarations and the API image reader; + # expected schemas must not be regenerated from the native parameter models. + root = Path(__file__).parent + path = ( + root / "fixtures/common_tool_contracts.json" + if group == "common" + else root / group / "fixtures/pre_native_tool_contracts.json" + ) + baseline = json.loads(path.read_text(encoding="utf-8")) + expected = baseline["schemas"] + if group == "common": + # list_dir now resolves the directory from the invocation context. + expected["list_dir"]["properties"]["path"]["description"] = ( + "Directory path. Defaults to the current task's output directory." + ) + elif group == "libero": + # Task-card replay added these parameters after the historical snapshot. + expected["pi0_pick"]["properties"].update( + { + "gripper_open_thresh": { + "type": "number", + "description": "Minimum finger separation accepted as a held object (default 0.0)", + }, + "descent_thresh": { + "type": "number", + "description": "Required descent before lift detection, m (default 0.10)", + }, + } + ) + elif group == "robocasa": + # The preparation PR removed these unimplemented perception tools. + del expected["view_camera_meta"] + del expected["back_project"] + elif group == "robotwin": + # The native finish declaration adds parameter documentation. + expected["finish"]["properties"]["status"]["description"] = ( + "Requested task outcome." + ) + expected["finish"]["properties"]["summary"]["description"] = ( + "Summary of what worked and what failed." + ) + + for name, parameters in NULLABLE_INPUTS.get(group, {}).items(): + for parameter in parameters: + schema = expected[name]["properties"][parameter] + schema["type"] = [schema["type"], "null"] -@pytest.mark.parametrize("robot_name", sorted(ROBOT_SCHEMAS)) + tools = _tools(group) + actual = {item.name: item.input_schema for item in tools} + assert len(actual) == len(tools) + assert actual.keys() == expected.keys() + for name, schema in actual.items(): + # Native tool calls now reject unknown top-level arguments. + expected[name]["additionalProperties"] = False + assert schema == expected[name], f"{group}.{name} input schema changed" + + if "descriptions" in baseline: + assert {item.name: item.description for item in tools} == baseline[ + "descriptions" + ] + + for item in tools: + for name, parameter in actual[item.name]["properties"].items(): + if "default" in parameter: + assert ( + item.args_schema.model_fields[name].default == parameter["default"] + ), f"{group}.{item.name}.{name} default differs from its schema" + + +@pytest.mark.parametrize("robot_name", ROBOT_NAMES) def test_robot_tool_schemas_have_valid_object_inputs(robot_name: str) -> None: - for spec in ROBOT_SCHEMAS[robot_name]: - assert set(spec) >= {"name", "description", "input_schema"} - assert isinstance(spec["description"], str) and spec["description"].strip() + for item in _tools(robot_name): + assert isinstance(item.description, str) and item.description.strip() - input_schema = spec["input_schema"] + input_schema = item.input_schema assert input_schema["type"] == "object" properties = input_schema.get("properties", {}) required = input_schema.get("required", []) @@ -98,12 +135,13 @@ def test_robot_tool_schemas_have_valid_object_inputs(robot_name: str) -> None: assert set(required) <= set(properties) -def test_robot_action_schemas_keep_bounded_vector_shapes() -> None: - schema_sets = [ - {spec["name"]: spec for spec in libero_tools.TOOLS_SPEC}["move_to"], - {spec["name"]: spec for spec in robotwin_tools.TOOLS_SPEC}["move_to"], - ] - for spec in schema_sets: - xyz = spec["input_schema"]["properties"]["xyz"] - assert xyz["type"] == "array" - assert xyz["minItems"] == xyz["maxItems"] == 3 +@pytest.mark.parametrize("robot_name", ROBOT_NAMES) +def test_owned_tool_collections_satisfy_executor_invariants(robot_name): + tools = (*COMMON_TOOLS, *_tools(robot_name)) + names = [item.name for item in tools] + assert len(names) == len(set(names)) + for item in tools: + parameter = inspect.signature(item.handler).parameters["ctx"] + assert parameter.kind is inspect.Parameter.KEYWORD_ONLY + assert parameter.default is inspect.Parameter.empty + assert "ctx" not in item.input_schema["properties"] diff --git a/tests/unit_tests/robots/test_toolkit_contracts.py b/tests/unit_tests/robots/test_toolkit_contracts.py index 889a70726..60a68f68a 100644 --- a/tests/unit_tests/robots/test_toolkit_contracts.py +++ b/tests/unit_tests/robots/test_toolkit_contracts.py @@ -14,18 +14,13 @@ from __future__ import annotations +from importlib import import_module from pathlib import Path from types import SimpleNamespace from typing import Any import pytest -from robots.libero import robot_spec as libero_robot_spec -from robots.libero import toolkit as libero_toolkit -from robots.robocasa import robot_spec as robocasa_robot_spec -from robots.robocasa import toolkit as robocasa_toolkit -from robots.robotwin import robot_spec as robotwin_robot_spec -from robots.robotwin import toolkit as robotwin_toolkit from rpent.dashboard.events import NullDashboardEventSink from rpent.robots import RunConfig @@ -39,21 +34,17 @@ def _run_config(memory_dir: Path, *, recipe_tag: str = "cell-s0") -> RunConfig: ) -@pytest.mark.parametrize( - ("robot_spec", "toolkit_module", "toolkit_name", "configured_leaf"), - [ - (robocasa_robot_spec, robocasa_toolkit, "RoboCasaToolkit", "memory"), - (robotwin_robot_spec, robotwin_toolkit, "RoboTwinToolkit", "memory"), - ], -) +@pytest.mark.parametrize("robot_name", ["robocasa", "robotwin"]) def test_evaluation_toolkit_factories_use_configured_read_only_memory( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, - robot_spec: Any, - toolkit_module: Any, - toolkit_name: str, - configured_leaf: str, + robot_name: str, ) -> None: + robot_spec = import_module(f"robots.{robot_name}.robot_spec") + toolkit_module = import_module(f"robots.{robot_name}.toolkit") + toolkit_name = {"robocasa": "RoboCasaToolkit", "robotwin": "RoboTwinToolkit"}[ + robot_name + ] captured: dict[str, Any] = {} def fake_toolkit(**kwargs: Any) -> SimpleNamespace: @@ -61,50 +52,42 @@ def fake_toolkit(**kwargs: Any) -> SimpleNamespace: return SimpleNamespace(**kwargs) monkeypatch.setattr(toolkit_module, toolkit_name, fake_toolkit) - resources_dir = tmp_path / robot_spec.__name__ - configured_dir = resources_dir / configured_leaf - memory_dir = resources_dir / "memory" - + memory_dir = tmp_path / robot_name / "memory" toolkit = robot_spec.get_toolkit( runtime_kwargs={"env": "offline"}, dashboard_events=NullDashboardEventSink(), - config=_run_config(configured_dir), + config=_run_config(memory_dir), ) assert toolkit.memory.root == memory_dir.resolve() - write = toolkit.memory.get_common_tool_bindings()["write_text_file"][1] with pytest.raises(PermissionError, match="writing to memory is denied"): - write(str(memory_dir / "global" / "strategy.md"), "changed") + toolkit.memory.authorize_write(memory_dir / "global" / "strategy.md") assert captured["runtime_kwargs"] == {"env": "offline"} + assert captured["output_dir"] == memory_dir.parent / "run" -@pytest.mark.parametrize( - ("robot_name", "robot_spec", "toolkit_module", "toolkit_name"), - [ - ("libero", libero_robot_spec, libero_toolkit, "LiberoToolkit"), - ("robotwin", robotwin_robot_spec, robotwin_toolkit, "RoboTwinToolkit"), - ], -) +@pytest.mark.parametrize("robot_name", ["libero", "robocasa", "robotwin"]) def test_toolkit_factories_fall_back_to_each_robot_memory_root( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, robot_name: str, - robot_spec: Any, - toolkit_module: Any, - toolkit_name: str, ) -> None: + robot_spec = import_module(f"robots.{robot_name}.robot_spec") + toolkit_module = import_module(f"robots.{robot_name}.toolkit") + toolkit_name = { + "libero": "LiberoToolkit", + "robocasa": "RoboCasaToolkit", + "robotwin": "RoboTwinToolkit", + }[robot_name] default_memory = tmp_path / robot_name / "memory" + + def get_memory_dir(requested_robot): + assert requested_robot == robot_name + return default_memory + + monkeypatch.setattr(robot_spec, "get_memory_dir", get_memory_dir) monkeypatch.setattr( - robot_spec, - "get_memory_dir", - lambda requested_robot: ( - default_memory if requested_robot == robot_name else None - ), - ) - monkeypatch.setattr( - toolkit_module, - toolkit_name, - lambda **kwargs: SimpleNamespace(**kwargs), + toolkit_module, toolkit_name, lambda **kwargs: SimpleNamespace(**kwargs) ) config = RunConfig( recipe_tag="cell-s0", @@ -112,11 +95,9 @@ def test_toolkit_factories_fall_back_to_each_robot_memory_root( prompt_vars={}, task_desc={}, ) - toolkit = robot_spec.get_toolkit( runtime_kwargs={}, dashboard_events=NullDashboardEventSink(), config=config, ) - assert toolkit.memory.root == default_memory.resolve() diff --git a/tests/unit_tests/rpent/cli/test_main_contracts.py b/tests/unit_tests/rpent/cli/test_main_contracts.py index f3e63ff79..3cd18eec6 100644 --- a/tests/unit_tests/rpent/cli/test_main_contracts.py +++ b/tests/unit_tests/rpent/cli/test_main_contracts.py @@ -393,9 +393,10 @@ def test_full_cli_exploration_finalizes_memory_without_starting_gpu_runtime( monkeypatch: pytest.MonkeyPatch, ) -> None: cli = _cli_module() + from rpent.planner.base import PlannerResult from rpent.robots import PromptBundle, RobotSpec, RunConfig - from rpent.tools.toolkit import ToolResult + from rpent.tools import ToolResult calls: dict[str, Any] = {} @@ -415,17 +416,12 @@ def __init__(self) -> None: self.calls: list[tuple[str, dict[str, Any]]] = [] self.closed = False self.memory = FakeMemoryManager() + self.finish_result = None def execute_tool(self, name: str, args: dict[str, Any]) -> ToolResult: self.calls.append((name, args)) - return ToolResult( - name, - { - "_finish": True, - "status": args["status"], - "summary": args["summary"], - }, - ) + self.finish_result = {"status": args["status"], "summary": args["summary"]} + return ToolResult(data={"_finish": True, **self.finish_result}) def close(self) -> None: self.closed = True @@ -455,12 +451,12 @@ def solve( "input_queue": input_queue, "dashboard_interaction": dashboard_interaction, } - finish = toolkit.execute_tool( + toolkit.execute_tool( "finish", {"status": "success", "summary": "simulated task complete"}, ) return PlannerResult( - finish_result=finish.result, + finish_result=toolkit.finish_result, messages=[{"role": "assistant", "content": "finished offline"}], stats={ "total_input_tokens": 0, @@ -568,7 +564,6 @@ def reject_memory_sync(*args: Any, **kwargs: Any) -> None: transcript = json.loads((tmp_path / "transcript_libero_s0.json").read_text()) assert transcript["robot"] == "libero" assert transcript["finish"] == { - "_finish": True, "status": "success", "summary": "simulated task complete", } diff --git a/tests/unit_tests/rpent/dashboard/test_state_contracts.py b/tests/unit_tests/rpent/dashboard/test_state_contracts.py index acf858350..fe891e85a 100644 --- a/tests/unit_tests/rpent/dashboard/test_state_contracts.py +++ b/tests/unit_tests/rpent/dashboard/test_state_contracts.py @@ -19,6 +19,7 @@ import numpy as np import pytest +from fastapi.testclient import TestClient from rpent.dashboard.events import ( RunStartedEvent, @@ -33,10 +34,12 @@ InteractionUnavailableError, UnknownDashboardMessageError, ) +from rpent.dashboard.server import DashboardServer from rpent.dashboard.spec import DashboardSpec from rpent.dashboard.state import DashboardState +from rpent.memory import MemoryManager from rpent.session import EnvState -from rpent.tools.toolkit import Toolkit, ToolResult +from rpent.tools import Toolkit, ToolResult, readonly, tool DASHBOARD_SPEC: DashboardSpec = { "task": { @@ -226,13 +229,14 @@ def test_dashboard_primitives_are_available_only_while_planner_is_idle( state = _ready_state(tmp_path) _claim_started_task(state) toolkit = MagicMock(spec=Toolkit) - toolkit.get_tools_spec.return_value = [ - { - "name": "move_to", - "input_schema": {"type": "object", "additionalProperties": False}, - } - ] - tool_result = ToolResult(name="move_to", result={"ok": True}) + + @tool + def move_to(*, ctx) -> ToolResult: + """Move the robot.""" + return ToolResult(data={"ok": True}) + + toolkit.list_tools.return_value = (move_to,) + tool_result = ToolResult(data={"ok": True}) toolkit.execute_tool.return_value = tool_result state.bind_toolkit(toolkit) @@ -244,7 +248,9 @@ def test_dashboard_primitives_are_available_only_while_planner_is_idle( state.set_planner_activity("idle", accepting_input=True) assert state.snapshot()["primitives_available"] is True - assert state.primitive_specs() == toolkit.get_tools_spec.return_value + assert state.primitive_specs() == [ + {"name": "move_to", "input_schema": move_to.input_schema} + ] assert state.execute_primitive("move_to", {}) is tool_result state.set_planner_activity("busy") @@ -383,3 +389,93 @@ def test_dashboard_step_events_offset_new_traces_and_resolve_action_video( assert detail["timeline"][1]["terminated"] is True assert state.frame("camera") == second_env.load_bytes("camera.png") assert state.action_video_path(0) == first_env.artifact_path("action.mp4", step=0) + + +def test_primitive_http_uses_native_validation_results_and_observation(tmp_path): + state = _ready_state(tmp_path) + _claim_started_task(state) + calls = [] + + @tool + def move_to(distance: int, *, ctx) -> ToolResult: + """Move a test robot.""" + calls.append(distance) + if distance < 0: + return ToolResult(error="motion rejected") + return ToolResult(data={"distance": distance}) + + @tool + @readonly + def finish(status: str, summary: str, *, ctx) -> ToolResult: + """Finish the test task.""" + return ToolResult(data={"status": status, "summary": summary}) + + class RobotToolkit(Toolkit): + def _capture_observation(self, *, command, result, elapsed_s): + with self.state.record_step( + state={"distance": calls[-1]}, + command=command, + result=result.to_dict(), + elapsed_s=elapsed_s, + ): + self.state.save("overhead.png", np.zeros((2, 2, 3), dtype=np.uint8)) + return {"distance": calls[-1]}, [] + + toolkit = RobotToolkit( + robot=None, + tools=(move_to, finish), + state=EnvState(tmp_path / "env"), + memory=MemoryManager(root=tmp_path / "memory"), + output_dir=tmp_path, + dashboard_events=state, + ) + state.bind_toolkit(toolkit) + state.set_planner_activity("idle", accepting_input=True) + with TestClient(DashboardServer(state=state)._app) as client: + specs = client.get("/api/session/primitives").json()["primitives"] + assert specs == [{"name": "move_to", "input_schema": move_to.input_schema}] + for arguments in ({}, {"distance": "bad"}): + response = client.post( + "/api/session/primitive", + json={"name": "move_to", "arguments": arguments}, + ) + assert response.status_code == 422 + assert "Invalid arguments for move_to" in response.json()["error"] + assert calls == [] + response = client.post( + "/api/session/primitive", json={"name": "finish", "arguments": {}} + ) + assert response.status_code == 403 + # Coercion is the same native validation used by planner tool calls. + response = client.post( + "/api/session/primitive", + json={"name": "move_to", "arguments": {"distance": "2"}}, + ) + assert response.json() == {"ok": True} + assert calls == [2] + snapshot = client.get("/api/session/state").json() + assert snapshot["frame_available"] == {"overhead": True} + frame = client.get("/api/session/frame", params={"kind": "overhead"}) + assert frame.headers["content-type"] == "image/png" + assert frame.content == toolkit.state.load_bytes("overhead.png") + assert ( + client.get("/api/session/frame", params={"kind": "missing"}).status_code + == 404 + ) + response = client.post( + "/api/session/primitive", + json={"name": "move_to", "arguments": {"distance": -1}}, + ) + assert response.status_code == 422 + assert response.json() == {"error": "motion rejected"} + state.set_planner_activity("busy") + assert ( + client.post( + "/api/session/primitive", + json={"name": "move_to", "arguments": {"distance": 3}}, + ).status_code + == 409 + ) + assert calls == [2, -1] + state.unbind_toolkit(toolkit) + toolkit.close() diff --git a/tests/unit_tests/rpent/memory/test_authorization.py b/tests/unit_tests/rpent/memory/test_authorization.py new file mode 100644 index 000000000..ffb9541a4 --- /dev/null +++ b/tests/unit_tests/rpent/memory/test_authorization.py @@ -0,0 +1,103 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import pytest + +from rpent.memory import MemoryManager + + +@pytest.fixture +def memory(tmp_path, monkeypatch): + monkeypatch.setenv("RPENT_REPO_ROOT", str(tmp_path)) + return MemoryManager( + tmp_path / "memory/libero", + memory_access="inbox_write", + inbox_cell_tag="current", + ) + + +@pytest.mark.parametrize( + "relative", + [ + "", + "MEMORY.md", + "notes.md", + "global/strategy.md", + "suite/task.md", + "task_only/audit.json", + "results/recipe.jsonl", + ], +) +def test_published_scopes_are_readable_and_never_writable(memory, relative): + path = memory.root / relative + assert memory.authorize_read(path) == path + with pytest.raises(PermissionError): + memory.authorize_write(path) + + +@pytest.mark.parametrize( + "relative", + [ + "_internal/inbox/other/draft.md", + "_internal/merged/current/draft.md", + "private/file.md", + ], +) +def test_other_private_paths_are_denied(memory, relative): + for authorize in (memory.authorize_read, memory.authorize_write): + with pytest.raises(PermissionError): + authorize(memory.root / relative) + + +def test_current_inbox_is_readable_and_writable_only_in_exploration(memory): + path = memory.root / "_internal/inbox/current/draft.md" + assert memory.authorize_read(path) == path + assert memory.authorize_write(path) == path + evaluation = MemoryManager(memory.root) + for authorize in (evaluation.authorize_read, evaluation.authorize_write): + with pytest.raises(PermissionError): + authorize(path) + + +def test_nonmemory_paths_and_relative_paths_resolve(memory, tmp_path): + assert memory.authorize_write("output/note.txt") == tmp_path / "output/note.txt" + assert memory.authorize_read("memory/libero/MEMORY.md") == memory.root / "MEMORY.md" + assert memory.authorize_read("") == tmp_path + + +def test_foreign_memory_and_symlink_escapes_are_denied(memory, tmp_path): + foreign = tmp_path / "memory/robotwin/global/note.md" + foreign.parent.mkdir(parents=True) + foreign.write_text("foreign") + link = tmp_path / "link.md" + link.symlink_to(foreign) + for authorize in (memory.authorize_read, memory.authorize_write): + for path in (foreign, link): + with pytest.raises(PermissionError, match="another robot"): + authorize(path) + alias = tmp_path / "memory_alias" + alias.symlink_to(tmp_path / "memory", target_is_directory=True) + with pytest.raises(PermissionError, match="another robot"): + memory.authorize_read(alias / "robotwin/global/note.md") + + +def test_authorization_returns_resolved_target_for_io(memory, tmp_path): + destination = tmp_path / "files/note.txt" + destination.parent.mkdir() + destination.write_text("note") + link = tmp_path / "note-link.txt" + link.symlink_to(destination) + assert memory.authorize_read(link) == destination + assert memory.authorize_write(link) == destination diff --git a/tests/unit_tests/rpent/planner/_native_helpers.py b/tests/unit_tests/rpent/planner/_native_helpers.py new file mode 100644 index 000000000..0583b79f8 --- /dev/null +++ b/tests/unit_tests/rpent/planner/_native_helpers.py @@ -0,0 +1,121 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import threading +from pathlib import Path +from typing import Any + +from mcp import types + +from rpent.memory import MemoryManager +from rpent.session import EnvState +from rpent.tools import ( + ToolContext, + Toolkit, + ToolResult, + readonly, + tool, +) + +PNG = b"\x89PNG\r\n\x1a\ncontract-image" + + +@tool +@readonly +def inspect_scene(detail: str = "low", *, ctx: ToolContext[Any]) -> ToolResult: + """Inspect the current scene.""" + return ctx.robot.response() + + +@tool +def finish(status: str, summary: str, *, ctx: ToolContext) -> ToolResult: + """Return the outcome configured by the planner test.""" + return ctx.robot.finish(status, summary) + + +class FakeToolkit(Toolkit): + def __init__( + self, + output: Path, + result: dict[str, Any] | None = None, + *, + images: list[bytes] | None = None, + accepted: dict[str, Any] | None = None, + ): + self.result = result if result is not None else {"value": "ok"} + self.images = images or [] + self.accepted = accepted + self.calls = [] + self.cancel_calls = 0 + super().__init__( + state=EnvState(output), + memory=MemoryManager(output / "memory"), + robot=self, + output_dir=output, + tools=(finish, inspect_scene), + ) + + def response(self): + data = dict(self.result) + error = data.pop("error", None) + return ToolResult(data=data, error=error, images=self.images) + + def finish(self, status: str, summary: str) -> ToolResult: + result = self.response() + if not result.is_error: + accepted = self.accepted or { + "status": status, + "summary": summary, + } + result.data.update({"_finish": True, **accepted}) + return result + + def execute_tool(self, name, args): + self.calls.append((name, args)) + return super().execute_tool(name, args) + + def cancel_active_and_wait(self): + self.cancel_calls += 1 + super().cancel_active_and_wait() + + +class RecordingSink: + def __init__(self) -> None: + self.events: list[Any] = [] + + @property + def enabled(self) -> bool: + return True + + def emit(self, event: Any) -> None: + self.events.append(event) + + +async def call_sdk_tool(config, name, arguments): + request = types.CallToolRequest( + method="tools/call", + params=types.CallToolRequestParams(name=name, arguments=arguments), + ) + response = await config["instance"].request_handlers[types.CallToolRequest](request) + return response.root + + +@tool +@readonly +def read(number: int = 0, *, ctx: ToolContext[threading.Barrier]) -> ToolResult: + """Read concurrently and return the validated number.""" + ctx.robot.wait(timeout=3) + return ToolResult(data={"number": number}) diff --git a/tests/unit_tests/rpent/planner/conftest.py b/tests/unit_tests/rpent/planner/conftest.py new file mode 100644 index 000000000..e9d7d3db5 --- /dev/null +++ b/tests/unit_tests/rpent/planner/conftest.py @@ -0,0 +1,31 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from ._native_helpers import FakeToolkit + + +@pytest.fixture +def make_toolkit(tmp_path): + instances = [] + + def make(*args, **kwargs): + toolkit = FakeToolkit(tmp_path / str(len(instances)), *args, **kwargs) + instances.append(toolkit) + return toolkit + + yield make + for toolkit in instances: + toolkit.close() diff --git a/tests/unit_tests/rpent/planner/test_api_contracts.py b/tests/unit_tests/rpent/planner/test_api_contracts.py index 1d7101363..d3dd57230 100644 --- a/tests/unit_tests/rpent/planner/test_api_contracts.py +++ b/tests/unit_tests/rpent/planner/test_api_contracts.py @@ -15,7 +15,7 @@ from __future__ import annotations import asyncio -import base64 +import json import queue from typing import Any @@ -34,54 +34,10 @@ from rpent.planner.api_loop import ( ApiAgentLoop, _build_tools, - _content_blocks_to_pydantic, _make_tool_function, ) -from rpent.tools.toolkit import ToolResult - -class RecordingSink: - def __init__(self) -> None: - self.events: list[Any] = [] - - @property - def enabled(self) -> bool: - return True - - def emit(self, event: Any) -> None: - self.events.append(event) - - -class FakeToolkit: - state = None - - def __init__(self, result: dict[str, Any] | None = None) -> None: - self.result = result or {"ok": True} - self.calls: list[tuple[str, dict[str, Any]]] = [] - self.cancel_calls = 0 - - def get_tools_spec(self) -> list[dict[str, Any]]: - return [ - { - "name": "finish", - "description": "Finish after the environment accepts the result.", - "input_schema": { - "type": "object", - "properties": { - "status": {"type": "string"}, - "summary": {"type": "string"}, - }, - "required": ["status", "summary"], - }, - } - ] - - def execute_tool(self, name: str, args: dict[str, Any]) -> ToolResult: - self.calls.append((name, args)) - return ToolResult(name, dict(self.result)) - - def cancel_active_and_wait(self) -> None: - self.cancel_calls += 1 +from ._native_helpers import PNG, FakeToolkit, RecordingSink def solve_with_model( @@ -105,7 +61,9 @@ def solve_with_model( ) -def test_successful_finish_waits_for_its_tool_result() -> None: +def test_successful_finish_waits_for_its_tool_result( + make_toolkit, +) -> None: seen_instructions: list[str | None] = [] def model(messages: list[Any], info: Any) -> ModelResponse: @@ -127,14 +85,14 @@ def model(messages: list[Any], info: Any) -> ModelResponse: usage=RequestUsage(input_tokens=7, output_tokens=3), ) - toolkit = FakeToolkit() + toolkit = make_toolkit() sink = RecordingSink() result = solve_with_model(model, toolkit, sink) assert seen_instructions == ["Use tools carefully."] assert toolkit.calls == [("finish", {"status": "success", "summary": "done"})] assert result.finish_result == { - "_finish": True, + "value": "ok", "status": "success", "summary": "done", } @@ -152,7 +110,16 @@ def model(messages: list[Any], info: Any) -> ModelResponse: assert any(isinstance(event, UsageEvent) for event in sink.events) -def test_rejected_finish_does_not_end_the_run() -> None: +@pytest.mark.parametrize( + "payload", + [ + {"error": "finish refused by environment"}, + {"reading": float("nan")}, + {"reading": float("inf")}, + {"nested": {(1, 2): "invalid JSON key"}}, + ], +) +def test_rejected_finish_does_not_end_the_run(make_toolkit, payload) -> None: def model(messages: list[Any], info: Any) -> ModelResponse: del info if any( @@ -171,38 +138,42 @@ def model(messages: list[Any], info: Any) -> ModelResponse: ] ) - toolkit = FakeToolkit({"error": "finish refused by environment"}) + toolkit = make_toolkit(payload) result = solve_with_model(model, toolkit, RecordingSink()) assert result.finish_result is None assert result.error is None assert result.stats["tool_calls"] == 1 + error = payload.get("error", "Tool result serialization failed") assert any( - message.get("role") == "tool" - and message.get("content") == '{\n "error": "finish refused by environment"\n}' + message.get("role") == "tool" and error in message.get("content", "") for message in result.messages ) -def test_backend_failure_is_returned_without_escaping() -> None: +def test_backend_failure_is_returned_without_escaping( + make_toolkit, +) -> None: def model(messages: list[Any], info: Any) -> ModelResponse: del messages, info raise RuntimeError("provider failed") - result = solve_with_model(model, FakeToolkit(), RecordingSink()) + result = solve_with_model(model, make_toolkit(), RecordingSink()) assert result.finish_result is None assert result.error == "RuntimeError: provider failed" - assert result.messages == [{"role": "user", "content": "complete the task"}] + assert result.messages[0] == {"role": "user", "content": "complete the task"} -def test_timeout_cancels_active_toolkit_work() -> None: +def test_timeout_cancels_active_toolkit_work( + make_toolkit, +) -> None: async def model(messages: list[Any], info: Any) -> ModelResponse: del messages, info await asyncio.sleep(10) return ModelResponse(parts=[TextPart("unreachable")]) - toolkit = FakeToolkit() + toolkit = make_toolkit() result = solve_with_model( model, toolkit, @@ -211,11 +182,13 @@ async def model(messages: list[Any], info: Any) -> ModelResponse: ) assert result.error == "API planner timed out after 0.01s" - assert toolkit.cancel_calls == 1 - assert result.messages == [{"role": "user", "content": "complete the task"}] + assert toolkit.cancel_calls >= 1 + assert result.messages[0] == {"role": "user", "content": "complete the task"} -def test_queue_and_dashboard_inputs_are_rejected_before_model_use() -> None: +def test_queue_and_dashboard_inputs_are_rejected_before_model_use( + make_toolkit, +) -> None: calls = 0 def model(messages: list[Any], info: Any) -> ModelResponse: @@ -233,7 +206,7 @@ def model(messages: list[Any], info: Any) -> ModelResponse: planner.solve( system_prompt="", user_message="task", - toolkit=FakeToolkit(), + toolkit=make_toolkit(), max_turns=1, input_queue=queue.Queue(), dashboard_interaction=object(), @@ -242,58 +215,41 @@ def model(messages: list[Any], info: Any) -> ModelResponse: assert calls == 0 -def test_tool_schema_and_dispatch_are_mapped_to_pydantic_ai() -> None: - toolkit = FakeToolkit() +def test_tool_schema_and_dispatch_are_mapped_to_pydantic_ai( + make_toolkit, +) -> None: + toolkit = make_toolkit() tools = _build_tools(toolkit) - assert [tool.name for tool in tools] == ["read_image", "finish"] + assert [tool.name for tool in tools] == [ + "read_image", + *[tool.name for tool in toolkit.list_tools() if tool.name != "read_image"], + ] + assert "read_image" in [tool.name for tool in tools] assert all(tool.sequential for tool in tools) - finish = tools[1] - assert finish.description == "Finish after the environment accepts the result." - assert ( - finish.function_schema.json_schema - == toolkit.get_tools_spec()[0]["input_schema"] + finish = next(tool for tool in tools if tool.name == "finish") + assert finish.function_schema.json_schema == next( + tool.input_schema for tool in toolkit.list_tools() if tool.name == "finish" ) -def test_tool_result_conversion_keeps_text_and_images_separate() -> None: - raw_image = b"\x89PNG\r\ncontract-image" - encoded = base64.b64encode(raw_image).decode() - blocks = [ - {"type": "text", "text": "observation"}, - { - "type": "image", - "source": { - "type": "base64", - "media_type": "image/png", - "data": encoded, - }, - }, - ] - - text, images = _content_blocks_to_pydantic(blocks) - - assert text == "observation" - assert images == [BinaryContent(data=raw_image, media_type="image/png")] - assert blocks[1]["source"]["data"] == encoded - +def test_tool_result_conversion_keeps_text_and_images_separate( + make_toolkit, +) -> None: + toolkit = make_toolkit({"value": "visible"}, images=[PNG]) + result = asyncio.run(_make_tool_function(toolkit, "inspect_scene")()) + assert isinstance(result, ToolReturn) + assert json.loads(result.return_value) == {"value": "visible"} + assert result.content == [BinaryContent(data=PNG, media_type="image/png")] -def test_no_images_mode_suppresses_binary_tool_content() -> None: - toolkit = FakeToolkit({"value": "visible", "_image_bytes": b"secret pixels"}) - multimodal = _make_tool_function(toolkit, "finish")( - status="success", - summary="done", +def test_no_images_mode_suppresses_binary_tool_content( + make_toolkit, +) -> None: + toolkit = make_toolkit({"value": "visible"}, images=[PNG]) + result = asyncio.run( + _make_tool_function(toolkit, "inspect_scene", no_images=True)() ) - text_only = _make_tool_function(toolkit, "finish", no_images=True)( - status="success", - summary="done", - ) - - assert isinstance(multimodal, ToolReturn) - assert multimodal.return_value == '{\n "value": "visible"\n}' - assert len(multimodal.content or []) == 1 - assert isinstance(multimodal.content[0], BinaryContent) - assert text_only == '{\n "value": "visible"\n}' - assert "secret" not in text_only + assert json.loads(result) == {"value": "visible"} + assert "contract-image" not in result diff --git a/tests/unit_tests/rpent/planner/test_claude_contracts.py b/tests/unit_tests/rpent/planner/test_claude_contracts.py index 866dda45b..ce65ca419 100644 --- a/tests/unit_tests/rpent/planner/test_claude_contracts.py +++ b/tests/unit_tests/rpent/planner/test_claude_contracts.py @@ -29,59 +29,11 @@ _build_rpent_server, _ClaudeSessionDriver, _Recorder, - _tool_result_to_mcp, ) -from rpent.tools.toolkit import ToolResult +from rpent.planner.utils.http_mcp_server import mcp_result +from rpent.tools import ToolResult - -class RecordingSink: - def __init__(self) -> None: - self.events: list[Any] = [] - - @property - def enabled(self) -> bool: - return True - - def emit(self, event: Any) -> None: - self.events.append(event) - - -class FakeToolkit: - def __init__(self, result: dict[str, Any] | None = None) -> None: - self.result = result or {"value": "ok"} - self.calls: list[tuple[str, dict[str, Any]]] = [] - self.cancel_calls = 0 - - def get_tools_spec(self) -> list[dict[str, Any]]: - return [ - { - "name": "inspect_scene", - "description": "Inspect the current scene.", - "input_schema": { - "type": "object", - "properties": {"detail": {"type": "string"}}, - }, - }, - { - "name": "finish", - "description": "Finish the task.", - "input_schema": { - "type": "object", - "properties": { - "status": {"type": "string"}, - "summary": {"type": "string"}, - }, - "required": ["status", "summary"], - }, - }, - ] - - def execute_tool(self, name: str, args: dict[str, Any]) -> ToolResult: - self.calls.append((name, args)) - return ToolResult(name, dict(self.result)) - - def cancel_active_and_wait(self) -> None: - self.cancel_calls += 1 +from ._native_helpers import PNG, RecordingSink, call_sdk_tool class FakeSdkTools: @@ -93,35 +45,9 @@ def ClaudeAgentOptions(self, **kwargs: Any) -> dict[str, Any]: self.options = kwargs return kwargs - @staticmethod - def tool(name: str, description: str, schema: dict[str, Any]): - def decorate(function: Any) -> Any: - function.sdk_name = name - function.sdk_description = description - function.sdk_schema = schema - return function - - return decorate - - def create_sdk_mcp_server( - self, - *, - name: str, - version: str, - tools: list[Any], - ) -> dict[str, Any]: - self.created_server = {"name": name, "version": version, "tools": tools} - return self.created_server - def patch_sdk_surface(monkeypatch: pytest.MonkeyPatch, fake: FakeSdkTools) -> None: monkeypatch.setattr(claude_agent_sdk, "ClaudeAgentOptions", fake.ClaudeAgentOptions) - monkeypatch.setattr(claude_agent_sdk, "tool", fake.tool) - monkeypatch.setattr( - claude_agent_sdk, - "create_sdk_mcp_server", - fake.create_sdk_mcp_server, - ) def make_planner(tmp_path: Path, sink: RecordingSink, *, timeout_s: float = 1): @@ -139,12 +65,13 @@ def make_planner(tmp_path: Path, sink: RecordingSink, *, timeout_s: float = 1): def test_options_translate_builtin_and_rpent_tools_without_mutating_specs( + make_toolkit, tmp_path: Path, ) -> None: sink = RecordingSink() planner = make_planner(tmp_path, sink) - toolkit = FakeToolkit() - original_specs = toolkit.get_tools_spec() + toolkit = make_toolkit() + original_specs = toolkit.list_tools() fake_sdk = FakeSdkTools() options = planner._build_options(fake_sdk, toolkit=toolkit, max_turns=4) @@ -159,20 +86,25 @@ def test_options_translate_builtin_and_rpent_tools_without_mutating_specs( "Read", "Grep", "mcp__external__keep", - "mcp__rpent__inspect_scene", - "mcp__rpent__finish", + *[ + f"mcp__rpent__{tool.name}" + for tool in toolkit.list_tools() + if tool.name != "read_image" + ], ] assert options["add_dirs"] == [str(tmp_path), str(tmp_path / "memory")] assert options["setting_sources"] == [] - assert toolkit.get_tools_spec() == original_specs + assert toolkit.list_tools() == original_specs -def test_options_construct_with_the_installed_claude_sdk(tmp_path: Path) -> None: +def test_options_construct_with_the_installed_claude_sdk( + make_toolkit, tmp_path: Path +) -> None: planner = make_planner(tmp_path, RecordingSink()) options = planner._build_options( claude_agent_sdk, - toolkit=FakeToolkit(), + toolkit=make_toolkit(), max_turns=4, ) @@ -180,54 +112,47 @@ def test_options_construct_with_the_installed_claude_sdk(tmp_path: Path) -> None assert options.cwd == str(tmp_path) assert options.model == "fake-claude" assert options.max_turns == 4 - assert options.allowed_tools[-2:] == [ - "mcp__rpent__inspect_scene", - "mcp__rpent__finish", - ] - + assert "mcp__rpent__inspect_scene" in options.allowed_tools + assert "mcp__rpent__finish" in options.allowed_tools + assert "mcp__rpent__read_image" not in options.allowed_tools -def test_in_process_mcp_bridge_maps_schema_dispatch_and_errors() -> None: - toolkit = FakeToolkit({"error": "rejected", "_image_bytes": b"contract-image"}) - fake_sdk = FakeSdkTools() - - server = _build_rpent_server(fake_sdk, toolkit=toolkit) - - assert server["name"] == "rpent" - assert server["version"] == "0.1.0" - tools = {tool.sdk_name: tool for tool in server["tools"]} - assert tools["inspect_scene"].sdk_description == "Inspect the current scene." - assert ( - tools["inspect_scene"].sdk_schema == toolkit.get_tools_spec()[0]["input_schema"] - ) - response = asyncio.run(tools["inspect_scene"]({"detail": "high"})) +def test_in_process_mcp_bridge_maps_schema_dispatch_and_errors( + make_toolkit, +) -> None: + toolkit = make_toolkit({"error": "rejected"}, images=[PNG]) + server = _build_rpent_server(toolkit=toolkit) + response = asyncio.run(call_sdk_tool(server, "inspect_scene", {"detail": "high"})) assert toolkit.calls == [("inspect_scene", {"detail": "high"})] - assert response["is_error"] is True - assert [block["type"] for block in response["content"]] == ["text", "image"] - assert response["content"][1]["mimeType"] == "image/png" - + assert response.isError is True + assert [block.type for block in response.content] == ["text", "image"] + assert response.content[1].mimeType == "image/png" -def test_tool_result_conversion_supports_plain_values_and_content_blocks() -> None: - assert _tool_result_to_mcp("plain") == { - "content": [{"type": "text", "text": "plain"}] - } +def test_tool_result_conversion_preserves_original_finish_payload() -> None: result = ToolResult( - "inspect_scene", - {"value": "visible", "_image_bytes": b"pixels"}, + data={ + "value": "visible", + "_finish": True, + "status": "failure", + "summary": "verified", + }, + images=[PNG], ) - converted = _tool_result_to_mcp(result) - - assert converted["content"][0] == { - "type": "text", - "text": '{\n "value": "visible"\n}', + converted = mcp_result(result) + assert json.loads(converted["content"][0]["text"]) == { + "value": "visible", + "_finish": True, + "status": "failure", + "summary": "verified", } assert converted["content"][1]["type"] == "image" - assert "is_error" not in converted + assert converted["isError"] is False def test_successful_fake_sdk_stream_accounts_for_finish_and_hides_image_payload( + make_toolkit, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -253,6 +178,11 @@ async def query(*, prompt: str, options: Any): }, ], } + await call_sdk_tool( + options["mcp_servers"]["rpent"], + "finish", + {"status": "success", "summary": "done"}, + ) yield { "type": "UserMessage", "parent_tool_use_id": "finish-1", @@ -266,12 +196,12 @@ async def query(*, prompt: str, options: Any): result = make_planner(tmp_path, sink).solve( system_prompt="system rules", user_message="user task", - toolkit=FakeToolkit(), + toolkit=make_toolkit(), max_turns=3, ) assert result.finish_result == { - "_finish": True, + "value": "ok", "status": "success", "summary": "done", } @@ -289,8 +219,10 @@ async def query(*, prompt: str, options: Any): assert any(isinstance(event, UsageEvent) for event in sink.events) -def test_rejected_finish_result_is_not_promoted(tmp_path: Path) -> None: - recorder = _Recorder(max_turns=2, dashboard_events=RecordingSink()) +def test_rejected_finish_result_is_not_promoted(make_toolkit, tmp_path: Path) -> None: + toolkit = make_toolkit({"error": "finish refused"}) + arguments = {"status": "success", "summary": "too early"} + recorder = _Recorder(toolkit=toolkit, max_turns=2, dashboard_events=RecordingSink()) recorder.observe( { "type": "AssistantMessage", @@ -301,12 +233,19 @@ def test_rejected_finish_result_is_not_promoted(tmp_path: Path) -> None: "type": "ToolUseBlock", "id": "finish-1", "name": "mcp__rpent__finish", - "input": {"status": "success", "summary": "too early"}, + "input": arguments, } ], } ) + response = asyncio.run( + call_sdk_tool(_build_rpent_server(toolkit=toolkit), "finish", arguments) + ) + assert response.isError + assert toolkit.calls == [("finish", arguments)] + assert toolkit.finish_result is None + rendered = recorder.observe( { "type": "UserMessage", @@ -315,8 +254,8 @@ def test_rejected_finish_result_is_not_promoted(tmp_path: Path) -> None: { "type": "ToolResultBlock", "tool_use_id": "finish-1", - "content": "finish refused", - "is_error": True, + "content": response.content[0].text, + "is_error": response.isError, } ], } @@ -329,6 +268,7 @@ def test_rejected_finish_result_is_not_promoted(tmp_path: Path) -> None: def test_fake_sdk_failure_and_timeout_are_returned( + make_toolkit, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -344,7 +284,7 @@ async def failed_query(*, prompt: str, options: Any): failure = make_planner(tmp_path / "failure", RecordingSink()).solve( system_prompt="", user_message="task", - toolkit=FakeToolkit(), + toolkit=make_toolkit(), max_turns=1, ) @@ -367,7 +307,7 @@ async def slow_query(*, prompt: str, options: Any): ).solve( system_prompt="", user_message="task", - toolkit=FakeToolkit(), + toolkit=make_toolkit(), max_turns=1, ) @@ -377,6 +317,7 @@ async def slow_query(*, prompt: str, options: Any): def test_terminal_timeout_cancels_active_toolkit_work( + make_toolkit, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -394,7 +335,7 @@ async def slow_query(*, prompt: str, options: Any): stream_closed = True monkeypatch.setattr(claude_agent_sdk, "query", slow_query) - toolkit = FakeToolkit() + toolkit = make_toolkit() result = make_planner( tmp_path, RecordingSink(), @@ -408,10 +349,12 @@ async def slow_query(*, prompt: str, options: Any): assert result.error == "Claude Agent SDK timed out after 0.01s" assert stream_closed is True - assert toolkit.cancel_calls == 1 + assert toolkit.cancel_calls >= 1 -def test_stateful_sdk_driver_closes_adapter_tasks_and_client() -> None: +def test_stateful_sdk_driver_closes_adapter_tasks_and_client( + make_toolkit, +) -> None: events: list[str] = [] class Client: @@ -456,7 +399,9 @@ async def close(self) -> None: driver = _ClaudeSessionDriver( sdk=Sdk(), options={"fake": "options"}, - recorder=_Recorder(max_turns=1, dashboard_events=RecordingSink()), + recorder=_Recorder( + toolkit=make_toolkit(), max_turns=1, dashboard_events=RecordingSink() + ), emit=lambda message: events.append(f"message:{message}"), ) @@ -473,6 +418,7 @@ async def close(self) -> None: def test_queue_and_dashboard_are_mutually_exclusive_before_sdk_use( + make_toolkit, tmp_path: Path, ) -> None: planner = make_planner(tmp_path, RecordingSink()) @@ -481,7 +427,7 @@ def test_queue_and_dashboard_are_mutually_exclusive_before_sdk_use( planner.solve( system_prompt="", user_message="task", - toolkit=FakeToolkit(), + toolkit=make_toolkit(), max_turns=1, input_queue=queue.Queue(), dashboard_interaction=object(), diff --git a/tests/unit_tests/rpent/planner/test_codex_contracts.py b/tests/unit_tests/rpent/planner/test_codex_contracts.py index 44221ae75..47a6a628b 100644 --- a/tests/unit_tests/rpent/planner/test_codex_contracts.py +++ b/tests/unit_tests/rpent/planner/test_codex_contracts.py @@ -14,6 +14,7 @@ from __future__ import annotations +import asyncio import json import os import queue @@ -35,47 +36,12 @@ ) from rpent.planner.utils.http_mcp_server import ( HttpMcpServer, - _toolkit_to_mcp_content, + build_mcp_server, + mcp_result, ) -from rpent.tools.toolkit import ToolResult +from rpent.tools import ToolResult - -class RecordingSink: - def __init__(self) -> None: - self.events: list[Any] = [] - - @property - def enabled(self) -> bool: - return True - - def emit(self, event: Any) -> None: - self.events.append(event) - - -class FakeToolkit: - def __init__(self) -> None: - self.cancel_calls = 0 - - def get_tools_spec(self) -> list[dict[str, Any]]: - return [ - { - "name": "finish", - "description": "Finish the task.", - "input_schema": { - "type": "object", - "properties": { - "status": {"type": "string"}, - "summary": {"type": "string"}, - }, - }, - } - ] - - def execute_tool(self, name: str, args: dict[str, Any]) -> ToolResult: - return ToolResult(name, {"name": name, "args": args}) - - def cancel_active_and_wait(self) -> None: - self.cancel_calls += 1 +from ._native_helpers import PNG, FakeToolkit, RecordingSink, call_sdk_tool class FakeMcpServer: @@ -102,7 +68,17 @@ def __init__(self, events: list[dict[str, Any]]) -> None: self.steered: list[str] = [] def stream(self): - yield from self.events + for event in self.events: + if event.get("method") == "item/completed": + item = event["payload"].get("item", {}) + if ( + item.get("tool") == "mcp__rpent__finish" + and item.get("status") == "completed" + ): + FakeMcpServer.instances[-1].toolkit.execute_tool( + "finish", item["arguments"] + ) + yield event def interrupt(self) -> None: self.interrupt_calls += 1 @@ -174,31 +150,23 @@ def codex_config(**kwargs: Any) -> dict[str, Any]: def test_mcp_content_conversion_preserves_text_images_and_error_status() -> None: - plain, plain_error = _toolkit_to_mcp_content("plain") - assert plain_error is False - assert plain[0].type == "text" - assert plain[0].text == "plain" - - result = ToolResult( - "finish", - {"error": "finish refused", "_image_bytes": b"image bytes"}, - ) - content, is_error = _toolkit_to_mcp_content(result) - - assert [block.type for block in content] == ["text", "image"] - assert json.loads(content[0].text) == {"error": "finish refused"} - assert content[1].mimeType == "image/png" - assert is_error is True + result = ToolResult(error="finish refused", images=[PNG]) + converted = mcp_result(result) + assert [block["type"] for block in converted["content"]] == ["text", "image"] + assert json.loads(converted["content"][0]["text"]) == {"error": "finish refused"} + assert converted["content"][1]["mimeType"] == "image/png" + assert converted["isError"] is True def test_http_mcp_readiness_ignores_environment_proxy( + make_toolkit, monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("HTTP_PROXY", "http://127.0.0.1:1") monkeypatch.setenv("HTTPS_PROXY", "http://127.0.0.1:1") monkeypatch.delenv("NO_PROXY", raising=False) monkeypatch.delenv("no_proxy", raising=False) - server = HttpMcpServer(FakeToolkit()) + server = HttpMcpServer(make_toolkit()) try: url = server.start(ready_timeout_s=3.0) @@ -367,7 +335,7 @@ def test_planner_forwards_configured_service_tier( make_planner(tmp_path, RecordingSink()).solve( system_prompt="system rules", user_message="user task", - toolkit=FakeToolkit(), + toolkit=FakeToolkit(tmp_path), max_turns=1, ) @@ -377,6 +345,7 @@ def test_planner_forwards_configured_service_tier( def test_successful_fake_codex_lifecycle_uses_fake_mcp_and_accounts_events( + make_toolkit, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -425,7 +394,7 @@ def test_successful_fake_codex_lifecycle_uses_fake_mcp_and_accounts_events( result = make_planner(tmp_path, sink).solve( system_prompt="system rules", user_message="user task", - toolkit=FakeToolkit(), + toolkit=make_toolkit(), max_turns=3, ) @@ -436,7 +405,7 @@ def test_successful_fake_codex_lifecycle_uses_fake_mcp_and_accounts_events( assert fake_codex.closed is True assert fake_codex.thread.turn_prompts[0][0] == "system rules\n\nuser task" assert result.finish_result == { - "_finish": True, + "value": "ok", "status": "success", "summary": "done", } @@ -454,10 +423,20 @@ def test_successful_fake_codex_lifecycle_uses_fake_mcp_and_accounts_events( assert any(isinstance(event, UsageEvent) for event in sink.events) -def test_rejected_finish_item_is_not_promoted() -> None: +def test_rejected_finish_item_is_not_promoted( + make_toolkit, +) -> None: from rpent.planner.codex import _Recorder - recorder = _Recorder(max_turns=2, dashboard_events=RecordingSink()) + toolkit = make_toolkit({"error": "finish refused"}) + arguments = {"status": "success", "summary": "too early"} + recorder = _Recorder(toolkit=toolkit, max_turns=2, dashboard_events=RecordingSink()) + response = asyncio.run( + call_sdk_tool({"instance": build_mcp_server(toolkit)}, "finish", arguments) + ) + assert response.isError + assert toolkit.calls == [("finish", arguments)] + assert toolkit.finish_result is None rendered = recorder.observe( { @@ -466,9 +445,9 @@ def test_rejected_finish_item_is_not_promoted() -> None: "item": { "type": "mcpToolCall", "tool": "mcp__rpent__finish", - "status": "failed", - "arguments": {"status": "success", "summary": "too early"}, - "error": "finish refused", + "status": "completed", + "arguments": arguments, + "result": response.model_dump(by_alias=True), } }, } @@ -480,6 +459,7 @@ def test_rejected_finish_item_is_not_promoted() -> None: def test_fake_codex_backend_failure_stops_mcp_server( + make_toolkit, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -499,7 +479,7 @@ def __exit__(self, *args: Any) -> None: result = make_planner(tmp_path, RecordingSink()).solve( system_prompt="", user_message="task", - toolkit=FakeToolkit(), + toolkit=make_toolkit(), max_turns=1, ) @@ -508,6 +488,7 @@ def __exit__(self, *args: Any) -> None: def test_timeout_interrupts_without_starting_a_worker_or_socket( + make_toolkit, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -539,7 +520,7 @@ def is_alive(self) -> bool: ).solve( system_prompt="", user_message="task", - toolkit=FakeToolkit(), + toolkit=make_toolkit(), max_turns=1, ) @@ -548,6 +529,7 @@ def is_alive(self) -> bool: def test_terminal_timeout_cancels_active_toolkit_work( + make_toolkit, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -571,7 +553,7 @@ def is_alive(self) -> bool: "threading", SimpleNamespace(Thread=TimeoutThread), ) - toolkit = FakeToolkit() + toolkit = make_toolkit() result = make_planner( tmp_path, RecordingSink(), @@ -585,10 +567,11 @@ def is_alive(self) -> bool: assert result.error == "Codex SDK timed out after 0.01s" assert FakeMcpServer.instances[0].stopped is True - assert toolkit.cancel_calls == 1 + assert toolkit.cancel_calls >= 1 def test_queue_and_dashboard_are_rejected_before_mcp_construction( + make_toolkit, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -603,7 +586,7 @@ def __init__(self, toolkit: Any) -> None: planner.solve( system_prompt="", user_message="task", - toolkit=FakeToolkit(), + toolkit=make_toolkit(), max_turns=1, input_queue=queue.Queue(), dashboard_interaction=object(), diff --git a/tests/unit_tests/rpent/planner/test_http_mcp_server.py b/tests/unit_tests/rpent/planner/test_http_mcp_server.py index c6a5bab3c..62180a05e 100644 --- a/tests/unit_tests/rpent/planner/test_http_mcp_server.py +++ b/tests/unit_tests/rpent/planner/test_http_mcp_server.py @@ -16,160 +16,82 @@ import asyncio import json -import time +import threading from pathlib import Path -from typing import Any import httpx from mcp import ClientSession from mcp.client.streamable_http import streamable_http_client -from rpent.dashboard.events import DashboardEventSink -from rpent.memory.manager import MemoryManager +from rpent.memory import MemoryManager from rpent.planner.utils.http_mcp_server import HttpMcpServer from rpent.session import EnvState -from rpent.tools.toolkit import Toolkit, readonly -from rpent.utils.logging import init_output_dir +from rpent.tools import ToolContext, Toolkit, ToolResult, tool -CONCURRENT_CALLS = [ - ("list_dir", {"path": "resources/libero/memory"}), - ("read_text_file", {"path": "robots/libero/guides/strict_hybrid_guide.md"}), - ("read_text_file", {"path": "robots/libero/guides/pro_hybrid_guide.md"}), - ("read_text_file", {"path": "robots/libero/guides/env_calibration.md"}), - ("view_env_state", {"step": 0}), -] +from ._native_helpers import read -class RecordingSink(DashboardEventSink): - def __init__(self) -> None: - self.events: list[Any] = [] +@tool +def finish(status: str, summary: str, *, ctx: ToolContext) -> ToolResult: + """Accept the requested outcome for this test toolkit.""" + return ToolResult(data={"_finish": True, "status": status, "summary": summary}) - @property - def enabled(self) -> bool: - return True - def emit(self, event: Any) -> None: - self.events.append(event) - - -class FakeToolkit(Toolkit): - """Minimal toolkit whose tools sleep to widen the overlap window.""" - - def __init__(self, state_dir: Path) -> None: - super().__init__( - dashboard_events=RecordingSink(), - state=EnvState(state_dir), - memory=MemoryManager(state_dir / "memory"), - ) - self.overlap_errors: list[tuple[str, dict[str, Any]]] = [] - self._register_fake_tools() - - def _register_fake_tools(self) -> None: - @readonly - def read_text_file(path: str, max_chars: int = 40000) -> dict: - time.sleep(0.05) - p = Path(path) - return {"path": str(p), "size": 0, "content": "fake content"} - - @readonly - def list_dir(path: str = "") -> dict: - time.sleep(0.05) - return {"path": path, "count": 0, "files": []} - - def view_env_state(step: int = -1) -> dict: - time.sleep(0.3) - return {"step": step, "mode": "evaluation"} - - self.add_tool( - "read_text_file", - { - "name": "read_text_file", - "description": "Read a UTF-8 text file.", - "input_schema": { - "type": "object", - "properties": {"path": {"type": "string"}}, - "required": ["path"], - }, - }, - read_text_file, - ) - self.add_tool( - "list_dir", - { - "name": "list_dir", - "description": "List files in a directory.", - "input_schema": { - "type": "object", - "properties": {"path": {"type": "string"}}, - }, - }, - list_dir, - ) - self.add_tool( - "view_env_state", - { - "name": "view_env_state", - "description": "View the current environment state.", - "input_schema": { - "type": "object", - "properties": {"step": {"type": "integer"}}, - }, - }, - view_env_state, - ) - - def execute_tool(self, name: str, input_dict: dict[str, Any]) -> Any: - result = super().execute_tool(name, input_dict) - if result.result.get("error") == "another tool operation is still active": - self.overlap_errors.append((name, dict(input_dict))) - return result - - def get_env_state( - self, - *, - command: dict[str, Any], - result: dict[str, Any], - elapsed_s: float, - ) -> dict[str, Any]: - return {"observed": True} - - def solved(self) -> bool: - return False - - -def test_http_mcp_server_serializes_concurrent_tool_calls(tmp_path: Path) -> None: - init_output_dir(tmp_path / "log") - toolkit = FakeToolkit(tmp_path) +def test_http_serialized_calls_keep_native_validation(tmp_path: Path) -> None: + toolkit = Toolkit( + state=EnvState(tmp_path), + memory=MemoryManager(tmp_path / "memory"), + robot=threading.Barrier(1), + output_dir=tmp_path, + tools=(finish, read), + ) server = HttpMcpServer(toolkit) - try: - url = server.start() - rejected = asyncio.run(_fire_concurrent(url)) - finally: - server.stop() - assert rejected == 0 - assert toolkit.overlap_errors == [] - - -async def _fire_concurrent(url: str) -> int: - rejected = 0 - # trust_env=False keeps the SDK client on loopback even when the host - # advertises an HTTP proxy (e.g. macOS system settings), matching the - # production readiness probe in http_mcp_server._wait_for_ready. - async with httpx.AsyncClient(trust_env=False) as http_client: - async with streamable_http_client(url, http_client=http_client) as ( - read, - write, - _get_session_id, + async def scenario(url): + async with ( + httpx.AsyncClient(trust_env=False) as client, + streamable_http_client(url, http_client=client) as (reader, writer, _), ): - async with ClientSession(read, write) as session: + async with ClientSession(reader, writer) as session: await session.initialize() + specs = await session.list_tools() + assert "read_image" not in {tool.name for tool in specs.tools} + for name in ("read_image", "mcp__rpent__read_image"): + rejected = await session.call_tool(name, {"name": "frame.png"}) + assert rejected.isError + assert json.loads(rejected.content[0].text) == { + "error": "Unknown tool: read_image" + } + assert ( + next(t for t in specs.tools if t.name == "read").inputSchema + == read.input_schema + ) results = await asyncio.gather( - *(session.call_tool(name, args) for name, args in CONCURRENT_CALLS) + session.call_tool("read", {"number": "1"}), + session.call_tool("read", {"number": 2}), ) - for result in results: - if "another tool operation is still active" in json.dumps( - result.content, default=str - ): - rejected += 1 - return rejected + assert all(not result.isError for result in results) + assert [json.loads(r.content[0].text)["number"] for r in results] == [ + 1, + 2, + ] + invalid = await session.call_tool("read", {"number": "bad"}) + assert invalid.isError + failure = json.loads(invalid.content[0].text) + assert set(failure) == {"error"} + assert failure["error"].startswith("Invalid arguments for read.") + assert "number" in failure["error"] + assert "valid integer" in failure["error"] + unknown = await session.call_tool("read", {"numbr": 1}) + assert unknown.isError + assert "extra_forbidden" in unknown.content[0].text + results = await asyncio.gather( + *[session.call_tool("list_dir", {}) for _ in range(4)] + ) + assert all(not r.isError for r in results) + + try: + asyncio.run(scenario(server.start())) + finally: + server.stop() + toolkit.close() diff --git a/tests/unit_tests/rpent/planner/test_native_adapters.py b/tests/unit_tests/rpent/planner/test_native_adapters.py new file mode 100644 index 000000000..03208d3cf --- /dev/null +++ b/tests/unit_tests/rpent/planner/test_native_adapters.py @@ -0,0 +1,116 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +import json + +import pytest +from mcp import types + +from rpent.dashboard.events import NullDashboardEventSink +from rpent.planner.api_loop import ( + _ApiRunObserver, + _build_tools, +) +from rpent.planner.claude_code import ( + _build_rpent_server, +) +from rpent.planner.claude_code import _Recorder as ClaudeRecorder +from rpent.planner.codex import _Recorder as CodexRecorder + +from ._native_helpers import call_sdk_tool + + +def test_read_image_is_available_to_api_and_not_callable_through_mcp(make_toolkit): + toolkit = make_toolkit() + assert "read_image" in {definition.name for definition in _build_tools(toolkit)} + + async def scenario(): + server = _build_rpent_server(toolkit=toolkit) + specs = await server["instance"].request_handlers[types.ListToolsRequest]( + types.ListToolsRequest(method="tools/list") + ) + assert "read_image" not in {definition.name for definition in specs.root.tools} + rejected = await call_sdk_tool(server, "read_image", {"name": "frame.png"}) + assert rejected.isError + assert json.loads(rejected.content[0].text) == { + "error": "Unknown tool: read_image" + } + assert toolkit.calls == [] + + asyncio.run(scenario()) + + +@pytest.mark.parametrize( + "recorder_type", [ClaudeRecorder, CodexRecorder, _ApiRunObserver] +) +def test_recorders_read_verified_finish_without_a_matching_provider_event( + make_toolkit, + recorder_type, +): + accepted = { + "status": "failure", + "summary": "environment verification failed", + "operator_aborted": True, + "operator_finished": True, + "operator_verdict": "abort", + "operator_notes": "Stop the run.", + "attempt": 2, + } + toolkit = make_toolkit({}, accepted=accepted) + recorder = recorder_type( + toolkit=toolkit, + max_turns=3, + dashboard_events=NullDashboardEventSink(), + **({"messages": []} if recorder_type is _ApiRunObserver else {}), + ) + assert recorder.finish_result is None + result = toolkit.execute_tool( + "finish", {"status": "success", "summary": "model claims success"} + ) + assert not result.is_error + assert result.data == {"_finish": True, **accepted} + assert recorder.finish_result == accepted + assert recorder.finish_result == toolkit.finish_result + returned = recorder.finish_result + returned["operator_aborted"] = False + assert toolkit.finish_result["operator_aborted"] is True + + +def test_api_finish_accepted_during_interrupt_seals_dashboard(): + from types import SimpleNamespace + + from rpent.planner.api_loop import _ApiDashboardSession + + async def scenario(): + ended = asyncio.Event() + control = SimpleNamespace(end=ended.set) + session = _ApiDashboardSession( + agent=None, + control=control, + observer=SimpleNamespace( + finish_result={"status": "failure", "summary": "verified"} + ), + max_turns=3, + no_images=False, + ) + session._active_prompt = True + session._run_task = asyncio.create_task(asyncio.Event().wait()) + assert await session.interrupt() == 1 + assert ended.is_set() + assert session._run_task is None + + asyncio.run(scenario()) diff --git a/tests/unit_tests/rpent/tools/test_common_tools.py b/tests/unit_tests/rpent/tools/test_common_tools.py new file mode 100644 index 000000000..dfe31552d --- /dev/null +++ b/tests/unit_tests/rpent/tools/test_common_tools.py @@ -0,0 +1,211 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import numpy as np +import pytest + +from rpent.memory import MemoryManager +from rpent.session import EnvState +from rpent.tools import ToolContext, Toolkit, ToolResult, tool + + +@tool +def finish(status: str, summary: str, *, ctx: ToolContext) -> ToolResult: + """Accept the requested outcome for this test toolkit.""" + return ToolResult(data={"_finish": True, "status": status, "summary": summary}) + + +@pytest.fixture +def toolkit(tmp_path, monkeypatch): + monkeypatch.setenv("RPENT_REPO_ROOT", str(tmp_path)) + instance = Toolkit( + state=EnvState(tmp_path / "sessions" / "session_001"), + memory=MemoryManager(tmp_path / "memory" / "libero"), + robot=None, + output_dir=tmp_path, + tools=(finish,), + ) + yield instance + instance.close() + + +def test_read_text_file_and_list_dir_use_native_results(toolkit, tmp_path): + (tmp_path / "note.txt").write_text("tool reads", encoding="utf-8") + text_result = toolkit.execute_tool("read_text_file", {"path": "note.txt"}) + directory_result = toolkit.execute_tool("list_dir", {}) + assert not text_result.is_error + assert not directory_result.is_error + assert text_result.data["content"] == "tool reads" + assert "note.txt" in directory_result.data["files"] + + +def test_file_contracts_preserve_truncation_overwrite_and_utf8_counts( + toolkit, tmp_path +): + written = toolkit.execute_tool( + "write_text_file", {"path": "files/note.txt", "content": "hello 世界"} + ) + assert written.data == { + "path": str(tmp_path / "files/note.txt"), + "bytes_written": 12, + } + read = toolkit.execute_tool( + "read_text_file", {"path": "files/note.txt", "max_chars": "7"} + ) + assert read.data["size"] == 8 + assert ( + read.data["content"] + == "hello 世\n\n[TRUNCATED — file is 8 chars, showed first 7]" + ) + toolkit.execute_tool("write_text_file", {"path": "files/note.txt", "content": "x"}) + assert (tmp_path / "files/note.txt").read_text() == "x" + (tmp_path / "files/a").mkdir() + listed = toolkit.execute_tool("list_dir", {"path": "files"}) + assert listed.data == { + "path": str(tmp_path / "files"), + "count": 2, + "files": ["a", "note.txt"], + } + assert toolkit.state.latest_step is None + + +def test_default_directory_is_bound_to_each_task_not_observation_or_global_output( + toolkit, tmp_path, monkeypatch +): + other_dir = tmp_path / "other" + monkeypatch.chdir(tmp_path) + other = Toolkit( + state=EnvState(other_dir / "state"), + memory=toolkit.memory, + robot=None, + output_dir="other", + tools=(finish,), + ) + monkeypatch.setattr( + "rpent.utils.logging.get_output_dir", lambda: tmp_path / "unrelated" + ) + try: + monkeypatch.chdir(other_dir) + supplied = {"ctx": {"output_dir": str(other_dir)}, "output_dir": str(other_dir)} + assert toolkit.execute_tool("list_dir", supplied).is_error + assert toolkit.execute_tool("list_dir", {}).data["path"] == str(tmp_path) + assert other.execute_tool("list_dir", {}).data["path"] == str(other_dir) + finally: + other.close() + + +@pytest.mark.parametrize( + "name,args", + [ + ("read_text_file", {"path": "missing"}), + ("write_text_file", {"path": "sessions", "content": "x"}), + ("list_dir", {"path": "missing"}), + ], +) +def test_file_failures_are_explicit_without_capture(toolkit, name, args): + result = toolkit.execute_tool(name, args) + assert result.is_error + assert "step" not in result.data + assert not toolkit.execute_tool("list_dir", {}).is_error + + +def test_png_bytes_preserve_original_artifact_and_step(toolkit): + image = np.full((3, 4, 3), 128, dtype=np.uint8) + name = "frame.png" + with toolkit.state.record_step(state={}): + toolkit.state.save(name, image) + path = toolkit.state.artifact_path(name) + before = path.read_bytes() + result = toolkit.execute_tool("read_image", {"name": name}) + assert not result.is_error + assert result.data == {"artifact": name, "step": 0} + assert result.images == [before] + assert path.read_bytes() == before + assert toolkit.state.latest_step == 0 + + +def test_read_image_with_relative_output_dir(toolkit, tmp_path, monkeypatch): + working_dir = tmp_path / "outside" + working_dir.mkdir() + monkeypatch.chdir(working_dir) + state = EnvState("run") + relative_toolkit = Toolkit( + state=state, + memory=toolkit.memory, + robot=None, + output_dir="run", + tools=(), + ) + with state.record_step(state={}): + state.save("frame.png", np.zeros((2, 2, 3), dtype=np.uint8)) + expected = (working_dir / "run/frame.png/00.png").read_bytes() + for cwd in (working_dir, tmp_path): + monkeypatch.chdir(cwd) + result = relative_toolkit.execute_tool("read_image", {"name": "frame.png"}) + assert not result.is_error + assert result.images == [expected] + + +@pytest.mark.parametrize("suffix", ["jpg", "jpeg"]) +def test_read_image_rejects_jpeg_artifacts(toolkit, suffix): + name = f"frame.{suffix}" + with toolkit.state.record_step(state={}): + toolkit.state.save(name, np.zeros((2, 2, 3), dtype=np.uint8)) + result = toolkit.execute_tool("read_image", {"name": name}) + assert result.is_error + assert "PNG" in result.error + assert result.images == [] + + +@pytest.mark.parametrize( + "name,step", + [ + ("frame.png", -1), + ("frame.png", 42), + ("../outside.png", 0), + ("file.txt", 0), + ("missing.png", 0), + ], +) +def test_missing_invalid_or_nonimage_artifacts_report_errors(toolkit, name, step): + with toolkit.state.record_step(state={}): + toolkit.state.save("file.txt", "text") + result = toolkit.execute_tool("read_image", {"name": name, "step": step}) + assert result.is_error + assert result.images == [] + assert toolkit.finish_result is None + + +def test_file_and_image_tools_apply_memory_permissions(toolkit, tmp_path): + published = tmp_path / "memory/libero/global/note.md" + published.parent.mkdir(parents=True) + published.write_text("published") + assert ( + toolkit.execute_tool("read_text_file", {"path": str(published)}).data["content"] + == "published" + ) + denied = toolkit.execute_tool( + "write_text_file", {"path": str(published), "content": "changed"} + ) + assert denied.is_error + foreign = tmp_path / "memory/robotwin/global/frame.png" + foreign.parent.mkdir(parents=True) + with toolkit.state.record_step(state={}): + toolkit.state.save("frame.png", np.zeros((1, 1, 3), dtype=np.uint8)) + path = toolkit.state.artifact_path("frame.png") + foreign.write_bytes(path.read_bytes()) + path.unlink() + path.symlink_to(foreign) + assert toolkit.execute_tool("read_image", {"name": "frame.png"}).is_error diff --git a/tests/unit_tests/rpent/tools/test_human_in_the_loop.py b/tests/unit_tests/rpent/tools/test_human_in_the_loop.py index 28e5eb2f2..45991cf3c 100644 --- a/tests/unit_tests/rpent/tools/test_human_in_the_loop.py +++ b/tests/unit_tests/rpent/tools/test_human_in_the_loop.py @@ -18,8 +18,8 @@ import pytest +from rpent.tools import ToolCancelled from rpent.tools.human_in_the_loop import HumanInTheLoopInput -from rpent.tools.toolkit import ToolCancelled def test_operator_replies_are_request_scoped_and_do_not_consume_steering(monkeypatch): diff --git a/tests/unit_tests/rpent/tools/test_native_protocol.py b/tests/unit_tests/rpent/tools/test_native_protocol.py new file mode 100644 index 000000000..7d9311dc9 --- /dev/null +++ b/tests/unit_tests/rpent/tools/test_native_protocol.py @@ -0,0 +1,209 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU-only examples and checks for signature-derived native tools.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Annotated, Any, Literal + +import pytest +from pydantic import Field + +from rpent.tools import ( + Tool, + ToolContext, + ToolResult, + readonly, + tool, +) +from rpent.tools.base import MAX_TOOL_TEXT_BYTES + +if TYPE_CHECKING: + from robots.libero.toolkit import LiberoRuntime + + +@tool +def move( + xyz: Annotated[list[float], Field(min_length=3, max_length=3)], + substeps: Annotated[int, Field(ge=0)] = 25, + *, + ctx: ToolContext, +) -> ToolResult: + """Move to a world-frame position in metres. + + Args: + xyz: World metres. + substeps: Number of steps. + """ + return ToolResult(data={"xyz": xyz, "substeps": substeps}) + + +def test_signature_and_docstring_define_schema_and_validation() -> None: + assert isinstance(move, Tool) + assert move.name == "move" + assert move.description == "Move to a world-frame position in metres." + assert not move.readonly + assert move.input_schema == { + "type": "object", + "additionalProperties": False, + "properties": { + "xyz": { + "type": "array", + "items": {"type": "number"}, + "minItems": 3, + "maxItems": 3, + "description": "World metres.", + }, + "substeps": { + "type": "integer", + "minimum": 0, + "description": "Number of steps.", + }, + }, + "required": ["xyz"], + } + assert move.args_schema.model_validate({"xyz": [0, 0, 0]}).substeps == 25 + assert "ctx" not in move.input_schema["properties"] + + +def test_export_preserves_nullable_constraints_and_explicit_defaults() -> None: + @tool + @readonly + def exported( + vector: Annotated[list[float] | None, Field(min_length=3, max_length=3)] = None, + choice: Literal["auto"] | None = None, + step: Annotated[int, Field(json_schema_extra={"default": -1})] = -1, + *, + ctx: ToolContext, + ) -> ToolResult: + """Export constrained nullable parameters.""" + return ToolResult() + + schema = exported.input_schema + assert "title" not in schema + assert "description" not in schema + properties = schema["properties"] + assert properties["vector"] == { + "type": ["array", "null"], + "items": {"type": "number"}, + "minItems": 3, + "maxItems": 3, + } + assert properties["step"] == {"type": "integer", "default": -1} + + assert properties["choice"] == { + "anyOf": [{"const": "auto", "type": "string"}, {"type": "null"}] + } + assert exported.args_schema.model_validate({}).step == -1 + + +def test_dictionary_schema_preserves_value_constraints() -> None: + @tool + def exported( + free: dict[str, Any], typed: dict[str, int], *, ctx: ToolContext + ) -> ToolResult: + """Accept free-form data and an integer mapping.""" + return ToolResult() + + properties = exported.input_schema["properties"] + assert properties["free"] == {"type": "object"} + assert properties["typed"] == { + "type": "object", + "additionalProperties": {"type": "integer"}, + } + validated = exported.args_schema.model_validate( + {"free": {"anything": [None, "value"]}, "typed": {"count": 3}} + ) + assert validated.free == {"anything": [None, "value"]} + assert validated.typed == {"count": 3} + + +@pytest.mark.parametrize( + ("markers", "is_readonly"), + [ + ([], False), + ([readonly], True), + ], +) +def test_readonly_marker_preserves_handler(markers, is_readonly) -> None: + def handler( + limit: Annotated[int, Field(ge=1)] = 10, *, ctx: ToolContext + ) -> ToolResult: + """Return the requested limit. + + Args: + limit: Maximum number of results. + """ + return ToolResult(data={"limit": limit}) + + original = tool(handler) + for marker in markers: + handler = marker(handler) + definition = tool(handler) + assert definition.handler is original.handler + assert definition.readonly is is_readonly + assert definition.description == original.description + assert definition.input_schema == original.input_schema + assert definition.args_schema.model_validate({}).limit == 10 + assert definition.handler(3, ctx=None).data == {"limit": 3} + + +def test_injected_and_return_annotations_need_not_resolve() -> None: + @tool + @readonly + def handler(value: int, *, ctx: ToolContext[LiberoRuntime]) -> LiberoRuntime: + """Context and return types are not model inputs. + + Args: + value: A model-supplied integer. + ctx: Internal resources, hidden from the model. + """ + return ToolResult(data={"value": value}) + + assert handler.handler(3, ctx=None).data == {"value": 3} + assert set(handler.input_schema["properties"]) == {"value"} + + +@pytest.mark.parametrize("error", ["", "Failed."]) +def test_explicit_error_serializes_without_mutating_data(error): + result = ToolResult(data={"step": 2}, error=error) + assert result.is_error + assert json.loads(result.to_text()) == {"step": 2, "error": error} + assert result.data == {"step": 2} + + +def test_business_error_key_does_not_mark_tool_failure(): + result = ToolResult(data={"error": {"count": 0}}) + assert not result.is_error + assert json.loads(result.to_text()) == result.data + + +@pytest.mark.parametrize( + ("summary", "truncated"), + [("任务未完成", False), ("x" * 60000, True), ("任务未完成" * 12000, True)], +) +def test_text_size_limit_preserves_original_result(summary, truncated): + data = {"status": "failure", "summary": summary} + result = ToolResult(data=dict(data)) + + text = result.to_text() + + assert len(text.encode("utf-8")) <= MAX_TOOL_TEXT_BYTES + if truncated: + assert text.endswith("[truncated]") + else: + assert json.loads(text) == data + assert result.data == data diff --git a/tests/unit_tests/rpent/tools/test_scheduling.py b/tests/unit_tests/rpent/tools/test_scheduling.py new file mode 100644 index 000000000..0795c35ea --- /dev/null +++ b/tests/unit_tests/rpent/tools/test_scheduling.py @@ -0,0 +1,61 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import threading +from concurrent.futures import ThreadPoolExecutor + +from rpent.memory import MemoryManager +from rpent.session import EnvState +from rpent.tools import ToolContext, Toolkit, ToolResult, readonly, tool + + +def test_overlapping_calls_are_rejected_and_cancel_waits_for_active_call(tmp_path): + entered = threading.Event() + release = threading.Event() + contexts = [] + + @tool + @readonly + def read(*, ctx: ToolContext) -> ToolResult: + """Read until the current operation reaches its cancellation boundary.""" + contexts.append(ctx) + entered.set() + assert release.wait(5) + ctx.check_cancelled() + return ToolResult(data={"ok": True}) + + toolkit = Toolkit( + state=EnvState(tmp_path), + memory=MemoryManager(tmp_path / "memory"), + robot=None, + output_dir=tmp_path, + tools=(read,), + ) + with ThreadPoolExecutor(max_workers=2) as pool: + active = pool.submit(toolkit.execute_tool, "read", {}) + try: + assert entered.wait(5) + rejected = toolkit.execute_tool("read", {}) + assert rejected.error == "another tool operation is still active" + assert len(contexts) == 1 + cancelled = pool.submit(toolkit.cancel_active_and_wait) + assert contexts[0]._cancel_event.wait(5) + assert not cancelled.done() + finally: + release.set() + assert active.result(5).error == "Tool call cancelled." + cancelled.result(5) + assert not toolkit.execute_tool("read", {}).is_error + assert len(contexts) == 2 + assert not contexts[1]._cancel_event.is_set() diff --git a/tests/unit_tests/rpent/tools/test_toolkit_contracts.py b/tests/unit_tests/rpent/tools/test_toolkit_contracts.py index 3d0a34d00..5287b5c65 100644 --- a/tests/unit_tests/rpent/tools/test_toolkit_contracts.py +++ b/tests/unit_tests/rpent/tools/test_toolkit_contracts.py @@ -14,542 +14,221 @@ from __future__ import annotations -import base64 -import copy import json -import threading -from functools import partial -from pathlib import Path -from typing import Any +from types import SimpleNamespace import pytest from rpent.dashboard.events import StepRecordEvent from rpent.memory import MemoryManager -from rpent.memory import tools as memory_tools from rpent.session import EnvState -from rpent.tools import common -from rpent.tools.toolkit import Toolkit, ToolResult, readonly +from rpent.tools import ToolCancelled, ToolContext, Toolkit, ToolResult, readonly, tool -class _RecordingEventSink: - def __init__(self) -> None: - self.events: list[Any] = [] +@tool +def finish(status: str, summary: str, *, ctx: ToolContext) -> ToolResult: + """Accept the requested outcome for this test toolkit.""" + return ToolResult(data={"_finish": True, "status": status, "summary": summary}) - @property - def enabled(self) -> bool: - return True - def emit(self, event: Any) -> None: - self.events.append(event) - - -class _ContractToolkit(Toolkit): - def __init__( - self, - output_dir: Path, - *, - memory: MemoryManager | None = None, - ) -> None: - self.events = _RecordingEventSink() - self.capture_calls: list[dict[str, Any]] = [] - self.capture_error: Exception | None = None - super().__init__( - dashboard_events=self.events, - state=EnvState(output_dir), - memory=memory or MemoryManager(output_dir / "memory"), - ) - - def get_env_state( - self, - *, - command: dict[str, Any], - result: dict[str, Any], - elapsed_s: float, - ) -> dict[str, Any]: - if self.capture_error is not None: - raise self.capture_error - call = { - "command": copy.deepcopy(command), - "result": copy.deepcopy(result), - "elapsed_s": elapsed_s, - } - self.capture_calls.append(call) - with self.state.record_step( - state={"capture_count": len(self.capture_calls)}, - command=command, - result=result, - elapsed_s=elapsed_s, - ): - pass - return {"observation": len(self.capture_calls)} - - def solved(self) -> bool: - return False - - -def test_tool_result_builds_text_and_images_without_mutating_result() -> None: - image_payloads = { - "_image_bytes": b"main", - "_image_cam_bytes": b"camera", - "_image_nav_bytes": b"navigation", - "_image_wrist_bytes": b"wrist", - } - result = {"status": "ok", "count": 2, **image_payloads} - original = copy.deepcopy(result) - - tool_result = ToolResult(name="observe", result=result, call_id="call-1") - - assert result == original - assert tool_result.call_id == "call-1" - assert tool_result.is_finish is False - assert json.loads(tool_result.content_blocks[0]["text"]) == { - "status": "ok", - "count": 2, - } - assert [block["type"] for block in tool_result.content_blocks] == [ - "text", - "image", - "image", - "image", - "image", - ] - assert [ - base64.b64decode(block["source"]["data"]) - for block in tool_result.content_blocks[1:] - ] == list(image_payloads.values()) - assert all( - block["source"]["media_type"] == "image/png" - for block in tool_result.content_blocks[1:] - ) +class EventSink: + enabled = True + def __init__(self): + self.events = [] -@pytest.mark.parametrize( - ("raw_result", "expected_text"), - [ - ("plain text", "plain text"), - (17, "17"), - (["one", "two"], "['one', 'two']"), - ], -) -def test_tool_result_converts_ordinary_results_to_text( - raw_result: Any, - expected_text: str, -) -> None: - tool_result = ToolResult(name="ordinary", result=raw_result) - - assert tool_result.content_blocks == [{"type": "text", "text": expected_text}] - - -def test_tool_result_recognizes_finish_only_from_truthy_dict_sentinel() -> None: - assert ToolResult("finish", {"_finish": True}).is_finish is True - assert ToolResult("finish", {"_finish": False}).is_finish is False - assert ToolResult("finish", "finished").is_finish is False + def emit(self, event): + self.events.append(event) -@pytest.mark.parametrize( - ("raw_result", "expected_plain_text"), - [ - pytest.param( - "界" * 10, - "界" * 6, - id="ordinary-unicode-text", - ), - pytest.param( - {"value": "界" * 10}, - None, - id="unicode-json-dict", - ), - ], -) -def test_tool_result_text_limit_counts_utf8_bytes( - monkeypatch: pytest.MonkeyPatch, - raw_result: Any, - expected_plain_text: str | None, -) -> None: - monkeypatch.setattr(ToolResult, "MAX_TEXT_BYTES_IN_RESULT", 20) - - text = ToolResult(name="unicode", result=raw_result).content_blocks[0]["text"] - - assert len(text.encode("utf-8")) <= 20 - assert text.encode("utf-8").decode("utf-8") == text - if expected_plain_text is not None: - assert text == expected_plain_text - - -def test_toolkit_registers_common_specs_with_fresh_placeholder_substitution( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - output_dir = tmp_path / "run-output" - original_specs = copy.deepcopy(common.TOOLS_SPEC) - monkeypatch.setattr("rpent.utils.templates.get_output_dir", lambda: output_dir) - toolkit = _ContractToolkit(tmp_path / "state") - - first = toolkit.get_tools_spec() - second = toolkit.get_tools_spec() - - assert [spec["name"] for spec in first] == [ - "read_text_file", - "write_text_file", - "list_dir", - "finish", - ] - list_dir_spec = next(spec for spec in first if spec["name"] == "list_dir") - assert str(output_dir) in list_dir_spec["description"] - assert memory_tools.MEMORY_BOUNDARY_NOTE in list_dir_spec["description"] - assert ( - str(output_dir) - in list_dir_spec["input_schema"]["properties"]["path"]["description"] - ) - assert common.TOOLS_SPEC == original_specs - assert first == second - assert first is not second - assert first[0] is not common.TOOLS_SPEC[0] +@tool +def action(value: int = 1, *, ctx: ToolContext) -> ToolResult: + """Execute an action with the supplied value.""" + ctx.robot.received.append((value, ctx)) + if ctx.robot.failure: + raise ctx.robot.failure + return ctx.robot.result -def test_common_file_tools_dispatch_offline_without_capturing_robot_state( - tmp_path: Path, -) -> None: - toolkit = _ContractToolkit(tmp_path / "state") - text_file = tmp_path / "files" / "note.txt" +@tool +@readonly +def inspect_state(*, ctx: ToolContext) -> ToolResult: + """Return data without capturing another observation.""" + return ToolResult(data={"ready": True}) - written = toolkit.execute_tool( - "write_text_file", - {"path": str(text_file), "content": "hello 世界"}, - ) - read = toolkit.execute_tool( - "read_text_file", - {"path": str(text_file), "max_chars": 7}, - ) - listed = toolkit.execute_tool("list_dir", {"path": str(text_file.parent)}) - finished = toolkit.execute_tool( - "finish", - {"status": "success", "summary": "done"}, - ) - assert written.result == { - "path": str(text_file), - "bytes_written": len("hello 世界".encode()), - } - assert read.result["path"] == str(text_file) - assert read.result["size"] == len("hello 世界") - assert read.result["content"].startswith("hello 世") - assert "[TRUNCATED" in read.result["content"] - assert listed.result == { - "path": str(text_file.parent), - "count": 1, - "files": ["note.txt"], - } - assert finished.result == { - "_finish": True, - "status": "success", - "summary": "done", - } - assert finished.is_finish is True - assert toolkit.capture_calls == [] - assert toolkit.events.events == [] - - -def test_common_file_tools_enforce_memory_manager_boundaries( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - repo_root = tmp_path / "repo" - memory_root = repo_root / "memory" / "libero" - published = memory_root / "global" / "strategy.md" - published.parent.mkdir(parents=True) - published.write_text("published") - (memory_root / "MEMORY.md").write_text("index") - root_leaf = memory_root / "notes.md" - root_leaf.write_text("root-level note") - evaluation_inbox = memory_root / "_internal" / "inbox" / "current-cell" - evaluation_inbox.mkdir(parents=True) - (evaluation_inbox / "draft.md").write_text("private draft") - foreign = repo_root / "memory" / "robotwin" / "global" / "x.md" - foreign.parent.mkdir(parents=True) - foreign.write_text("foreign") - monkeypatch.setattr(memory_tools, "get_repo_root", lambda: repo_root) - monkeypatch.setattr(common, "get_repo_root", lambda: repo_root) - - read_only_memory = MemoryManager(memory_root) - evaluation = _ContractToolkit( - tmp_path / "evaluation-state", - memory=read_only_memory, - ) - read_published = evaluation.execute_tool( - "read_text_file", - {"path": "memory/libero/global/strategy.md"}, - ) - read_root_leaf = evaluation.execute_tool( - "read_text_file", - {"path": str(root_leaf)}, - ) - list_published = evaluation.execute_tool( - "list_dir", - {"path": str(published.parent)}, - ) - write_published = evaluation.execute_tool( - "write_text_file", - {"path": str(published), "content": "changed"}, - ) - read_foreign = evaluation.execute_tool( - "read_text_file", - {"path": str(foreign)}, - ) - read_evaluation_inbox = evaluation.execute_tool( - "read_text_file", - {"path": str(evaluation_inbox / "draft.md")}, - ) - - assert evaluation.memory is read_only_memory - assert read_published.result["content"] == "published" - assert read_root_leaf.result["content"] == "root-level note" - assert list_published.result["files"] == ["strategy.md"] - assert "writing to memory is denied" in write_published.result["error"] - assert "another robot's memory is denied" in read_foreign.result["error"] - assert "reading this memory path is denied" in read_evaluation_inbox.result["error"] - - exploration = _ContractToolkit( - tmp_path / "exploration-state", - memory=MemoryManager( - memory_root, - memory_access="inbox_write", - inbox_cell_tag="current-cell", +class RecordingToolkit(Toolkit): + def _capture_observation(self, *, command, result, elapsed_s): + with self.state.record_step( + state={"position": 3}, + command=command, + result=result.to_dict(), + elapsed_s=elapsed_s, + ) as step: + return {"step": step, "position": 3}, [b"observation"] + + +@pytest.fixture +def toolkit(tmp_path): + instance = RecordingToolkit( + state=EnvState(tmp_path / "states"), + memory=MemoryManager(tmp_path / "memory"), + robot=SimpleNamespace( + received=[], + failure=None, + result=ToolResult(data={"moved": True}, images=[b"action"]), ), - ) - own_draft = memory_root / "_internal" / "inbox" / "current-cell" / "draft.md" - other_draft = memory_root / "_internal" / "inbox" / "other-cell" / "draft.md" - write_own = exploration.execute_tool( - "write_text_file", - {"path": str(own_draft), "content": "draft"}, - ) - read_own = exploration.execute_tool( - "read_text_file", - {"path": str(own_draft)}, - ) - read_other = exploration.execute_tool( - "read_text_file", - {"path": str(other_draft)}, - ) - inbox_escape = own_draft.parent / "published-link.md" - inbox_escape.symlink_to(published) - write_through_symlink = exploration.execute_tool( - "write_text_file", - {"path": str(inbox_escape), "content": "escaped"}, - ) - - assert write_own.result["bytes_written"] == 5 - assert read_own.result["content"] == "draft" - assert "reading this memory path is denied" in read_other.result["error"] - assert "writing to memory is denied" in write_through_symlink.result["error"] - assert published.read_text() == "published" - assert evaluation.capture_calls == [] - assert exploration.capture_calls == [] - - -def test_toolkit_reports_unknown_tools_and_invalid_arguments(tmp_path: Path) -> None: - toolkit = _ContractToolkit(tmp_path) - - unknown = toolkit.execute_tool("missing", {"value": 1}) - invalid = toolkit.execute_tool("read_text_file", {"unexpected": True}) - - assert unknown.result == {"error": "unknown tool: missing"} - assert "bad arguments for read_text_file" in invalid.result["error"] - assert invalid.result["got"] == {"unexpected": True} - assert toolkit.capture_calls == [] - - -def test_readonly_marker_handles_functions_bound_methods_and_nested_partials( - tmp_path: Path, -) -> None: - toolkit = _ContractToolkit(tmp_path) - - @readonly - def readonly_function(value: str) -> dict[str, str]: - return {"value": value} - - class Handler: - @readonly - def readonly_method(self, *, prefix: str, value: str) -> dict[str, str]: - return {"value": prefix + value} - - handler = Handler() - toolkit.add_tool("function", {"name": "function"}, readonly_function) - toolkit.add_tool("method", {"name": "method"}, handler.readonly_method) - toolkit.add_tool( - "partial", - {"name": "partial"}, - partial(partial(handler.readonly_method, prefix="pre-"), value="bound"), - ) - - assert toolkit.execute_tool("function", {"value": "plain"}).result == { - "value": "plain" - } - assert toolkit.execute_tool( - "method", {"prefix": "pre-", "value": "bound"} - ).result == {"value": "pre-bound"} - assert toolkit.execute_tool("partial", {}).result == {"value": "pre-bound"} - assert toolkit.capture_calls == [] - assert toolkit.events.events == [] - - -def test_stateful_dispatch_captures_state_and_emits_the_record( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - clock = iter((10.0, 10.25)) - monkeypatch.setattr( - "rpent.tools.toolkit.time.perf_counter", - lambda: next(clock), - ) - toolkit = _ContractToolkit(tmp_path) - handler_result = {"moved": True} - - def move(*, distance: int) -> dict[str, bool]: - assert distance == 3 - return handler_result - - toolkit.add_tool("move", {"name": "move"}, move) - - result = toolkit.execute_tool("move", {"distance": 3}) - - assert result.result == {"observation": 1} - assert handler_result == {"moved": True} - assert toolkit.capture_calls[0]["command"] == { - "action": "move", - "distance": 3, - } - assert toolkit.capture_calls[0]["result"] == {"moved": True} - assert toolkit.capture_calls[0]["elapsed_s"] == 0.25 - assert len(toolkit.events.events) == 1 - event = toolkit.events.events[0] + output_dir=tmp_path, + tools=(finish, action, inspect_state), + dashboard_events=EventSink(), + ) + yield instance + instance.close() + + +def test_capture_replaces_action_data_appends_images_and_publishes_once(toolkit): + result = toolkit.execute_tool("action", {"value": "2"}) + assert not result.is_error + assert result.data == {"step": 0, "position": 3} + assert result.images == [b"action", b"observation"] + record = toolkit.state.latest_record() + assert record.command == {"action": "action", "value": 2} + assert record.result == {"moved": True} + assert record.elapsed_s >= 0 + (event,) = toolkit._dashboard_events.events assert isinstance(event, StepRecordEvent) - assert event.record.step_idx == 0 - assert event.record.command == {"action": "move", "distance": 3} + assert event.record is record assert event.env_state is toolkit.state + assert toolkit.execute_tool("inspect_state", {}).data == {"ready": True} + assert len(toolkit.state.records()) == 1 + assert len(toolkit._dashboard_events.events) == 1 + + +@pytest.mark.parametrize("arguments", [{"value": "invalid"}, {"value": None}]) +def test_validation_rejects_external_arguments_before_execution(toolkit, arguments): + result = toolkit.execute_tool("action", arguments) + assert result.is_error + message, details = result.error.split("\n", 1) + assert message == "Invalid arguments for action." + (error,) = json.loads(details)["errors"] + assert error["loc"] == ["value"] + assert "input" not in error + assert toolkit._robot.received == [] + assert toolkit.state.latest_record() is None + assert toolkit._dashboard_events.events == [] + + +def test_unknown_tool_does_not_execute_or_capture(toolkit): + assert toolkit.execute_tool("missing", {}).error == "Unknown tool: missing" + assert toolkit.state.latest_record() is None + assert toolkit._robot.received == [] + + +@pytest.mark.parametrize("field", ["valu", "ctx"]) +def test_unknown_arguments_are_rejected_before_execution(toolkit, field): + result = toolkit.execute_tool("action", {field: 3}) + assert result.is_error + details = json.loads(result.error.split("\n", 1)[1]) + assert details["errors"][0]["type"] == "extra_forbidden" + assert details["errors"][0]["loc"] == [field] + assert toolkit._robot.received == [] + assert toolkit.state.latest_record() is None + assert toolkit._dashboard_events.events == [] -def test_handler_error_is_retained_when_state_capture_also_fails( - tmp_path: Path, -) -> None: - toolkit = _ContractToolkit(tmp_path) - toolkit.capture_error = RuntimeError("capture exploded") - - def fail() -> dict[str, Any]: - raise ValueError("handler exploded") - - @readonly - def probe() -> dict[str, bool]: - return {"ready": True} - - toolkit.add_tool("fail", {"name": "fail"}, fail) - toolkit.add_tool("probe", {"name": "probe"}, probe) - - failed = toolkit.execute_tool("fail", {}) - - assert failed.result["error"] == "handler exploded" - assert failed.result["state_capture_error"] == "capture exploded" - assert "ValueError: handler exploded" in failed.result["traceback"] - assert toolkit.events.events == [] - assert toolkit.execute_tool("probe", {}).result == {"ready": True} - - -@pytest.mark.timeout(5) -def test_toolkit_rejects_overlapping_operations_and_cleans_up_after_success( - tmp_path: Path, -) -> None: - toolkit = _ContractToolkit(tmp_path) - started = threading.Event() - release = threading.Event() - results: list[ToolResult] = [] - worker_errors: list[BaseException] = [] - - @readonly - def blocking() -> dict[str, bool]: - started.set() - assert release.wait(2), "test did not release the blocking handler" - return {"released": True} - - toolkit.add_tool("blocking", {"name": "blocking"}, blocking) - - def run_blocking() -> None: - try: - results.append(toolkit.execute_tool("blocking", {})) - except BaseException as error: - worker_errors.append(error) - - worker = threading.Thread(target=run_blocking, daemon=True) - worker.start() - try: - assert started.wait(2), "blocking handler did not start" - overlap = toolkit.execute_tool("finish", {"status": "failure", "summary": "x"}) - assert overlap.result == {"error": "another tool operation is still active"} - finally: - release.set() - worker.join(2) - - assert not worker.is_alive() - assert worker_errors == [] - assert len(results) == 1 - assert results[0].result == {"released": True} - assert toolkit.execute_tool( - "finish", {"status": "success", "summary": "clean"} - ).is_finish - - -def test_toolkit_cleans_up_operation_after_handler_failure(tmp_path: Path) -> None: - toolkit = _ContractToolkit(tmp_path) - - @readonly - def fail() -> dict[str, Any]: - raise RuntimeError("tool failed") - - toolkit.add_tool("fail", {"name": "fail"}, fail) - - failed = toolkit.execute_tool("fail", {}) - - assert failed.result["error"] == "tool failed" - assert "RuntimeError: tool failed" in failed.result["traceback"] - assert toolkit.execute_tool( - "finish", {"status": "failure", "summary": "recovered"} - ).is_finish - - -@pytest.mark.timeout(5) -def test_toolkit_cooperatively_cancels_and_cleans_up_active_operation( - tmp_path: Path, -) -> None: - toolkit = _ContractToolkit(tmp_path) - started = threading.Event() - stop_polling = threading.Event() - results: list[ToolResult] = [] - - def cancellable() -> dict[str, bool]: - started.set() - while not stop_polling.wait(0.01): - toolkit.raise_if_cancelled() - return {"unexpected": True} - - toolkit.add_tool("cancellable", {"name": "cancellable"}, cancellable) - worker = threading.Thread( - target=lambda: results.append(toolkit.execute_tool("cancellable", {})), - daemon=True, - ) - worker.start() - try: - assert started.wait(2), "cancellable handler did not start" - toolkit.cancel_active_and_wait() - finally: - stop_polling.set() - worker.join(2) - - assert not worker.is_alive() - assert results[0].result["code"] == "tool_cancelled" - assert results[0].result["interrupted"] is True - assert results[0].result["error"] == "tool operation interrupted" - assert toolkit.capture_calls[0]["result"]["code"] == "tool_cancelled" - assert len(toolkit.events.events) == 1 - assert toolkit.execute_tool( - "finish", {"status": "failure", "summary": "cancelled"} - ).is_finish +@pytest.mark.parametrize( + "error_type", [ToolCancelled, PermissionError, TypeError, ValueError] +) +def test_handler_failure_is_bounded_logged_and_still_captured( + toolkit, error_type, caplog +): + message = "failure " * 100 + toolkit._robot.failure = error_type(message) + result = toolkit.execute_tool("action", {}) + assert result.error == message[:500] + assert result.data == {"step": 0, "position": 3} + assert toolkit.state.latest_record().result == {"error": message[:500]} + assert message in caplog.text + assert not toolkit.execute_tool("inspect_state", {}).is_error + + +@pytest.mark.parametrize("error", [None, "", "action failed"]) +def test_capture_failure_preserves_action_data_images_and_error( + toolkit, monkeypatch, error +): + toolkit._robot.result.error = error + + def fail(**kwargs): + raise OSError("camera offline") + + monkeypatch.setattr(toolkit, "_capture_observation", fail) + result = toolkit.execute_tool("action", {}) + prefix = f"{error}\n" if error is not None else "" + assert result.error == prefix + "State capture failed: camera offline" + assert result.data == {"moved": True} + assert result.images == [b"action"] + assert toolkit.state.latest_record() is None + assert not toolkit.execute_tool("inspect_state", {}).is_error + + +@pytest.mark.parametrize("error", ["", "business failure"]) +def test_successful_capture_retains_explicit_action_error(toolkit, error): + toolkit._robot.result.error = error + result = toolkit.execute_tool("action", {}) + assert result.is_error + assert result.error == error + assert result.data == {"step": 0, "position": 3} + assert toolkit.state.latest_record().result == {"moved": True, "error": error} + + +def test_capture_failure_after_recording_still_publishes_saved_step( + toolkit, monkeypatch +): + capture = toolkit._capture_observation + + def fail(**kwargs): + capture(**kwargs) + raise OSError("image unavailable") + + monkeypatch.setattr(toolkit, "_capture_observation", fail) + result = toolkit.execute_tool("action", {}) + assert result.error == "State capture failed: image unavailable" + (event,) = toolkit._dashboard_events.events + assert event.record is toolkit.state.latest_record() + + +def test_dashboard_failure_does_not_retry_or_change_result( + toolkit, monkeypatch, caplog +): + calls = [] + + def fail(event): + calls.append(event) + raise RuntimeError("dashboard offline") + + monkeypatch.setattr(toolkit._dashboard_events, "emit", fail) + result = toolkit.execute_tool("action", {}) + assert not result.is_error + assert len(calls) == len(toolkit.state.records()) == 1 + assert "dashboard offline" in caplog.text + + +def test_base_exception_propagates_and_releases_admission(toolkit): + toolkit._robot.failure = KeyboardInterrupt() + with pytest.raises(KeyboardInterrupt): + toolkit.execute_tool("action", {}) + assert not toolkit.execute_tool("inspect_state", {}).is_error + + +def test_context_is_fresh_for_each_call(toolkit): + for _ in range(2): + toolkit.execute_tool("action", {"value": "3"}) + first, second = [ctx for value, ctx in toolkit._robot.received] + assert first is not second + assert first._cancel_event is not second._cancel_event + for ctx in (first, second): + assert ctx.robot is toolkit._robot + assert ctx.state is toolkit.state + assert ctx.memory is toolkit.memory + assert ctx.output_dir == toolkit._task_output_dir + assert ctx.record_frame == toolkit.record_frame From b2ccc52ea41e27861d3d3a0945b1c5acfbe5b7de Mon Sep 17 00:00:00 2001 From: wilburx813 Date: Thu, 10 Sep 2026 17:48:21 +0800 Subject: [PATCH 2/2] refactor: unify tool scheduling, recording, and recipe export --- .../rst_source/development/add_primitive.rst | 4 +- .../rst_source/development/add_robot.rst | 10 +- .../rst_source/development/architecture.rst | 10 +- .../rst_source/development/interfaces.rst | 30 +- .../rst_source/development/memory.rst | 9 +- .../rst_source/usage/configure_planner.rst | 2 +- .../rst_source/development/add_primitive.rst | 4 +- .../rst_source/development/add_robot.rst | 5 +- .../rst_source/development/architecture.rst | 4 +- .../rst_source/development/interfaces.rst | 12 +- .../rst_source/development/memory.rst | 6 +- .../rst_source/usage/configure_planner.rst | 2 +- robots/dual_franka/toolkit.py | 12 +- robots/libero/toolkit.py | 118 +--- robots/robocasa/toolkit.py | 65 --- robots/robotwin/toolkit.py | 50 -- rpent/dashboard/planner_control.py | 30 +- rpent/planner/api_loop.py | 11 +- rpent/planner/base.py | 22 +- rpent/planner/claude_code.py | 84 ++- rpent/planner/codex.py | 28 +- rpent/planner/utils/http_mcp_server.py | 5 +- rpent/tools/base.py | 7 +- rpent/tools/toolkit.py | 190 +++++-- .../dual_franka/test_dual_franka_tools.py | 21 + tests/unit_tests/robots/franka/test_tools.py | 5 + tests/unit_tests/robots/libero/conftest.py | 3 +- .../robots/libero/test_libero_integration.py | 3 + .../libero/test_libero_toolkit_contracts.py | 4 +- tests/unit_tests/robots/libero/test_recipe.py | 94 ++++ .../robots/libero/test_recording.py | 143 +++++ tests/unit_tests/robots/robocasa/conftest.py | 3 +- .../test_robocasa_toolkit_contracts.py | 7 +- tests/unit_tests/robots/robotwin/conftest.py | 3 +- .../test_planner_control_contracts.py | 30 ++ tests/unit_tests/rpent/planner/conftest.py | 3 +- .../rpent/planner/test_api_contracts.py | 2 +- .../rpent/planner/test_execution_lifecycle.py | 138 +++++ .../rpent/planner/test_http_mcp_server.py | 4 +- .../rpent/planner/test_native_adapters.py | 127 ++++- .../rpent/tools/test_common_tools.py | 29 +- tests/unit_tests/rpent/tools/test_recipe.py | 72 +++ .../unit_tests/rpent/tools/test_recording.py | 77 +++ .../unit_tests/rpent/tools/test_scheduling.py | 504 ++++++++++++++++-- 44 files changed, 1579 insertions(+), 413 deletions(-) create mode 100644 tests/unit_tests/robots/libero/test_recipe.py create mode 100644 tests/unit_tests/robots/libero/test_recording.py create mode 100644 tests/unit_tests/rpent/planner/test_execution_lifecycle.py create mode 100644 tests/unit_tests/rpent/tools/test_recipe.py create mode 100644 tests/unit_tests/rpent/tools/test_recording.py diff --git a/docs/source-en/rst_source/development/add_primitive.rst b/docs/source-en/rst_source/development/add_primitive.rst index 4d92527c7..e7b20afbb 100644 --- a/docs/source-en/rst_source/development/add_primitive.rst +++ b/docs/source-en/rst_source/development/add_primitive.rst @@ -84,10 +84,10 @@ it. The toolkit handles capture through ``_capture_observation``; handlers submit frames but do not save their own episode video or duplicate the state dump. For a tool that reads existing observations, place ``@readonly`` below -``@tool`` to skip automatic capture. Calls still execute one at a time. +``@tool`` to allow concurrent execution and skip automatic capture. ``write_text_file`` and ``finish`` have no readonly marker and run exclusively; the executor skips observation capture for common tools and ``finish``. -See :doc:`interfaces` for execution and cancellation. +See :doc:`interfaces` for cancellation and scheduling. .. _add-primitive-model-based: diff --git a/docs/source-en/rst_source/development/add_robot.rst b/docs/source-en/rst_source/development/add_robot.rst index 23742d422..e56a933e5 100644 --- a/docs/source-en/rst_source/development/add_robot.rst +++ b/docs/source-en/rst_source/development/add_robot.rst @@ -324,9 +324,9 @@ for both action responses and ``view_env_state``. executor stores accepted results in ``toolkit.finish_result``. Tools submit environment-step RGB frames via ``ctx.record_frame(rgb)``. -The base toolkit collects the frames and signals cancellation. Robot toolkits -save Dashboard action clips during observation capture, override ``close()`` -to save the episode video, and implement ``write_recipe()`` from their trace. +The base toolkit owns the frame buffer, Dashboard action clips, episode video, +recipe export, and cancellation lifecycle. Use the inherited ``close()`` for +recording cleanup; if a robot needs additional cleanup, preserve the base call. Conventions worth keeping ------------------------- @@ -334,8 +334,8 @@ Conventions worth keeping - ``output_dir`` is the runner-created working directory. Environment artifacts are managed by ``EnvState`` through logical names; transcripts share the same run directory. -- ``@readonly`` below ``@tool`` skips automatic observation capture. Calls - execute one at a time. Call +- ``@readonly`` below ``@tool`` enables shared execution and skips automatic + observation capture. Keep actions and ``finish`` exclusive and call ``ctx.check_cancelled()`` at safe boundaries in long loops. - Server-side values use transport-supported Python / NumPy types and remain torch-free. diff --git a/docs/source-en/rst_source/development/architecture.rst b/docs/source-en/rst_source/development/architecture.rst index 680e7433d..e48b3e6bd 100644 --- a/docs/source-en/rst_source/development/architecture.rst +++ b/docs/source-en/rst_source/development/architecture.rst @@ -175,15 +175,15 @@ invokes them through ``execute_tool``. It adapts schemas and ``ToolResult`` text / PNG images to its SDK; Claude Code and Codex use MCP adapters at this boundary. Robot handlers are independent of the planner transport. -``Toolkit`` validates arguments, admits one call, injects ``ToolContext``, and +``Toolkit`` validates arguments, schedules calls, injects ``ToolContext``, and captures post-action observations. Handlers use ``ctx.robot`` to access their session runtime and issue ``reset`` / ``step`` / ``predict`` requests through environment and model clients. HTTP or socket RPC carries those calls and NumPy observations between processes. -The toolkit also owns cancellation, the frame buffer, and ``finish_result``. -Robot-specific subclasses build observations, save videos and recipes, and -report native success through ``solved()``. See :doc:`interfaces` for the contracts. +The toolkit also owns cancellation, frame recording, recipe export, and +``finish_result``. Robot-specific subclasses build observations and report +native success through ``solved()``. See :doc:`interfaces` for the contracts. Dashboard (optional) -------------------- @@ -196,7 +196,7 @@ the CLI before it calls ``robot_spec.init_runtime`` once with the shared component names. The environment must provide ``robot_spec.dashboard``; it defines the task command and fields, runtime components, and allowed primitive controls. -Camera tabs use the ``frame_channels`` mapping to recorded image artifacts. The Session controller waits for that robot-defined command +Camera tabs are discovered from the PNG artifacts in each recorded step. The Session controller waits for that robot-defined command (``/rpent-task`` for LIBERO). For every claimed TaskRun, the Dashboard calls ``parse_config`` and the same ``robot_spec.init_runtime`` hook with the unique component names, merges the shared and unique diff --git a/docs/source-en/rst_source/development/interfaces.rst b/docs/source-en/rst_source/development/interfaces.rst index 137645446..ca15783dc 100644 --- a/docs/source-en/rst_source/development/interfaces.rst +++ b/docs/source-en/rst_source/development/interfaces.rst @@ -85,7 +85,7 @@ Contract: read ``toolkit.list_tools()`` and adapt each ``Tool``'s ``name``, ``toolkit.execute_tool(name, arguments)`` and return ``PlannerResult`` when ``toolkit.finish_result`` is set or a run limit is reached. For asynchronous adapters, use ``rpent.planner.base.execute_tool`` to run the synchronous executor -in a worker. The API and MCP adapters serialize tool calls. +in a worker and retain it through cancellation. Native tools and Toolkit ------------------------ @@ -144,12 +144,12 @@ error handling, and validated argument logs follow the shared native executor. Arm normalization is declared in the Pydantic parameter type. -Execution and lifecycle -~~~~~~~~~~~~~~~~~~~~~~~ +Scheduling and lifecycle +~~~~~~~~~~~~~~~~~~~~~~~~ -Each toolkit permits one active call. Overlapping direct calls return an error; -API and MCP adapters serialize their calls. Place ``@readonly`` below ``@tool`` -to skip automatic observation capture. +Tools run exclusively by default. Place ``@readonly`` below ``@tool`` to allow +execution alongside other readonly tools and skip automatic observation capture. +Pending exclusive calls take priority and keep their queue order. Non-readonly robot tools capture a new observation after execution. Common tools and ``finish`` are excluded from capture: ``write_text_file`` and ``finish`` run @@ -165,19 +165,19 @@ Capture also runs after handler errors; the call remains active until capture and Dashboard publication finish. Long-running handlers call ``ctx.check_cancelled()`` at safe boundaries. -``cancel_active_and_wait()`` signals the active call and waits for it to exit. -Subsequent calls receive a fresh cancellation signal. Tools submit RGB frames -with ``ctx.record_frame``; robot toolkits save per-action clips during capture -and override ``close()`` to save their episode video. +``cancel_active_and_wait()`` pauses admission, cancels pending and active calls, +and waits for cleanup; ``resume_calls()`` reopens admission. +``close()`` permanently closes admission, drains calls, and saves collected +frames as ``episode.mp4``. Tools submit RGB frames with ``ctx.record_frame``; +when Dashboard events are enabled, the executor also saves per-action clips. Each robot supplies its own ``finish`` tool. A successful call stores its ``status`` and ``summary`` in ``toolkit.finish_result``; it does not close admission. ``solved()`` reports environment success independently of the -planner's requested finish status. Robot toolkits implement -``write_recipe(recipe_tag)`` using their recorded state trace. LIBERO exports -the successful attempt after the last reset; RoboCasa and RoboTwin retain -their action filters. The runner decides whether the run qualifies for memory -publication. +planner's requested finish status. ``write_recipe(recipe_tag)`` exports +successful robot calls, including perception and resets, in completion order; +common file/image tools and ``finish`` are excluded. The runner decides whether +the run qualifies for memory publication. Inter-process communication --------------------------- diff --git a/docs/source-en/rst_source/development/memory.rst b/docs/source-en/rst_source/development/memory.rst index 09fff8df3..14786b8ba 100644 --- a/docs/source-en/rst_source/development/memory.rst +++ b/docs/source-en/rst_source/development/memory.rst @@ -85,7 +85,8 @@ exploration may write to its configured ``_internal/inbox//``. Access to another robot's repository memory is denied. These checks govern memory access; they do not restrict all files to the output directory. -Each robot toolkit implements ``write_recipe(recipe_tag)`` from its state -trace. LIBERO exports the successful attempt after the last reset; RoboCasa -and RoboTwin filter recorded actions using their existing recipe rules. The -runner decides whether the audit and recipe qualify for publication to memory. +``Toolkit.write_recipe(recipe_tag)`` exports successful robot action and +perception calls in completion order, including resets across attempts in the +same session. Common file/image tools and ``finish`` are excluded. A failed +call is not exported. The runner decides whether the resulting audit and +recipe qualify for publication to memory. diff --git a/docs/source-en/rst_source/usage/configure_planner.rst b/docs/source-en/rst_source/usage/configure_planner.rst index ead12df35..801419f46 100644 --- a/docs/source-en/rst_source/usage/configure_planner.rst +++ b/docs/source-en/rst_source/usage/configure_planner.rst @@ -306,7 +306,7 @@ Any planner must: 2. Read native tools from ``toolkit.list_tools()`` and adapt their ``name``, ``description``, and ``input_schema`` to the SDK. Execute calls through ``toolkit.execute_tool(name, arguments)``; asynchronous adapters use - ``rpent.planner.base.execute_tool`` to execute tools in a worker thread. + ``rpent.planner.base.execute_tool`` for cancellation-aware execution. 3. Convert ``ToolResult.to_text()`` and the PNG bytes in ``ToolResult.images`` to the SDK format, preserving ``ToolResult.is_error``. 4. Check ``toolkit.finish_result`` and stop according to ``max_turns`` and diff --git a/docs/source-zh/rst_source/development/add_primitive.rst b/docs/source-zh/rst_source/development/add_primitive.rst index 3047ef1f8..b5b2dc88b 100644 --- a/docs/source-zh/rst_source/development/add_primitive.rst +++ b/docs/source-zh/rst_source/development/add_primitive.rst @@ -79,14 +79,14 @@ 工具执行后，toolkit 会自动保存新的状态快照。对于 ``view_env_state``、 ``back_project`` 等读取已有观测的工具，可以在 ``@tool`` 下方添加 - ``@readonly``，省去这次状态捕获；调用仍按顺序执行。公共工具和 ``finish`` + ``@readonly``，允许并行执行并省去这次状态捕获。公共工具和 ``finish`` 不触发状态捕获；``write_text_file`` 和 ``finish`` 不设置 readonly，独占执行。 2. **将工具加入 toolkit。** 把函数声明加入该机器人的工具集合，例如 LIBERO 的 ``LIBERO_TOOLS``。Toolkit 在构造时接收这组工具，并统一处理参数校验和调用。 完成以上步骤后，``api``、``claude_code`` 和 ``codex`` 三种 planner -都可以调用该工具，无需分别编写适配代码。执行和取消的约定参见 +都可以调用该工具，无需分别编写适配代码。需要支持并发或取消时，参见 :doc:`interfaces` 中的工具集说明。 .. _add-primitive-model-based: diff --git a/docs/source-zh/rst_source/development/add_robot.rst b/docs/source-zh/rst_source/development/add_robot.rst index 1344d710b..e4a5da6b2 100644 --- a/docs/source-zh/rst_source/development/add_robot.rst +++ b/docs/source-zh/rst_source/development/add_robot.rst @@ -295,9 +295,8 @@ VLA 客户端和本次会话的状态。工具函数通过 ``ctx.robot`` 访问 - 实现 ``_capture_observation(*, command, result, elapsed_s)``，调用上述 状态保存与观测整理函数，返回观测数据和 PNG 图片。Toolkit 会在动作执行后 自动调用它；原始执行结果可通过 ``result.to_dict()`` 保存到步骤日志中。 -- 实现 ``solved()``，根据环境状态判断任务是否成功。工具调用与取消由基类处理。 - 机器人 toolkit 在捕获观测时保存动作视频，重写 ``close()`` 保存回合录像， - 并通过 ``write_recipe()`` 从记录的状态导出 recipe。 +- 实现 ``solved()``，根据环境状态判断任务是否成功。工具调用、取消和录像收尾 + 由基类处理；若需要额外的关闭逻辑，在重写 ``close()`` 时保留对基类的调用。 ``runtime_kwargs`` 由 ``robot_spec.py:get_toolkit`` 转发给 toolkit，用于构造 运行时对象。其中通常包含 ``{"env": MyEnvClient(...), "model": VLAClient(...)}`` diff --git a/docs/source-zh/rst_source/development/architecture.rst b/docs/source-zh/rst_source/development/architecture.rst index 484107e4a..402dc4147 100644 --- a/docs/source-zh/rst_source/development/architecture.rst +++ b/docs/source-zh/rst_source/development/architecture.rst @@ -168,8 +168,8 @@ Dashboard（可选） ``--dashboard-host`` 和 ``--dashboard-port`` 启动 Dashboard。Session 配置全部来自 命令行，然后用共享 component 名称调用一次 ``robot_spec.init_runtime``。环境必须 提供 ``robot_spec.dashboard``，由它定义 -前端使用的任务命令与字段、runtime components 和允许执行的原语。相机标签通过 -``frame_channels`` 映射到每步记录的图片工件。Session +前端使用的任务命令与字段、runtime components 和允许执行的原语。相机标签从 +每步记录的 PNG 工件中自动发现。Session controller 随后等待该环境定义的命令（LIBERO 使用 ``/rpent-task``）；每次取得一个 TaskRun 后，Dashboard 会调用 ``parse_config``，再用 unique component 名称调用 同一个 ``robot_spec.init_runtime``，合并两次返回的客户端参数，并为本次任务新建 toolkit diff --git a/docs/source-zh/rst_source/development/interfaces.rst b/docs/source-zh/rst_source/development/interfaces.rst index dfba02f0e..badb1d46d 100644 --- a/docs/source-zh/rst_source/development/interfaces.rst +++ b/docs/source-zh/rst_source/development/interfaces.rst @@ -145,18 +145,18 @@ arm 的归一化声明在 Pydantic 参数类型中。 应写入观测中的日志；动作错误仍会保留。即使工具函数出错，toolkit 也会尝试 捕获观测，让 planner 了解当前环境。 -读取已有观测的工具可以在 ``@tool`` 下方添加 ``@readonly``，跳过自动捕获。 -每个 toolkit 同时只允许一个调用；重叠的直接调用会返回错误。API 和 MCP -适配器按顺序执行工具调用。 +读取已有观测的工具可以在 ``@tool`` 下方添加 ``@readonly``，允许与其他 +readonly 工具并行，并跳过自动捕获。其余工具独占执行；等待中的独占调用 +优先执行，并保持登记顺序。 公共工具和 ``finish`` 不触发观测捕获。``write_text_file`` 和 ``finish`` 不设置 readonly，因此独占执行但不新增观测。LIBERO 的 ``segment`` 使用 ``@readonly``，将分割附件保存到源 step 并直接返回分割结果，不新增观测。 长时间运行的工具应在安全的动作边界调用 ``ctx.check_cancelled()``。收到中断后， -``cancel_active_and_wait()`` 向当前调用发送取消信号，并等待它退出。后续调用 -使用新的取消信号。工具通过 ``ctx.record_frame`` 提交录像帧；机器人 toolkit -在捕获观测时保存动作片段，并重写 ``close()`` 保存回合录像。 +``cancel_active_and_wait()`` 会暂停新调用，并等待已有调用取消和清理；后续可以用 +``resume_calls()`` 恢复。运行结束时，``close()`` 关闭工具集并保存通过 +``ctx.record_frame`` 收集的录像帧。 每个机器人提供自己的 ``finish`` 工具。调用成功后，toolkit 将其中的 ``status`` 和 ``summary`` 保存到 ``finish_result``，供 planner 结束循环；环境是否真正成功， diff --git a/docs/source-zh/rst_source/development/memory.rst b/docs/source-zh/rst_source/development/memory.rst index f45b428e0..bd580f8a2 100644 --- a/docs/source-zh/rst_source/development/memory.rst +++ b/docs/source-zh/rst_source/development/memory.rst @@ -61,9 +61,9 @@ memory 同步到 ``memory//``。数据集是公开的，无需 token 即 公共文件工具已接入这些检查；新增需要访问 memory 的工具时，可通过 ``ctx.memory.authorize_read(path)`` 或 ``authorize_write(path)`` 获取允许访问的路径。 -各机器人 toolkit 通过 ``write_recipe(recipe_tag)`` 从状态记录导出 recipe。 -LIBERO 只导出最后一次 reset 后的成功尝试；RoboCasa 和 RoboTwin 保留各自的 -动作筛选规则。是否将生成的 audit 和 recipe 纳入 memory，由运行流程根据结果决定。 +运行中的 recipe 由 ``Toolkit.write_recipe(recipe_tag)`` 统一导出，记录成功的 +机器人动作与感知调用，也保留同一会话中的 reset。文件操作和 ``finish`` 不会写入 +recipe。是否将生成的 audit 和 recipe 纳入 memory，仍由运行流程根据结果决定。 贡献 memory ----------- diff --git a/docs/source-zh/rst_source/usage/configure_planner.rst b/docs/source-zh/rst_source/usage/configure_planner.rst index 4417a7bf2..5a5000073 100644 --- a/docs/source-zh/rst_source/usage/configure_planner.rst +++ b/docs/source-zh/rst_source/usage/configure_planner.rst @@ -279,7 +279,7 @@ agent SDK，可以实现 ``rpent.planner.base.Planner`` 协议，并在 1. 接收已经渲染好的 ``system_prompt`` 和 ``user_message``。 2. 从 ``toolkit.list_tools()`` 读取原生工具，并将其 ``name``、``description`` 和 ``input_schema`` 转为 SDK 格式。通过 ``toolkit.execute_tool(name, arguments)`` - 执行调用；异步适配器通过 ``rpent.planner.base.execute_tool`` 在线程中执行工具。 + 执行调用；异步适配器使用支持取消清理的 ``rpent.planner.base.execute_tool``。 3. 将 ``ToolResult.to_text()`` 和 ``ToolResult.images`` 中的 PNG 字节转换为 SDK 格式，并保留 ``ToolResult.is_error``。 4. 检查 ``toolkit.finish_result``，按 ``max_turns`` 等限制终止循环； diff --git a/robots/dual_franka/toolkit.py b/robots/dual_franka/toolkit.py index 52c8ed489..c7f849cd7 100644 --- a/robots/dual_franka/toolkit.py +++ b/robots/dual_franka/toolkit.py @@ -140,7 +140,7 @@ def request_direct_verdict(self, verdict: str) -> bool: """Seal this attempt immediately; cancel in-flight work at its next boundary.""" if verdict not in {"success", "failure", "abort"}: raise ValueError("verdict must be success, failure or abort") - with self._operation_lock: + with self._scheduler._condition: if self._direct_verdict_event.is_set(): return self._direct_verdict == verdict if self._mode != "exploration" or ( @@ -150,8 +150,7 @@ def request_direct_verdict(self, verdict: str) -> bool: return False self._direct_verdict = verdict self._direct_verdict_event.set() - if self._active_operation is not None: - self._active_operation.cancel_event.set() + self._scheduler.cancel() return True def raise_if_cancelled(self) -> None: @@ -159,9 +158,10 @@ def raise_if_cancelled(self) -> None: raise ToolCancelled( "operator submitted a terminal verdict; stopping exploration" ) - with self._operation_lock: - operation = self._active_operation - cancelled = operation is not None and operation.cancel_event.is_set() + with self._scheduler._condition: + cancelled = any( + call.cancel_event.is_set() for call in self._scheduler._active_calls + ) if cancelled: raise ToolCancelled("Tool call cancelled.") diff --git a/robots/libero/toolkit.py b/robots/libero/toolkit.py index 9b653a344..165a06765 100644 --- a/robots/libero/toolkit.py +++ b/robots/libero/toolkit.py @@ -16,7 +16,6 @@ from __future__ import annotations -import json from dataclasses import replace from functools import partial from pathlib import Path @@ -144,14 +143,11 @@ def __init__( ) except Exception: logger.exception("Dashboard failed to publish step %s", record.step_idx) - self._action_frame_cursor = 0 def _capture_observation( self, *, command: dict[str, Any], result: ToolResult, elapsed_s: float ) -> tuple[dict[str, Any], list[bytes]]: """Save the full action log, then assemble the current observation response.""" - frame_start = self._action_frame_cursor - self._action_frame_cursor = len(self._frames) logged_result = result.to_dict() record = dump_state( self._robot, @@ -163,21 +159,6 @@ def _capture_observation( }, ) self._robot.solved |= record.terminated - if self._dashboard_events.enabled: - try: - frames = self._frames[frame_start:] - if frames: - self._state.save( - f"action_{command['action']}.mp4", - frames, - step=record.step_idx, - fps=20, - ) - except Exception as exc: - logger.warning( - "failed to save action clip for step %s: %s", record.step_idx, exc - ) - record = self._state.get(record.step_idx) data, images = build_observation(self._state, record) if result.is_error: # Report the error once in this response; retain it in the saved log @@ -197,7 +178,8 @@ def solved(self) -> bool: return self._robot.solved def close(self) -> None: - """Finalize collected data and save the episode video independently.""" + """Drain calls and save the video before finalizing collected data.""" + super().close() try: episode = self._robot.finalize_flywheel() if episode is not None: @@ -205,18 +187,6 @@ def close(self) -> None: except Exception as exc: logger.warning("failed to finalize flywheel episode: %s", exc) - try: - if self._frames: - self._state.save("episode.mp4", self._frames, step=None, fps=20) - except Exception as exc: - logger.warning("failed to save episode video: %s", exc) - - def write_recipe(self, recipe_tag: str) -> str: - """Export the successful attempt from the recorded LIBERO trace.""" - return write_recipe_from_states( - self._state, recipe_tag, output_dir=self._task_output_dir - ) - def dump_state( runtime: LiberoRuntime, @@ -423,87 +393,3 @@ def _world_from_depth(depth_metric: np.ndarray, camera_meta: dict) -> np.ndarray axis=-1, ) return (camera_points @ extrinsic.T)[..., :3] - - -def _is_primitive_action(name: object) -> bool: - return name in { - "reset", - "pi0_pick", - "pi0_doubled", - "move_to", - "rotate_wrist", - "rotate_pitch", - "move_pose", - "release", - "set_gripper", - } - - -def write_recipe_from_states( - state: EnvState, recipe_tag: str, *, output_dir: Path | str -) -> str: - """Find a command sequence that gets ``terminated=True``. - - Export non-error LIBERO primitive commands and successful segment calls. - """ - records = state.records() - last_reset = max( - ( - record.step_idx - for record in records - if ( - (record.command or {}).get("action") == "reset" - and not (isinstance(record.result, dict) and record.result.get("error")) - ) - ), - default=-1, - ) - command_events = [] - for record in records: - if record.step_idx <= last_reset: - continue - 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)) - - # Never publish a failed trajectory as a recipe. The environment trace is - # authoritative; an agent's self-reported finish status is not. - solved = any( - record.terminated for record in records if record.step_idx > last_reset - ) - if not solved: - return "" - command_events.sort(key=lambda event: event[0]) - recipe_name = f"{recipe_tag}_recipe.jsonl" - recipe_path = Path(output_dir) / recipe_name - recipe_path.parent.mkdir(parents=True, exist_ok=True) - recipe_path.write_text( - "".join(json.dumps(command) + "\n" for _, command in command_events) - ) - return recipe_name diff --git a/robots/robocasa/toolkit.py b/robots/robocasa/toolkit.py index a67581e3d..a902e6f20 100644 --- a/robots/robocasa/toolkit.py +++ b/robots/robocasa/toolkit.py @@ -100,7 +100,6 @@ def __init__( self._dashboard_events.emit(StepRecordEvent(record=record, env_state=state)) except Exception: logger.exception("Dashboard failed to publish step %s", record.step_idx) - self._action_frame_cursor = 0 def _capture_observation( self, @@ -109,29 +108,12 @@ def _capture_observation( result: ToolResult, elapsed_s: float, ) -> tuple[dict[str, Any], list[bytes]]: - frame_start = self._action_frame_cursor - self._action_frame_cursor = len(self._frames) logged_result = result.to_dict() record = dump_state( self._robot, self._state, log={"command": command, "result": logged_result, "elapsed_s": elapsed_s}, ) - if self._dashboard_events.enabled: - try: - frames = self._frames[frame_start:] - if frames: - self._state.save( - f"action_{command['action']}.mp4", - frames, - step=record.step_idx, - fps=20, - ) - except Exception as exc: - logger.warning( - "failed to save action clip for step %s: %s", record.step_idx, exc - ) - record = self._state.get(record.step_idx) data, images = build_observation(self._state, record) if result.is_error: data["log"]["result"] = { @@ -145,18 +127,6 @@ def solved(self) -> bool: record = self._state.latest_record() return bool(record is not None and record.extras.get("success", False)) - def close(self) -> None: - """Save this robot's accumulated episode frames.""" - try: - if self._frames: - self._state.save("episode.mp4", self._frames, step=None, fps=20) - except Exception as exc: - logger.warning("failed to save episode video: %s", exc) - - def write_recipe(self, recipe_tag: str) -> str: - """Export non-error RoboCasa commands from the recorded trace.""" - return write_recipe_from_states(self._state, recipe_tag) - # Heavy npy artifacts pruned after the ``_keep_heavy`` window elapses (the # agent localizes from the latest frame; old world/depth maps are dead weight @@ -319,38 +289,3 @@ def build_observation( break return out, images - - -_PRIMITIVE_ACTIONS = frozenset( - { - "move_to", - "move_delta", - "rotate_pitch", - "set_gripper", - "release", - "scripted_grasp", - "rldx_skill", - "rldx_arm", - "navigate_to", - "move_base", - "reset", - } -) - - -def write_recipe_from_states(state: EnvState, recipe_tag: str) -> str: - """Export non-error RoboCasa primitive commands from the state trace as JSONL.""" - commands = [] - for record in state.records(): - command = record.command - if not isinstance(command, dict): - continue - if command.get("action") not in _PRIMITIVE_ACTIONS: - continue - result = record.result - if isinstance(result, dict) and result.get("error"): - continue - commands.append(command) - recipe_name = f"{recipe_tag}_recipe.jsonl" - state.save(recipe_name, commands, step=None) - return recipe_name diff --git a/robots/robotwin/toolkit.py b/robots/robotwin/toolkit.py index f4a0cefe6..8e320857a 100644 --- a/robots/robotwin/toolkit.py +++ b/robots/robotwin/toolkit.py @@ -98,7 +98,6 @@ def __init__( self._dashboard_events.emit(StepRecordEvent(record=record, env_state=state)) except Exception: logger.exception("Dashboard failed to publish step %s", record.step_idx) - self._action_frame_cursor = 0 def _capture_observation( self, @@ -107,29 +106,12 @@ def _capture_observation( result: ToolResult, elapsed_s: float, ) -> tuple[dict[str, Any], list[bytes]]: - frame_start = self._action_frame_cursor - self._action_frame_cursor = len(self._frames) logged_result = result.to_dict() record = dump_state( self._robot, self._state, log={"command": command, "result": logged_result, "elapsed_s": elapsed_s}, ) - if self._dashboard_events.enabled: - try: - frames = self._frames[frame_start:] - if frames: - self._state.save( - f"action_{command['action']}.mp4", - frames, - step=record.step_idx, - fps=20, - ) - except Exception as exc: - logger.warning( - "failed to save action clip for step %s: %s", record.step_idx, exc - ) - record = self._state.get(record.step_idx) data, images = build_observation(self._state, record) if result.is_error: data["log"]["result"] = { @@ -146,35 +128,6 @@ def solved(self) -> bool: and record.state["episode_status"].get("eval_success") is True ) - def close(self) -> None: - """Save this robot's accumulated episode frames.""" - try: - if self._frames: - self._state.save("episode.mp4", self._frames, step=None, fps=20) - except Exception as exc: - logger.warning("failed to save episode video: %s", exc) - - def write_recipe(self, recipe_tag: str) -> str: - """Export state-advancing RoboTwin primitives with no error and no - explicit ``success=False`` from ``EnvState.records()``.""" - recipe = [ - record.command - for record in self._state.records() - if isinstance(record.command, dict) - and record.command.get("action") in _RECIPE_ACTIONS - and not ( - isinstance(record.result, dict) - and ( - record.result.get("error") or record.result.get("success") is False - ) - ) - ] - name = f"{recipe_tag}_recipe.jsonl" - saved = self._state.save(name, recipe, step=None) - if saved is None: - raise RuntimeError(f"failed to save RoboTwin recipe artifact: {name}") - return str(self._state.artifact_path(name, step=None)) - def _world_from_depth( depth_metric: np.ndarray, camera_meta: dict[str, Any] @@ -310,6 +263,3 @@ def build_observation( except FileNotFoundError: pass return data, images - - -_RECIPE_ACTIONS = {"lingbot_act", "move_to", "rotate_wrist", "set_gripper", "release"} diff --git a/rpent/dashboard/planner_control.py b/rpent/dashboard/planner_control.py index ac9a0126c..872254aa5 100644 --- a/rpent/dashboard/planner_control.py +++ b/rpent/dashboard/planner_control.py @@ -21,6 +21,7 @@ from typing import Any from rpent.dashboard.interaction import DashboardInteractionPort +from rpent.planner.base import cancel_and_wait class DashboardPlannerControl: @@ -31,17 +32,20 @@ def __init__( *, interaction: DashboardInteractionPort, cancel_active_and_wait: Callable[[], None], + resume_calls: Callable[[], None], emit_user: Callable[[str], None], emit_initial_user: Callable[[], None], defer_message_ack: bool = False, ) -> None: self._interaction = interaction self._cancel_active_and_wait = cancel_active_and_wait + self._resume_calls = resume_calls self._emit_user = emit_user self._emit_initial_user = emit_initial_user self._defer_message_ack = defer_message_ack self._lock = asyncio.Lock() self._outstanding_completions = 0 + self._interrupting = False async def start(self) -> None: """Open Dashboard input after the initial backend submission succeeds.""" @@ -64,6 +68,9 @@ async def run(self, driver: Any) -> None: async def complete(self, driver: Any) -> None: """Record one completed backend request and flush queued input.""" + if self._interrupting: + self._outstanding_completions = max(0, self._outstanding_completions - 1) + return async with self._lock: if self._interaction.planner_activity == "ended": return @@ -75,6 +82,8 @@ async def complete(self, driver: Any) -> None: async def tool_completed(self, driver: Any) -> None: """Flush input queued while the backend was running a tool.""" + if self._interrupting: + return async with self._lock: await self._flush(driver) @@ -94,7 +103,17 @@ def message_discarded(self, message_id: str) -> None: async def cancel_active_toolkit(self) -> None: """Cancel and drain the active toolkit operation off the event loop.""" - await asyncio.to_thread(self._cancel_active_and_wait) + await cancel_and_wait(self._cancel_active_and_wait) + + async def _interrupt_driver(self, driver: Any) -> int: + # SDK result events must drain while _process holds the input lock. + # They must not wait on that lock or flush new input during cancellation. + self._interrupting = True + try: + await self.cancel_active_toolkit() + return await driver.interrupt() + finally: + self._interrupting = False async def _process(self, driver: Any) -> None: async with self._lock: @@ -102,8 +121,7 @@ async def _process(self, driver: Any) -> None: return if self._interaction.task_replacement_requested: try: - await self.cancel_active_toolkit() - await driver.interrupt() + await self._interrupt_driver(driver) except Exception as exc: self._interaction.complete_task_replacement( error=f"planner interrupt failed: {_exception_text(exc)}" @@ -113,8 +131,10 @@ async def _process(self, driver: Any) -> None: return if self._interaction.claim_interrupt_request(): try: - await self.cancel_active_toolkit() - completed = await driver.interrupt() + completed = await self._interrupt_driver(driver) + if self._interaction.planner_activity == "ended": + return + self._resume_calls() self._outstanding_completions = max( 0, self._outstanding_completions - completed ) diff --git a/rpent/planner/api_loop.py b/rpent/planner/api_loop.py index c4d59f07b..fba9e22f0 100644 --- a/rpent/planner/api_loop.py +++ b/rpent/planner/api_loop.py @@ -57,6 +57,7 @@ from rpent.planner.base import ( REASONING_EFFORTS, PlannerResult, + cancel_and_wait, execute_tool, ) from rpent.tools.toolkit import Toolkit @@ -277,6 +278,14 @@ async def _await_next() -> str | None: last_error = _api_error_text(e, no_images=self._no_images) logger.error("agent run failed: %s", last_error) + finally: + try: + await cancel_and_wait(toolkit.cancel_active_and_wait) + except Exception as exc: + logger.exception("API toolkit cleanup failed") + last_error = last_error or f"Toolkit cleanup failed: {exc}" + messages.append({"role": "toolkit", "finish": toolkit.finish_result}) + return PlannerResult( finish_result=observer.finish_result, messages=messages, @@ -311,6 +320,7 @@ def emit_user(text: str, *, initial: bool = False) -> None: control = DashboardPlannerControl( interaction=interaction, cancel_active_and_wait=toolkit.cancel_active_and_wait, + resume_calls=toolkit.resume_calls, emit_user=emit_user, emit_initial_user=lambda: emit_user(user_message, initial=True), defer_message_ack=True, @@ -738,7 +748,6 @@ def _build_tools(toolkit: Toolkit, *, no_images: bool = False) -> list[Tool]: ), json_schema=tool.input_schema, takes_ctx=False, - sequential=True, ) for tool in tools ] diff --git a/rpent/planner/base.py b/rpent/planner/base.py index bfd6a8af7..4135a2d12 100644 --- a/rpent/planner/base.py +++ b/rpent/planner/base.py @@ -19,6 +19,8 @@ import asyncio import os import queue +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import TYPE_CHECKING, Protocol @@ -51,9 +53,25 @@ def strip_mcp_prefix(name: str) -> str: return name.removeprefix(MCP_TOOL_PREFIX) +async def cancel_and_wait(cancel: Callable[[], None]) -> None: + """Keep cancellation independent of a tool pool full of waiting calls.""" + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="tool-control") as pool: + await asyncio.get_running_loop().run_in_executor(pool, cancel) + + async def execute_tool(toolkit: Toolkit, name: str, arguments: dict) -> ToolResult: - """Dispatch native tool execution off the event loop.""" - return await asyncio.to_thread(toolkit.execute_tool, name, arguments) + """Retain the worker through cancellation, including queued executor jobs.""" + worker = asyncio.create_task( + asyncio.to_thread(toolkit.execute_tool, name, arguments) + ) + try: + return await asyncio.shield(worker) + except asyncio.CancelledError: + await cancel_and_wait(toolkit.cancel_active_and_wait) + # Jobs still in the executor queue must observe paused admission before + # the adapter can resume the next turn. + await asyncio.shield(worker) + raise class PlannerResult: diff --git a/rpent/planner/claude_code.py b/rpent/planner/claude_code.py index dc6e7e4d1..d39a08655 100644 --- a/rpent/planner/claude_code.py +++ b/rpent/planner/claude_code.py @@ -47,6 +47,7 @@ REASONING_EFFORTS, PlannerResult, add_mcp_prefix, + cancel_and_wait, strip_mcp_prefix, ) from rpent.planner.utils.http_mcp_server import build_mcp_server, list_mcp_tools @@ -230,13 +231,14 @@ async def consume_stream() -> None: options, recorder, input_queue, + toolkit=toolkit, emit=_emit, emit_user=_emit_user, ) except asyncio.TimeoutError: error = f"Claude Agent SDK timed out after {self._timeout_s}s" try: - await asyncio.to_thread(toolkit.cancel_active_and_wait) + await cancel_and_wait(toolkit.cancel_active_and_wait) except Exception as cancel_error: logger.warning( "failed to cancel toolkit work after Claude timeout: %s", @@ -257,6 +259,16 @@ async def consume_stream() -> None: _write_jsonl(raw_f, {"type": "error", "message": error}) logger.info(rendered.rstrip()) + finally: + try: + await cancel_and_wait(toolkit.cancel_active_and_wait) + except Exception as exc: + logger.exception("Claude toolkit cleanup failed") + error = error or f"Toolkit cleanup failed: {exc}" + _write_jsonl( + raw_f, {"type": "toolkit_finish", "finish": toolkit.finish_result} + ) + elapsed = time.time() - started text = "".join(rendered_chunks) or output_path.read_text(errors="replace") error = error or recorder.error @@ -287,6 +299,7 @@ async def _run_interactive( recorder: "_Recorder", input_queue, *, + toolkit: Toolkit, emit, emit_user, ) -> None: @@ -298,7 +311,9 @@ async def _run_interactive( sentinel (or ``/quit``) interrupts the run; the ``finish`` tool ends it normally. Because a human supervises, there is no wall-clock cap here. """ - adapter = _TerminalSessionAdapter(input_queue=input_queue, emit_user=emit_user) + adapter = _TerminalSessionAdapter( + toolkit=toolkit, input_queue=input_queue, emit_user=emit_user + ) driver = _ClaudeSessionDriver( sdk=sdk, options=options, @@ -324,6 +339,7 @@ async def _run_dashboard_session( adapter = _ClaudeDashboardAdapter( interaction=dashboard_interaction, cancel_active_and_wait=toolkit.cancel_active_and_wait, + resume_calls=toolkit.resume_calls, emit_user=emit_user, emit_initial_user=lambda: emit_user( initial_user_text, @@ -392,6 +408,10 @@ def __init__( self._recorder = recorder self._emit = emit self._client: Any | None = None + self._pending_results = 0 + self._turns_drained = asyncio.Event() + self._turns_drained.set() + self._interrupting = False async def run(self, prompt: str, adapter: Any) -> None: """Open one client, submit the first query, and run one input adapter.""" @@ -430,7 +450,15 @@ async def query(self, text: str) -> None: """Submit one user turn to the owned client.""" if self._client is None: raise RuntimeError("Claude session is not connected") - await self._client.query(text) + self._pending_results += 1 + self._turns_drained.clear() + try: + await self._client.query(text) + except BaseException: + self._pending_results -= 1 + if not self._pending_results: + self._turns_drained.set() + raise async def submit(self, text: str) -> int: """Submit Dashboard input as a new Claude query.""" @@ -441,34 +469,54 @@ async def interrupt(self) -> int: """Interrupt the owned client and suppress its expected error result.""" if self._client is None: raise RuntimeError("Claude session is not connected") + interrupted = self._pending_results + if not interrupted: + return 0 + self._interrupting = True self._recorder.suppress_next_result_error = True try: await self._client.interrupt() + # The SDK acknowledgement alone does not close old tool requests. + # The consumer signals before calling control hooks (which may be + # waiting for this interrupt to release the Dashboard lock). + await asyncio.wait_for(self._turns_drained.wait(), timeout=15) except BaseException: - # A failed interrupt must not hide a later unrelated SDK error. self._recorder.suppress_next_result_error = False raise - # Claude emits a ResultMessage for the interrupted query, so the - # matching completion is accounted for by the normal message path. - return 0 + finally: + self._interrupting = False + return interrupted async def _consume(self, adapter: Any) -> None: if self._client is None: raise RuntimeError("Claude session is not connected") async for message in self._client.receive_messages(): self._emit(message) + if _kind(message) == "ResultMessage": + self._pending_results = max(0, self._pending_results - 1) + if not self._pending_results: + self._turns_drained.set() # A successful finish tool result owns the boundary: end the # session without giving queued Dashboard input a chance to flush. if self._recorder.finish_result is not None: logger.info("FINISH called: %s", self._recorder.finish_result) return + if self._interrupting: + continue await adapter.on_message(self, message) class _TerminalSessionAdapter: """Preserve the terminal TUI's interrupt-then-query steering policy.""" - def __init__(self, *, input_queue: Any, emit_user) -> None: + def __init__( + self, + *, + toolkit: Toolkit, + input_queue: queue.Queue[str | None], + emit_user: Callable[[str], None], + ) -> None: + self._toolkit = toolkit self._input_queue = input_queue self._emit_user = emit_user @@ -477,18 +525,14 @@ async def initial_query_succeeded(self, driver: _ClaudeSessionDriver) -> None: async def run(self, driver: _ClaudeSessionDriver) -> None: while True: - nxt = await asyncio.to_thread(next_user_line, self._input_queue) - if nxt is None: - with contextlib.suppress(Exception): - await driver.interrupt() + user_text = await asyncio.to_thread(next_user_line, self._input_queue) + await cancel_and_wait(self._toolkit.cancel_active_and_wait) + await driver.interrupt() + if user_text is None or self._toolkit.finish_result is not None: return - self._emit_user(nxt) - # Keep the current terminal semantics: every steering line - # interrupts the in-flight turn, then enters the same session. - with contextlib.suppress(Exception): - await driver.interrupt() - with contextlib.suppress(Exception): - await driver.query(nxt) + self._toolkit.resume_calls() + self._emit_user(user_text) + await driver.query(user_text) async def on_message( self, @@ -510,12 +554,14 @@ def __init__( *, interaction: DashboardInteractionPort, cancel_active_and_wait: Callable[[], None], + resume_calls: Callable[[], None], emit_user: Callable[[str], None], emit_initial_user: Callable[[], None], ) -> None: self._control = DashboardPlannerControl( interaction=interaction, cancel_active_and_wait=cancel_active_and_wait, + resume_calls=resume_calls, emit_user=emit_user, emit_initial_user=emit_initial_user, ) diff --git a/rpent/planner/codex.py b/rpent/planner/codex.py index c7c4b9623..75b82eb22 100644 --- a/rpent/planner/codex.py +++ b/rpent/planner/codex.py @@ -51,6 +51,7 @@ from rpent.planner.base import ( REASONING_EFFORTS, PlannerResult, + cancel_and_wait, strip_mcp_prefix, ) from rpent.planner.utils.http_mcp_server import HttpMcpServer @@ -223,7 +224,13 @@ def solve( _write_jsonl(raw_f, {"type": "error", "message": error}) logger.info(rendered.rstrip()) finally: - mcp_server.stop() + try: + toolkit.cancel_active_and_wait() + except Exception as exc: + logger.exception("Codex toolkit cleanup failed") + error = error or f"Toolkit cleanup failed: {exc}" + finally: + mcp_server.stop() elapsed = time.time() - started text = state.get("text", "") or output_path.read_text(errors="replace") @@ -303,6 +310,7 @@ def _steer() -> None: if stop_steer.is_set(): return if nxt is None: + recorder.toolkit.cancel_active_and_wait() try: turn.interrupt() except Exception: @@ -420,6 +428,7 @@ def emit_user(text: str, *, initial: bool = False) -> None: control = DashboardPlannerControl( interaction=interaction, cancel_active_and_wait=toolkit.cancel_active_and_wait, + resume_calls=toolkit.resume_calls, emit_user=emit_user, emit_initial_user=lambda: emit_user( initial_user_text, initial=True @@ -459,7 +468,13 @@ def emit_user(text: str, *, initial: bool = False) -> None: {"type": "toolkit_finish", "finish": recorder.finish_result}, ) finally: - mcp_server.stop() + try: + await cancel_and_wait(toolkit.cancel_active_and_wait) + except Exception as exc: + logger.exception("Codex toolkit cleanup failed") + error = error or f"Toolkit cleanup failed: {exc}" + finally: + await asyncio.to_thread(mcp_server.stop) if recorder.final_response is not None: last_message_path.write_text(recorder.final_response) @@ -547,14 +562,7 @@ async def interrupt(self) -> int: try: await asyncio.wait_for(done.wait(), timeout=15) except asyncio.TimeoutError: - if self._turn_task is not None: - self._turn_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await self._turn_task - self._turn = None - self._turn_done = None - self._turn_task = None - return 1 + raise RuntimeError("Codex did not finish the interrupted turn") from None # ``_consume_turn`` reports the matching completed turn boundary. return 0 diff --git a/rpent/planner/utils/http_mcp_server.py b/rpent/planner/utils/http_mcp_server.py index 5e50825d4..057abd3bb 100644 --- a/rpent/planner/utils/http_mcp_server.py +++ b/rpent/planner/utils/http_mcp_server.py @@ -32,7 +32,6 @@ from __future__ import annotations -import asyncio import base64 import socket import threading @@ -88,7 +87,6 @@ def mcp_result(result: ToolResult) -> dict[str, Any]: def build_mcp_server(toolkit: Toolkit) -> Server: """Build the MCP service shared by Claude and Codex with native validation.""" mcp_app: Server = Server(SERVER_NAME, version="0.1.0") - tool_execution_lock = asyncio.Lock() exported_tools = list_mcp_tools(toolkit) exported_names = {tool.name for tool in exported_tools} @@ -111,8 +109,7 @@ async def _call_tool(name: str, arguments: dict[str, Any]) -> types.CallToolResu if lookup not in exported_names: result = ToolResult(error=f"Unknown tool: {lookup}") else: - async with tool_execution_lock: - result = await execute_tool(toolkit, lookup, arguments or {}) + result = await execute_tool(toolkit, lookup, arguments or {}) return types.CallToolResult(**mcp_result(result)) return mcp_app diff --git a/rpent/tools/base.py b/rpent/tools/base.py index 300086ace..33c550268 100644 --- a/rpent/tools/base.py +++ b/rpent/tools/base.py @@ -141,7 +141,7 @@ def check_cancelled(self) -> None: class Tool(Generic[ParamsT]): """A handler and its generated parameter model, fixed before execution. - readonly skips automatic observation capture. + readonly allows shared execution and skips automatic observation capture. """ name: str @@ -193,7 +193,7 @@ def _parameter_model( def readonly(handler: Callable[ParamsT, ToolResult]) -> Callable[ParamsT, ToolResult]: """Skip automatic environment observation capture; file writes remain allowed. - Place this marker below @tool. Calls execute one at a time. + Place this marker below @tool. Readonly calls can execute concurrently. """ setattr(handler, "_rpent_readonly", True) return handler @@ -207,7 +207,8 @@ def tool(function: Callable[ParamsT, ToolResult], /) -> Tool[ParamsT]: Every handler declares a required keyword-only ctx, injected by the executor and excluded from the schema. Place @tool above @readonly. By default calls are exclusive; robot tools - other than finish capture observations afterward. @readonly disables capture. + other than finish capture observations afterward. @readonly allows shared + execution and disables capture. """ # Resolve annotations in factories/tests as well as at module scope, without # retaining the caller's frame or namespace in the resulting Tool. diff --git a/rpent/tools/toolkit.py b/rpent/tools/toolkit.py index 83d8d2720..0220b9065 100644 --- a/rpent/tools/toolkit.py +++ b/rpent/tools/toolkit.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Native tool execution with one active invocation per toolkit.""" +"""Native tool execution, per-toolkit scheduling, and lifecycle hooks.""" from __future__ import annotations @@ -21,7 +21,7 @@ import time from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Generic +from typing import Any, Generic, Literal import numpy as np from pydantic import ValidationError @@ -45,17 +45,103 @@ logger = get_logger("tools") -@dataclass(slots=True) -class _ToolOperation: +@dataclass(eq=False) +class _Call: + """Scheduling state for one call; arguments stay with the executor.""" + + tool: Tool cancel_event: threading.Event = field(default_factory=threading.Event) - done_event: threading.Event = field(default_factory=threading.Event) + done: bool = False + + +class _Scheduler: + """Run readonly tools together and exclusive tools one at a time. + + Exclusive calls take priority and keep their queue order. Calls stay active + through capture; the condition lock only protects scheduling state, never + tool execution. + """ + + def __init__(self) -> None: + self._condition = threading.Condition() + self._pending: list[_Call] = [] + self._active_calls: set[_Call] = set() + self._state: Literal["open", "paused", "closed"] = "open" + + def acquire(self, tool: Tool) -> _Call: + with self._condition: + if self._state == "closed": + raise RuntimeError("Toolkit is closed.") + if self._state == "paused": + raise RuntimeError("Tool calls are paused.") + call = _Call(tool) + self._pending.append(call) + try: + while True: + if call.cancel_event.is_set(): + raise RuntimeError("Tool call cancelled.") + if self._can_start(call): + self._pending.remove(call) + self._active_calls.add(call) + return call + self._condition.wait() + except BaseException: + if call in self._pending: + self._pending.remove(call) + call.done = True + self._condition.notify_all() + raise + + def _can_start(self, call: _Call) -> bool: + first_exclusive = next( + (pending for pending in self._pending if not pending.tool.readonly), None + ) + + if not call.tool.readonly: + return not self._active_calls and call is first_exclusive + + return first_exclusive is None and all( + active.tool.readonly for active in self._active_calls + ) + + def release(self, call: _Call) -> None: + with self._condition: + self._active_calls.remove(call) + call.done = True + self._condition.notify_all() + + def cancel(self, *, close: bool = False) -> list[_Call]: + """Stop admission and signal queued and active calls without waiting.""" + with self._condition: + if close: + self._state = "closed" + elif self._state != "closed": + self._state = "paused" + calls = [*self._pending, *self._active_calls] + self._pending.clear() + for call in calls: + call.cancel_event.set() + self._condition.notify_all() + return calls + + def cancel_and_wait(self, *, close: bool = False) -> None: + with self._condition: + calls = self.cancel(close=close) + self._condition.wait_for(lambda: all(call.done for call in calls)) + + def resume(self) -> None: + with self._condition: + if self._active_calls: + raise RuntimeError("Wait for call cleanup before resuming.") + if self._state == "paused": + self._state = "open" class Toolkit(Generic[RobotT]): """A fixed tool collection and its execution resources for one planner session. - Robot tools submit RGB frames through ctx.record_frame(). Robot toolkits - save their action clips, episode video, and replay recipes. + Robot tools submit RGB frames through ctx.record_frame(). This toolkit owns + the frame buffer and saves action clips and the episode video. """ def __init__( @@ -76,9 +162,9 @@ def __init__( self._tools: dict[str, Tool] = { item.name: item for item in (*COMMON_TOOLS, *tools) } - self._operation_lock = threading.Lock() - self._active_operation: _ToolOperation | None = None + self._scheduler = _Scheduler() self._finish_result: dict[str, Any] | None = None + self._recipe_commands: list[dict[str, Any]] = [] self._frames: list[np.ndarray] = [] @property @@ -103,7 +189,7 @@ def record_frame(self, rgb: np.ndarray) -> None: self._frames.append(np.ascontiguousarray(np.asarray(rgb))) def execute_tool(self, name: str, arguments: dict) -> ToolResult: - """Validate and execute one call, then capture its observation.""" + """Validate, wait, execute, and capture before releasing the call.""" tool = self._tools.get(name) if tool is None: return ToolResult(error=f"Unknown tool: {name}") @@ -115,14 +201,13 @@ def execute_tool(self, name: str, arguments: dict) -> ToolResult: ) details = json.dumps({"errors": errors}) return ToolResult(error=f"Invalid arguments for {name}.\n{details}") - with self._operation_lock: - if self._active_operation is not None: - return ToolResult(error="another tool operation is still active") - operation = _ToolOperation() - self._active_operation = operation + try: + call = self._scheduler.acquire(tool) + except RuntimeError as exc: + return ToolResult(error=str(exc)[:500]) try: - if operation.cancel_event.is_set(): + if call.cancel_event.is_set(): return ToolResult(error="Tool call cancelled.") ctx = ToolContext( state=self._state, @@ -130,11 +215,13 @@ def execute_tool(self, name: str, arguments: dict) -> ToolResult: robot=self._robot, output_dir=self._task_output_dir, record_frame=self.record_frame, - _cancel_event=operation.cancel_event, + _cancel_event=call.cancel_event, ) capture = ( not tool.readonly and tool not in COMMON_TOOLS and name != "finish" ) + if capture and self._dashboard_events.enabled: + frame_start = len(self._frames) started = time.perf_counter() # Read fields directly so nested models reach the handler intact. kwargs = {name: getattr(args, name) for name in type(args).model_fields} @@ -146,6 +233,7 @@ def execute_tool(self, name: str, arguments: dict) -> ToolResult: if capture: elapsed_s = time.perf_counter() - started previous = self._state.latest_record() + observation_data = None try: observation_data, observation_images = self._capture_observation( command={"action": tool.name, **args.model_dump()}, @@ -163,6 +251,26 @@ def execute_tool(self, name: str, arguments: dict) -> ToolResult: result.error = error record = self._state.latest_record() if record is not None and record is not previous: + if self._dashboard_events.enabled: + try: + frames = self._frames[frame_start:] + if frames: + self._state.save( + f"action_{tool.name}.mp4", + frames, + step=record.step_idx, + fps=20, + ) + if observation_data is not None: + observation_data["artifacts"] = sorted( + record.artifacts + ) + except Exception as exc: + logger.warning( + "failed to save action clip for step %s: %s", + record.step_idx, + exc, + ) try: self._dashboard_events.emit( StepRecordEvent(record=record, env_state=self._state) @@ -190,12 +298,14 @@ def execute_tool(self, name: str, arguments: dict) -> ToolResult: self._finish_result = { key: value for key, value in result.data.items() if key != "_finish" } + elif tool not in COMMON_TOOLS and not result.is_error: + self._recipe_commands.append( + {"action": tool.name, **args.model_dump(mode="json")} + ) logger.info("Tool %s result: %s", name, result_text) return result finally: - with self._operation_lock: - self._active_operation = None - operation.done_event.set() + self._scheduler.release(call) def _capture_observation( self, *, command: dict[str, Any], result: ToolResult, elapsed_s: float @@ -204,27 +314,43 @@ def _capture_observation( Returned data replaces the action's data and must include the recorded step and its artifact names. Include action details in that data where - needed (e.g. log.result). Robot toolkits save action video artifacts. + needed (e.g. log.result). The executor adds action videos to the artifacts. The executor appends the images and retains the action's error. Raise if capture fails; an already saved step is still published to the Dashboard. """ raise NotImplementedError("This toolkit does not capture robot observations.") def cancel_active_and_wait(self) -> None: - """Request cancellation and wait for the active tool to return.""" - with self._operation_lock: - operation = self._active_operation - if operation is None: - return - operation.cancel_event.set() - operation.done_event.wait() + self._scheduler.cancel_and_wait() + + def resume_calls(self) -> None: + self._scheduler.resume() def close(self) -> None: - """Release robot resources at the end of a run. Default: no-op.""" + """Called once by the runner to drain calls and save the episode video.""" + self._scheduler.cancel_and_wait(close=True) + frames = self._frames + self._frames = [] + try: + if frames: + self._state.save("episode.mp4", frames, step=None, fps=20) + except Exception as exc: + logger.warning("failed to save episode video: %s", exc) def solved(self) -> bool: raise NotImplementedError - def write_recipe(self, recipe_tag: str) -> str | None: - """Write a replay recipe for this robot, if supported.""" - return None + def write_recipe(self, recipe_tag: str) -> str: + """Export successful robot action and perception calls in completion order. + + Keep the full session, including resets. File tools and finish are + excluded. The caller decides whether the task qualifies for publication. + """ + name = f"{recipe_tag}_recipe.jsonl" + path = self._task_output_dir / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "".join(json.dumps(command) + "\n" for command in self._recipe_commands), + encoding="utf-8", + ) + return name diff --git a/tests/unit_tests/robots/dual_franka/test_dual_franka_tools.py b/tests/unit_tests/robots/dual_franka/test_dual_franka_tools.py index 5b2b3144e..0c746c70c 100644 --- a/tests/unit_tests/robots/dual_franka/test_dual_franka_tools.py +++ b/tests/unit_tests/robots/dual_franka/test_dual_franka_tools.py @@ -660,3 +660,24 @@ def test_normal_return_contract(case, tmp_path): "operator_verdict": "success", **case["arguments"], } + + +def test_toolkit_cancellation_pauses_until_resume(tmp_path): + env = FakeEnv() + toolkit = DualFrankaToolkit( + runtime_kwargs={"env": env, "model": None, "task_description": "default task"}, + dashboard_events=NullDashboardEventSink(), + memory=MemoryManager(tmp_path / "memory"), + state_output_dir=tmp_path, + ) + try: + toolkit.cancel_active_and_wait() + assert toolkit.execute_tool( + "move_delta", {"arm": "left", "delta_xyz": [0, 0, 0]} + ).is_error + assert env.moves == [] + toolkit.resume_calls() + assert not toolkit.execute_tool("view_env_state", {}).is_error + assert "finish" in {item.name for item in toolkit.list_tools()} + finally: + toolkit.close() diff --git a/tests/unit_tests/robots/franka/test_tools.py b/tests/unit_tests/robots/franka/test_tools.py index 59a2ec426..1c2541cf6 100644 --- a/tests/unit_tests/robots/franka/test_tools.py +++ b/tests/unit_tests/robots/franka/test_tools.py @@ -178,6 +178,11 @@ def test_toolkit_factory_validation_and_capture(tmp_path, monkeypatch): ) assert read.images == result.images assert len(events) == 2 + toolkit.cancel_active_and_wait() + assert toolkit.execute_tool("move_delta", {"delta_xyz": [0, 0, 0]}).is_error + assert len(env.moves) == 1 + toolkit.resume_calls() + assert not toolkit.execute_tool("view_env_state", {}).is_error assert "finish" in {tool.name for tool in toolkit.list_tools()} assert toolkit.finish_result is None toolkit.close() diff --git a/tests/unit_tests/robots/libero/conftest.py b/tests/unit_tests/robots/libero/conftest.py index 6a71f056e..7d922dd97 100644 --- a/tests/unit_tests/robots/libero/conftest.py +++ b/tests/unit_tests/robots/libero/conftest.py @@ -135,4 +135,5 @@ def make(*, mode="evaluation", attempts=0, molmo_client=None, flywheel_config=No for toolkit in instances: # Avoid video encoding in CPU contract tests. toolkit._frames.clear() - toolkit.close() + if toolkit._scheduler._state != "closed": + toolkit.close() diff --git a/tests/unit_tests/robots/libero/test_libero_integration.py b/tests/unit_tests/robots/libero/test_libero_integration.py index cff521ec1..34e77937b 100644 --- a/tests/unit_tests/robots/libero/test_libero_integration.py +++ b/tests/unit_tests/robots/libero/test_libero_integration.py @@ -183,12 +183,15 @@ def test_close_handles_collection_and_video_independently( ) toolkit._robot = SimpleNamespace(finalize_flywheel=finalize) toolkit._frames = frames + toolkit._scheduler = Mock() toolkit._state = SimpleNamespace(save=save) logger = Mock() monkeypatch.setattr(libero_toolkit, "logger", logger) + monkeypatch.setattr("rpent.tools.toolkit.logger", logger) toolkit.close() + toolkit._scheduler.cancel_and_wait.assert_called_once_with(close=True) finalize.assert_called_once_with() save.assert_called_once_with("episode.mp4", frames, step=None, fps=20) if failure == "finalize": diff --git a/tests/unit_tests/robots/libero/test_libero_toolkit_contracts.py b/tests/unit_tests/robots/libero/test_libero_toolkit_contracts.py index dbfb92905..9c6cb90cd 100644 --- a/tests/unit_tests/robots/libero/test_libero_toolkit_contracts.py +++ b/tests/unit_tests/robots/libero/test_libero_toolkit_contracts.py @@ -172,7 +172,8 @@ def after_step(): assert stepped.wait(3) cancellation = pool.submit(toolkit.cancel_active_and_wait) # Confirm cancellation was delivered before allowing another action boundary. - assert toolkit._active_operation.cancel_event.wait(3) + for call in list(toolkit._scheduler._active_calls): + assert call.cancel_event.wait(3) release_step.set() result = action.result(3) cancellation.result(3) @@ -180,6 +181,7 @@ def after_step(): assert len(env.actions) == toolkit._robot.executed_steps == completed assert toolkit.state.latest_record().result == {"error": result.error} assert result.data["step"] == 1 + toolkit.resume_calls() env.after_step = lambda: None assert not toolkit.execute_tool("set_gripper", {"steps": 1}).is_error assert len(env.actions) == completed + 1 diff --git a/tests/unit_tests/robots/libero/test_recipe.py b/tests/unit_tests/robots/libero/test_recipe.py new file mode 100644 index 000000000..5c548e394 --- /dev/null +++ b/tests/unit_tests/robots/libero/test_recipe.py @@ -0,0 +1,94 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + + +def export(toolkit): + name = toolkit.write_recipe("cell") + assert name == "cell_recipe.jsonl" + return [ + json.loads(line) + for line in (toolkit._task_output_dir / name).read_text().splitlines() + ] + + +def test_common_files_finish_validation_and_execution_errors_are_excluded( + make_toolkit, monkeypatch +): + toolkit, env, _ = make_toolkit() + path = str(toolkit._task_output_dir / "note.txt") + for name, arguments in [ + ("write_text_file", {"path": path, "content": "note"}), + ("read_text_file", {"path": path}), + ("list_dir", {}), + ("read_image", {"name": "agentview_policy.png", "step": 0}), + ]: + assert not toolkit.execute_tool(name, arguments).is_error + for name, arguments in [ + ("unknown", {}), + ("set_gripper", {"steps": "invalid"}), + ("segment", {}), + ]: + assert toolkit.execute_tool(name, arguments).is_error + + def fail(action): + raise RuntimeError() + + with monkeypatch.context() as patch: + patch.setattr(env, "step", fail) + failed = toolkit.execute_tool("set_gripper", {"steps": 1}) + assert failed.is_error and failed.error == "" + assert not toolkit.execute_tool("set_gripper", {"steps": 1}).is_error + assert not toolkit.execute_tool( + "finish", {"status": "success", "summary": "done"} + ).is_error + assert [call["action"] for call in export(toolkit)] == ["set_gripper"] + + +def test_reset_is_kept_with_both_attempts_and_refused_finish_is_excluded(make_toolkit): + toolkit, _, _ = make_toolkit(mode="exploration", attempts=2) + assert toolkit.execute_tool( + "finish", {"status": "stuck", "summary": "first"} + ).is_error + for name, arguments in [ + ("segment", {"prompt": "bowl"}), + ("set_gripper", {"steps": 1}), + ("reset", {"reason": "try again"}), + ("view_env_state", {}), + ("segment", {"prompt": "lid"}), + ("set_gripper", {"steps": 1}), + ]: + result = toolkit.execute_tool(name, arguments) + assert not result.is_error, result.error + recipe = export(toolkit) + assert recipe[1] == {"action": "set_gripper", "gripper": -1.0, "steps": 1} + assert [call["action"] for call in recipe] == [ + "segment", + "set_gripper", + "reset", + "view_env_state", + "segment", + "set_gripper", + ] + + +def test_recipe_does_not_infer_tool_errors_from_environment_success(make_toolkit): + toolkit, env, _ = make_toolkit() + result = toolkit.execute_tool( + "pi0_doubled", {"prompt": "touch bowl", "max_chunks": 1} + ) + assert not result.is_error and result.data["log"]["result"]["success"] is False + assert not env.terminated + assert [call["action"] for call in export(toolkit)] == ["pi0_doubled"] diff --git a/tests/unit_tests/robots/libero/test_recording.py b/tests/unit_tests/robots/libero/test_recording.py new file mode 100644 index 000000000..c937d70ac --- /dev/null +++ b/tests/unit_tests/robots/libero/test_recording.py @@ -0,0 +1,143 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import numpy as np +import pytest + + +class Events: + enabled = True + + def __init__(self): + self.records = [] + + def emit(self, event): + self.records.append(event.record) + + +@pytest.fixture +def recording(make_toolkit, monkeypatch): + toolkit, env, _ = make_toolkit() + events = Events() + toolkit._dashboard_events = events + saved = [] + original_save = toolkit.state.save + + def save(name, value, **kwargs): + if name.endswith(".mp4"): + saved.append( + (name, kwargs["step"], [int(frame[0, 0, 0]) for frame in value]) + ) + value = b"video" + return original_save(name, value, **kwargs) + + monkeypatch.setattr(toolkit.state, "save", save) + return toolkit, env, events, saved + + +def image(env, value): + env.image = np.full((8, 8, 3), value, dtype=np.uint8) + + +def test_each_action_saves_its_frames_and_updates_response_and_dashboard(recording): + toolkit, env, events, saved = recording + image(env, 1) + first = toolkit.execute_tool("set_gripper", {"steps": 2}) + image(env, 2) + second = toolkit.execute_tool("set_gripper", {"steps": 3}) + assert saved == [ + ("action_set_gripper.mp4", 1, [1, 1]), + ("action_set_gripper.mp4", 2, [2, 2, 2]), + ] + for result, event in zip([first, second], events.records): + assert not result.is_error + assert "action_set_gripper.mp4" in result.data["artifacts"] + assert result.data["artifacts"] == sorted(event.artifacts) + observed = toolkit.execute_tool("view_env_state", {"step": result.data["step"]}) + assert observed.data == { + key: value for key, value in result.data.items() if key != "agent_elapsed_s" + } + assert observed.images == result.images + assert [int(frame[0, 0, 0]) for frame in toolkit._frames] == [1, 1, 2, 2, 2] + + +def test_zero_step_and_readonly_calls_do_not_save_clips(recording): + toolkit, _, _, saved = recording + for name, args in [ + ("set_gripper", {"steps": 0}), + ("view_env_state", {}), + ("finish", {"status": "stuck", "summary": "done"}), + ]: + result = toolkit.execute_tool(name, args) + assert not result.is_error + assert not any( + name.endswith(".mp4") for name in result.data.get("artifacts", []) + ) + assert saved == [] + + +@pytest.mark.parametrize("record_saved", [False, True]) +def test_capture_failure_never_mixes_action_frames( + recording, monkeypatch, record_saved +): + toolkit, env, events, saved = recording + + def fail(*args, **kwargs): + raise RuntimeError("observation failure") + + image(env, 1) + with monkeypatch.context() as patch: + if record_saved: + patch.setattr("robots.libero.toolkit.build_observation", fail) + else: + patch.setattr(env, "raw_obs", fail) + failed = toolkit.execute_tool("set_gripper", {"steps": 2}) + assert failed.is_error and "State capture failed:" in failed.error + recipe = toolkit._task_output_dir / toolkit.write_recipe("failed") + assert recipe.read_text() == "" + assert "step" not in failed.data + if record_saved: + assert saved == [("action_set_gripper.mp4", 1, [1, 1])] + assert "action_set_gripper.mp4" in events.records[-1].artifacts + else: + assert saved == [] + assert events.records == [] + + image(env, 2) + succeeded = toolkit.execute_tool("set_gripper", {"steps": 1}) + assert not succeeded.is_error + assert saved[-1] == ("action_set_gripper.mp4", 2 if record_saved else 1, [2]) + + +@pytest.mark.parametrize("raises", [False, True]) +def test_video_save_failure_keeps_action_success_and_does_not_add_artifact( + recording, monkeypatch, raises +): + toolkit, _, events, saved = recording + original_save = toolkit.state.save + + def fail_video(name, value, **kwargs): + if name.endswith(".mp4"): + if raises: + raise OSError("video failure") + return None + return original_save(name, value, **kwargs) + + monkeypatch.setattr(toolkit.state, "save", fail_video) + result = toolkit.execute_tool("set_gripper", {"steps": 1}) + assert not result.is_error + assert "action_set_gripper.mp4" not in result.data["artifacts"] + assert result.data["artifacts"] == sorted(events.records[-1].artifacts) + assert saved == [] diff --git a/tests/unit_tests/robots/robocasa/conftest.py b/tests/unit_tests/robots/robocasa/conftest.py index f5cbb3760..0c9ac9096 100644 --- a/tests/unit_tests/robots/robocasa/conftest.py +++ b/tests/unit_tests/robots/robocasa/conftest.py @@ -196,4 +196,5 @@ def make(**settings): yield make for run in runs: run.toolkit._frames.clear() - run.toolkit.close() + if run.toolkit._scheduler._state != "closed": + run.toolkit.close() diff --git a/tests/unit_tests/robots/robocasa/test_robocasa_toolkit_contracts.py b/tests/unit_tests/robots/robocasa/test_robocasa_toolkit_contracts.py index 5b2599ac4..a213e6bb8 100644 --- a/tests/unit_tests/robots/robocasa/test_robocasa_toolkit_contracts.py +++ b/tests/unit_tests/robots/robocasa/test_robocasa_toolkit_contracts.py @@ -152,7 +152,11 @@ def on_step(): assert entered.wait(5) # Request cancellation while the first step is in progress. stop = pool.submit(run.toolkit.cancel_active_and_wait) - assert run.toolkit._active_operation.cancel_event.wait(5) + # Wait for the scheduler to signal cancellation before releasing the step. + with run.toolkit._scheduler._condition: + assert run.toolkit._scheduler._condition.wait_for( + lambda: run.toolkit._scheduler._state == "paused", timeout=5 + ) resume.set() result = action.result(timeout=5) stop.result(timeout=5) @@ -161,6 +165,7 @@ def on_step(): assert len(run.env.actions) == 1 assert result.data["task_progress"]["steps"] == 1 run.env.on_step = lambda: None + run.toolkit.resume_calls() assert not run.toolkit.execute_tool("release", {"steps": 1}).is_error assert len(run.env.actions) == 2 diff --git a/tests/unit_tests/robots/robotwin/conftest.py b/tests/unit_tests/robots/robotwin/conftest.py index 217924abb..fb68b49f0 100644 --- a/tests/unit_tests/robots/robotwin/conftest.py +++ b/tests/unit_tests/robots/robotwin/conftest.py @@ -155,4 +155,5 @@ def robotwin(tmp_path): env=env, model=model, toolkit=toolkit, output_dir=tmp_path / "run" ) toolkit._frames.clear() - toolkit.close() + if toolkit._scheduler._state != "closed": + toolkit.close() diff --git a/tests/unit_tests/rpent/dashboard/test_planner_control_contracts.py b/tests/unit_tests/rpent/dashboard/test_planner_control_contracts.py index 80ac0056a..337b25858 100644 --- a/tests/unit_tests/rpent/dashboard/test_planner_control_contracts.py +++ b/tests/unit_tests/rpent/dashboard/test_planner_control_contracts.py @@ -222,6 +222,7 @@ def cancel_active_and_wait() -> None: return DashboardPlannerControl( interaction=interaction, cancel_active_and_wait=cancel_active_and_wait, + resume_calls=lambda: events.append("toolkit-resume"), emit_user=lambda text: events.append(f"user:{text}"), emit_initial_user=lambda: events.append("initial-user"), defer_message_ack=defer_message_ack, @@ -346,6 +347,7 @@ async def scenario() -> None: assert events[:2] == ["initial-user", "toolkit-cancel"] assert events[2:] == [ "backend-interrupt", + "toolkit-resume", "submit:continue after interrupt", "user:continue after interrupt", ] @@ -425,3 +427,31 @@ def test_end_seals_pending_messages_as_unsent() -> None: assert interaction.activity == "ended" assert interaction.messages[0].status == "unsent" + + +def test_interrupt_allows_sdk_result_events_to_drain_without_flushing_input() -> None: + interaction = FakeInteraction() + interaction.submit("after", "new turn") + interaction.end_on_wait = True + events: list[str] = [] + control = make_control(interaction, events) + driver = FakeDriver(events) + + async def interrupt() -> int: + events.append("backend-interrupt") + await control.tool_completed(driver) + await control.complete(driver) + events.append("backend-drained") + assert not any(event.startswith("submit:") for event in events) + return 0 + + driver.interrupt = interrupt + + async def scenario() -> None: + await control.start() + interaction.interrupt_requested = True + await asyncio.wait_for(control.run(driver), timeout=2) + + asyncio.run(scenario()) + assert events.index("backend-drained") < events.index("toolkit-resume") + assert events.index("toolkit-resume") < events.index("submit:new turn") diff --git a/tests/unit_tests/rpent/planner/conftest.py b/tests/unit_tests/rpent/planner/conftest.py index e9d7d3db5..a0796a05d 100644 --- a/tests/unit_tests/rpent/planner/conftest.py +++ b/tests/unit_tests/rpent/planner/conftest.py @@ -28,4 +28,5 @@ def make(*args, **kwargs): yield make for toolkit in instances: - toolkit.close() + if toolkit._scheduler._state != "closed": + toolkit.close() diff --git a/tests/unit_tests/rpent/planner/test_api_contracts.py b/tests/unit_tests/rpent/planner/test_api_contracts.py index d3dd57230..81e2529a0 100644 --- a/tests/unit_tests/rpent/planner/test_api_contracts.py +++ b/tests/unit_tests/rpent/planner/test_api_contracts.py @@ -227,7 +227,7 @@ def test_tool_schema_and_dispatch_are_mapped_to_pydantic_ai( *[tool.name for tool in toolkit.list_tools() if tool.name != "read_image"], ] assert "read_image" in [tool.name for tool in tools] - assert all(tool.sequential for tool in tools) + assert all(not tool.sequential for tool in tools) finish = next(tool for tool in tools if tool.name == "finish") assert finish.function_schema.json_schema == next( tool.input_schema for tool in toolkit.list_tools() if tool.name == "finish" diff --git a/tests/unit_tests/rpent/planner/test_execution_lifecycle.py b/tests/unit_tests/rpent/planner/test_execution_lifecycle.py new file mode 100644 index 000000000..8c4e91dd8 --- /dev/null +++ b/tests/unit_tests/rpent/planner/test_execution_lifecycle.py @@ -0,0 +1,138 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace + +import pytest + +from rpent.memory import MemoryManager +from rpent.planner.base import execute_tool +from rpent.session import EnvState +from rpent.tools import ToolContext, Toolkit, ToolResult, readonly, tool + + +@tool +def finish(status: str, summary: str, *, ctx: ToolContext) -> ToolResult: + """Accept the requested outcome for this test toolkit.""" + return ToolResult(data={"_finish": True, "status": status, "summary": summary}) + + +@tool +@readonly +def moving(*, ctx: ToolContext) -> ToolResult: + """Wait at a cooperative cancellation boundary.""" + ctx.robot.entered() + assert ctx._cancel_event.wait(3) + ctx.check_cancelled() + raise AssertionError("cancelled handler continued") + + +def test_cancel_drains_jobs_even_when_default_tool_pool_is_full(tmp_path): + async def scenario(): + loop = asyncio.get_running_loop() + entered = asyncio.Event() + queued = asyncio.Event() + + class ObservedExecutor(ThreadPoolExecutor): + submissions = 0 + + def submit(self, *args, **kwargs): + future = super().submit(*args, **kwargs) + self.submissions += 1 + if self.submissions == 2: + queued.set() + return future + + loop.set_default_executor(ObservedExecutor(max_workers=1)) + toolkit = Toolkit( + state=EnvState(tmp_path), + memory=MemoryManager(tmp_path / "memory"), + robot=SimpleNamespace( + entered=lambda: loop.call_soon_threadsafe(entered.set) + ), + output_dir=tmp_path, + tools=(finish, moving), + ) + try: + active = asyncio.create_task(execute_tool(toolkit, "moving", {})) + await asyncio.wait_for(entered.wait(), timeout=3) + pending = asyncio.create_task(execute_tool(toolkit, "moving", {})) + await asyncio.wait_for(queued.wait(), timeout=3) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout=3) + assert ( + await asyncio.wait_for(active, timeout=3) + ).error == "Tool call cancelled." + assert ( + toolkit.execute_tool("list_dir", {}).error == "Tool calls are paused." + ) + toolkit.resume_calls() + assert not toolkit.execute_tool("list_dir", {}).is_error + finally: + toolkit.close() + + asyncio.run(scenario()) + + +def test_request_cancellation_waits_for_capture_before_returning(tmp_path): + async def scenario(): + loop = asyncio.get_running_loop() + capturing = asyncio.Event() + cancelled = asyncio.Event() + release = threading.Event() + + @tool + def action(*, ctx: ToolContext) -> ToolResult: + """Perform an action before its final observation.""" + return ToolResult() + + class CapturingToolkit(Toolkit): + def _capture_observation(self, **kwargs): + loop.call_soon_threadsafe(capturing.set) + assert release.wait(4) + return {"step": 0}, [] + + def cancel_active_and_wait(self): + loop.call_soon_threadsafe(cancelled.set) + super().cancel_active_and_wait() + + toolkit = CapturingToolkit( + state=EnvState(tmp_path), + memory=MemoryManager(tmp_path / "memory"), + robot=None, + output_dir=tmp_path, + tools=(finish, action), + ) + request = asyncio.create_task(execute_tool(toolkit, "action", {})) + try: + await asyncio.wait_for(capturing.wait(), timeout=3) + request.cancel() + await asyncio.wait_for(cancelled.wait(), timeout=3) + assert not request.done() + with pytest.raises(RuntimeError, match="cleanup"): + toolkit.resume_calls() + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(request, timeout=3) + toolkit.resume_calls() + assert not toolkit.execute_tool("list_dir", {}).is_error + finally: + release.set() + toolkit.close() + + asyncio.run(scenario()) diff --git a/tests/unit_tests/rpent/planner/test_http_mcp_server.py b/tests/unit_tests/rpent/planner/test_http_mcp_server.py index 62180a05e..d90eadcf4 100644 --- a/tests/unit_tests/rpent/planner/test_http_mcp_server.py +++ b/tests/unit_tests/rpent/planner/test_http_mcp_server.py @@ -37,11 +37,11 @@ def finish(status: str, summary: str, *, ctx: ToolContext) -> ToolResult: return ToolResult(data={"_finish": True, "status": status, "summary": summary}) -def test_http_serialized_calls_keep_native_validation(tmp_path: Path) -> None: +def test_http_shared_calls_overlap_and_keep_native_validation(tmp_path: Path) -> None: toolkit = Toolkit( state=EnvState(tmp_path), memory=MemoryManager(tmp_path / "memory"), - robot=threading.Barrier(1), + robot=threading.Barrier(2), output_dir=tmp_path, tools=(finish, read), ) diff --git a/tests/unit_tests/rpent/planner/test_native_adapters.py b/tests/unit_tests/rpent/planner/test_native_adapters.py index 03208d3cf..2987d8024 100644 --- a/tests/unit_tests/rpent/planner/test_native_adapters.py +++ b/tests/unit_tests/rpent/planner/test_native_adapters.py @@ -16,22 +16,30 @@ import asyncio import json +import threading import pytest from mcp import types +from pydantic_ai.messages import ModelResponse, ToolCallPart, ToolReturnPart +from pydantic_ai.models.function import FunctionModel from rpent.dashboard.events import NullDashboardEventSink +from rpent.memory import MemoryManager from rpent.planner.api_loop import ( + ApiAgentLoop, _ApiRunObserver, _build_tools, ) from rpent.planner.claude_code import ( _build_rpent_server, + _ClaudeSessionDriver, ) from rpent.planner.claude_code import _Recorder as ClaudeRecorder from rpent.planner.codex import _Recorder as CodexRecorder +from rpent.session import EnvState +from rpent.tools import ToolContext, Toolkit, ToolResult, tool -from ._native_helpers import call_sdk_tool +from ._native_helpers import call_sdk_tool, read def test_read_image_is_available_to_api_and_not_callable_through_mcp(make_toolkit): @@ -90,6 +98,123 @@ def test_recorders_read_verified_finish_without_a_matching_provider_event( assert toolkit.finish_result["operator_aborted"] is True +@tool +def finish(status: str, summary: str, *, ctx: ToolContext) -> ToolResult: + """Accept the requested outcome for this test toolkit.""" + return ToolResult(data={"_finish": True, "status": status, "summary": summary}) + + +def test_claude_sdk_dispatch_is_parallel(tmp_path): + toolkit = Toolkit( + state=EnvState(tmp_path), + memory=MemoryManager(tmp_path / "memory"), + robot=threading.Barrier(2), + output_dir=tmp_path, + tools=(finish, read), + ) + + async def scenario(): + server = _build_rpent_server(toolkit=toolkit) + results = await asyncio.gather( + call_sdk_tool(server, "read", {"number": "1"}), + call_sdk_tool(server, "read", {"number": "2"}), + ) + assert all(not result.isError for result in results) + assert [json.loads(result.content[0].text)["number"] for result in results] == [ + 1, + 2, + ] + + try: + asyncio.run(scenario()) + finally: + toolkit.close() + + +def test_api_sdk_dispatch_is_parallel_and_accepts_verified_finish(tmp_path): + toolkit = Toolkit( + state=EnvState(tmp_path), + memory=MemoryManager(tmp_path / "memory"), + robot=threading.Barrier(2), + output_dir=tmp_path, + tools=(finish, read), + ) + + def model(messages, info): + returned = [ + part + for message in messages + for part in message.parts + if isinstance(part, ToolReturnPart) + ] + if not returned: + return ModelResponse( + parts=[ + ToolCallPart("read", {"number": "1"}, "r1"), + ToolCallPart("read", {"number": 2}, "r2"), + ] + ) + assert all("error" not in json.loads(part.content) for part in returned) + return ModelResponse( + parts=[ + ToolCallPart( + "finish", {"status": "failure", "summary": "complete"}, "f" + ) + ] + ) + + result = ApiAgentLoop( + FunctionModel(model), dashboard_events=NullDashboardEventSink() + ).solve(system_prompt="", user_message="read", toolkit=toolkit, max_turns=3) + assert result.error is None + assert result.finish_result == {"status": "failure", "summary": "complete"} + assert result.stats["tool_calls"] == 3 + + +def test_claude_interrupt_waits_for_result_message_before_resuming( + make_toolkit, +): + async def scenario(): + toolkit = make_toolkit() + recorder = ClaudeRecorder( + toolkit=toolkit, max_turns=3, dashboard_events=NullDashboardEventSink() + ) + driver = _ClaudeSessionDriver( + sdk=None, options=None, recorder=recorder, emit=recorder.observe + ) + messages = asyncio.Queue() + acknowledgement = asyncio.Event() + + class Client: + async def query(self, text): + pass + + async def interrupt(self): + acknowledgement.set() + + async def receive_messages(self): + while (message := await messages.get()) is not None: + yield message + + class Adapter: + async def on_message(self, driver, message): + raise AssertionError("old SDK completion must not flush new input") + + driver._client = Client() + await driver.query("old turn") + consumer = asyncio.create_task(driver._consume(Adapter())) + interrupted = asyncio.create_task(driver.interrupt()) + await acknowledgement.wait() + assert not interrupted.done() + await messages.put({"type": "ResultMessage", "is_error": True}) + assert await asyncio.wait_for(interrupted, timeout=2) == 1 + assert recorder.error is None + await messages.put(None) + await consumer + + asyncio.run(scenario()) + + def test_api_finish_accepted_during_interrupt_seals_dashboard(): from types import SimpleNamespace diff --git a/tests/unit_tests/rpent/tools/test_common_tools.py b/tests/unit_tests/rpent/tools/test_common_tools.py index dfe31552d..4b9312050 100644 --- a/tests/unit_tests/rpent/tools/test_common_tools.py +++ b/tests/unit_tests/rpent/tools/test_common_tools.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier import numpy as np import pytest @@ -41,13 +43,30 @@ def toolkit(tmp_path, monkeypatch): instance.close() -def test_read_text_file_and_list_dir_use_native_results(toolkit, tmp_path): - (tmp_path / "note.txt").write_text("tool reads", encoding="utf-8") - text_result = toolkit.execute_tool("read_text_file", {"path": "note.txt"}) - directory_result = toolkit.execute_tool("list_dir", {}) +def test_read_text_file_and_list_dir_execute_in_parallel( + toolkit, tmp_path, monkeypatch +): + (tmp_path / "note.txt").write_text("parallel reads", encoding="utf-8") + both_reading = Barrier(2) + authorize_read = toolkit.memory.authorize_read + + def synchronize_reads(path): + resolved = authorize_read(path) + # Both handlers must enter before either can complete. A serial + # executor breaks the barrier and returns tool errors. + both_reading.wait(timeout=3) + return resolved + + monkeypatch.setattr(toolkit.memory, "authorize_read", synchronize_reads) + with ThreadPoolExecutor(max_workers=2) as pool: + text = pool.submit(toolkit.execute_tool, "read_text_file", {"path": "note.txt"}) + directory = pool.submit(toolkit.execute_tool, "list_dir", {}) + text_result = text.result(timeout=5) + directory_result = directory.result(timeout=5) + assert not text_result.is_error assert not directory_result.is_error - assert text_result.data["content"] == "tool reads" + assert text_result.data["content"] == "parallel reads" assert "note.txt" in directory_result.data["files"] diff --git a/tests/unit_tests/rpent/tools/test_recipe.py b/tests/unit_tests/rpent/tools/test_recipe.py new file mode 100644 index 000000000..2b0929696 --- /dev/null +++ b/tests/unit_tests/rpent/tools/test_recipe.py @@ -0,0 +1,72 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import threading +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace + +from rpent.memory import MemoryManager +from rpent.session import EnvState +from rpent.tools import Toolkit, ToolResult, readonly, tool + + +def export(toolkit): + path = toolkit._task_output_dir / toolkit.write_recipe("cell") + return [json.loads(line) for line in path.read_text().splitlines()] + + +@tool +@readonly +def sense(label: str, *, ctx) -> ToolResult: + """Read a sensor.""" + if label == "slow": + ctx.robot.started.set() + assert ctx.robot.release.wait(5) + return ToolResult(data={"label": label}) + + +@tool +def finish(status: str, summary: str, *, ctx) -> ToolResult: + """Finish this session.""" + return ToolResult(data={"status": status, "summary": summary}) + + +def test_generic_parallel_perception_is_recorded_without_new_state_steps(tmp_path): + runtime = SimpleNamespace(started=threading.Event(), release=threading.Event()) + state = EnvState(tmp_path / "observations") + toolkit = Toolkit( + state=state, + memory=MemoryManager(tmp_path / "memory"), + robot=runtime, + output_dir=tmp_path / "recipe", + tools=(sense, finish), + ) + try: + with ThreadPoolExecutor(max_workers=2) as executor: + slow = executor.submit(toolkit.execute_tool, "sense", {"label": "slow"}) + try: + assert runtime.started.wait(5) + fast = executor.submit(toolkit.execute_tool, "sense", {"label": "fast"}) + assert not fast.result(timeout=5).is_error + finally: + runtime.release.set() + assert not slow.result(timeout=5).is_error + assert state.records() == [] + assert export(toolkit) == [ + {"action": "sense", "label": "fast"}, + {"action": "sense", "label": "slow"}, + ] + finally: + toolkit.close() diff --git a/tests/unit_tests/rpent/tools/test_recording.py b/tests/unit_tests/rpent/tools/test_recording.py new file mode 100644 index 000000000..012591625 --- /dev/null +++ b/tests/unit_tests/rpent/tools/test_recording.py @@ -0,0 +1,77 @@ +# Copyright 2026 The RPent Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import Mock + +import numpy as np + +from rpent.memory import MemoryManager +from rpent.session import EnvState +from rpent.tools import ToolContext, Toolkit, ToolResult, tool + + +@tool +def finish(status: str, summary: str, *, ctx: ToolContext) -> ToolResult: + """Accept the requested outcome for this test toolkit.""" + return ToolResult(data={"_finish": True, "status": status, "summary": summary}) + + +def common_toolkit(tmp_path): + return Toolkit( + state=EnvState(tmp_path), + memory=MemoryManager(tmp_path / "memory"), + robot=object(), + output_dir=tmp_path, + tools=(finish,), + ) + + +def values(frames): + return [int(frame[0, 0, 0]) for frame in frames] + + +def rgb(value): + return np.full((8, 8, 3), value, dtype=np.uint8)[:, ::-1] + + +def test_close_uses_toolkit_frames_and_clears_before_save(tmp_path, monkeypatch): + toolkit = common_toolkit(tmp_path) + toolkit.record_frame(rgb(1)) + toolkit.record_frame(rgb(2)) + saved = [] + + def save(name, frames, **kwargs): + assert toolkit._frames == [] + assert all(frame.flags.c_contiguous for frame in frames) + saved.append((name, values(frames), kwargs)) + + monkeypatch.setattr(toolkit.state, "save", save) + toolkit.close() + assert saved == [("episode.mp4", [1, 2], {"step": None, "fps": 20})] + assert toolkit.execute_tool("list_dir", {}).error == "Toolkit is closed." + + +def test_empty_close_and_save_failure_keep_calls_closed(tmp_path, monkeypatch, caplog): + toolkit = common_toolkit(tmp_path) + save = Mock(side_effect=OSError("disk unavailable")) + monkeypatch.setattr(toolkit.state, "save", save) + toolkit.close() + save.assert_not_called() + toolkit = common_toolkit(tmp_path / "second") + monkeypatch.setattr(toolkit.state, "save", save) + toolkit.record_frame(rgb(1)) + toolkit.close() + assert toolkit._frames == [] + assert "failed to save episode video" in caplog.text + assert toolkit.execute_tool("list_dir", {}).error == "Toolkit is closed." diff --git a/tests/unit_tests/rpent/tools/test_scheduling.py b/tests/unit_tests/rpent/tools/test_scheduling.py index 0795c35ea..dbfe5ebcd 100644 --- a/tests/unit_tests/rpent/tools/test_scheduling.py +++ b/tests/unit_tests/rpent/tools/test_scheduling.py @@ -12,50 +12,482 @@ # See the License for the specific language governing permissions and # limitations under the License. -import threading +"""Deterministic CPU scheduling checks using events and queue registration.""" + from concurrent.futures import ThreadPoolExecutor +from threading import Event, get_ident +from types import SimpleNamespace + +import pytest from rpent.memory import MemoryManager from rpent.session import EnvState -from rpent.tools import ToolContext, Toolkit, ToolResult, readonly, tool +from rpent.tools import ( + ToolContext, + Toolkit, + ToolResult, + readonly, + tool, +) + + +@tool +@readonly +def read(key: str, *, ctx: ToolContext) -> ToolResult: + """Read through the test's controlled callback.""" + return ctx.robot.invoke(key, ctx) + + +@tool +def write(key: str, *, ctx: ToolContext) -> ToolResult: + """Act through the test's controlled callback.""" + return ctx.robot.invoke(key, ctx) + + +@tool +def finish(status: str, summary: str, *, ctx: ToolContext) -> ToolResult: + """Finish with this robot's test outcome.""" + ctx.robot.invoke("finish", ctx) + return ToolResult(data={"_finish": True, "status": "failure", "summary": summary}) + + +class ScheduledToolkit(Toolkit): + def __init__(self, root, invoke): + self.captures = 0 + self.capture_hook = None + super().__init__( + state=EnvState(root), + memory=MemoryManager(root / "memory"), + robot=SimpleNamespace(invoke=invoke), + output_dir=root, + tools=( + finish, + read, + write, + ), + ) + + def _capture_observation(self, *, result, **kwargs): + if self.capture_hook: + self.capture_hook() + step = self.captures + self.captures += 1 + data = {"step": step} + return data, [] + + +class Harness: + def __init__(self, root): + self.started = {} + self.hooks = {} + self.order = [] + self.gates = [] + self.pool = ThreadPoolExecutor(max_workers=12) + self.toolkit = ScheduledToolkit(root, self.invoke) + self.queue_changed = Event() + scheduler = self.toolkit._scheduler + original_can_start = scheduler._can_start + + def observe_queue(call): + allowed = original_can_start(call) + if not allowed: + self.queue_changed.set() + return allowed + + scheduler._can_start = observe_queue + + def event(self, key): + return self.started.setdefault(key, Event()) + + def gate(self): + gate = Event() + self.gates.append(gate) + return gate + + def invoke(self, key, ctx): + self.order.append(key) + self.event(key).set() + if key in self.hooks: + self.hooks[key](ctx) + return ToolResult(data={"key": key}) + + def submit(self, name, key): + self.event(key) + return self.pool.submit(self.toolkit.execute_tool, name, {"key": key}) + + def wait_for_queued(self, count=1): + """Wait for submissions to queue while active calls are held by test gates.""" + scheduler = self.toolkit._scheduler + while True: + with scheduler._condition: + if len(scheduler._pending) == count: + return + self.queue_changed.clear() + assert self.queue_changed.wait(2), f"Expected {count} queued calls." + + def wait(self, key): + assert self.event(key).wait(2), f"Call {key} never started." + + def block(self, key, gate): + def wait(ctx): + assert gate.wait(4), f"Gate for {key} not released." + + self.hooks[key] = wait + + +@pytest.fixture +def harness(tmp_path): + value = Harness(tmp_path) + yield value + for gate in value.gates: + gate.set() + if value.toolkit._scheduler._state != "closed": + value.toolkit.close() + value.pool.shutdown(wait=True) + + +def test_shared_calls_overlap_and_waiting_writer_blocks_new_readers(harness): + gate = harness.gate() + for key in ("a", "b"): + harness.block(key, gate) + a = harness.submit("read", "a") + b = harness.submit("read", "b") + harness.wait("a") + harness.wait("b") + writer = harness.submit("write", "w") + harness.wait_for_queued() + c = harness.submit("read", "c") + harness.wait_for_queued(2) + assert not harness.event("w").is_set() + assert not harness.event("c").is_set() + gate.set() + assert all(not f.result(3).is_error for f in (a, b, writer, c)) + assert harness.order[-2:] == ["w", "c"] + + +def test_writers_keep_registration_order_and_can_pass_waiting_reads(harness): + gate = harness.gate() + harness.block("active", gate) + active = harness.submit("write", "active") + harness.wait("active") + futures = [] + for name, key in (("read", "a"), ("write", "b"), ("read", "c"), ("write", "d")): + futures.append(harness.submit(name, key)) + harness.wait_for_queued(len(futures)) + gate.set() + for future in [active, *futures]: + assert not future.result(3).is_error + assert harness.order[:3] == ["active", "b", "d"] + assert set(harness.order[3:]) == {"a", "c"} + + +def test_finish_uses_exclusive_order_and_keeps_admission_open(harness): + gate = harness.gate() + harness.block("active", gate) + active = harness.submit("write", "active") + harness.wait("active") + reader = harness.submit("read", "reader") + harness.wait_for_queued() + finishing = harness.pool.submit( + harness.toolkit.execute_tool, + "finish", + {"status": "success", "summary": "requested"}, + ) + harness.wait_for_queued(2) + writer = harness.submit("write", "writer") + harness.wait_for_queued(3) + gate.set() + for future in (active, reader, writer): + assert not future.result(3).is_error + result = finishing.result(3) + assert not result.is_error + assert harness.order == ["active", "finish", "writer", "reader"] + expected = {"status": "failure", "summary": "requested"} + assert harness.toolkit.finish_result == expected + assert harness.toolkit.captures == 2 + result.data["status"] = "changed" + saved = harness.toolkit.finish_result + saved["summary"] = "changed" + assert harness.toolkit.finish_result == expected + assert not harness.toolkit.execute_tool("write", {"key": "new"}).is_error + + +def test_file_write_is_exclusive_and_does_not_capture(harness, tmp_path): + gate = harness.gate() + harness.block("active", gate) + active = harness.submit("read", "active") + harness.wait("active") + path = tmp_path / "note.txt" + write = harness.pool.submit( + harness.toolkit.execute_tool, + "write_text_file", + {"path": str(path), "content": "saved"}, + ) + harness.wait_for_queued() + assert not path.exists() + gate.set() + assert not active.result(3).is_error + result = write.result(3) + assert not result.is_error + assert result.data == {"path": str(path), "bytes_written": 5} + assert path.read_text() == "saved" + assert harness.toolkit.captures == 0 + + +def test_other_tools_cannot_record_a_finish_result(harness): + harness.toolkit._robot.invoke = lambda key, ctx: ToolResult( + data={"_finish": True, "status": "success", "summary": "ordinary result"} + ) + assert not harness.toolkit.execute_tool("read", {"key": "probe"}).is_error + assert harness.toolkit.finish_result is None -def test_overlapping_calls_are_rejected_and_cancel_waits_for_active_call(tmp_path): - entered = threading.Event() - release = threading.Event() - contexts = [] +@pytest.mark.parametrize("name", ["read", "write"]) +def test_cancel_stops_waiting_calls_and_waits_for_active_cleanup(harness, name): + cancelled = harness.gate() + cleanup_gate = harness.gate() - @tool - @readonly - def read(*, ctx: ToolContext) -> ToolResult: - """Read until the current operation reaches its cancellation boundary.""" - contexts.append(ctx) - entered.set() - assert release.wait(5) + def loop(ctx): + try: + assert ctx._cancel_event.wait(3) + ctx.check_cancelled() + finally: + cancelled.set() + assert cleanup_gate.wait(4) + + harness.hooks["active"] = loop + active = harness.submit(name, "active") + harness.wait("active") + waiting = harness.submit("write", "waiting") + harness.wait_for_queued() + stop = harness.pool.submit(harness.toolkit.cancel_active_and_wait) + assert cancelled.wait(2) + assert waiting.result(2).error == "Tool call cancelled." + assert not harness.event("waiting").is_set() + assert not stop.done() + assert ( + harness.toolkit.execute_tool("read", {"key": "paused"}).error + == "Tool calls are paused." + ) + with pytest.raises(RuntimeError, match="cleanup"): + harness.toolkit.resume_calls() + cleanup_gate.set() + assert active.result(3).error == "Tool call cancelled." + stop.result(3) + assert harness.toolkit.captures == (1 if name == "write" else 0) + harness.toolkit.resume_calls() + assert not harness.toolkit.execute_tool("read", {"key": "new"}).is_error + + +def test_all_active_shared_calls_receive_cancellation(harness): + def loop(ctx): + assert ctx._cancel_event.wait(3) ctx.check_cancelled() - return ToolResult(data={"ok": True}) - - toolkit = Toolkit( - state=EnvState(tmp_path), - memory=MemoryManager(tmp_path / "memory"), - robot=None, - output_dir=tmp_path, - tools=(read,), + + harness.hooks.update(a=loop, b=loop) + futures = [harness.submit("read", key) for key in ("a", "b")] + harness.wait("a") + harness.wait("b") + harness.toolkit.cancel_active_and_wait() + assert [f.result(2).error for f in futures] == ["Tool call cancelled."] * 2 + + +def test_exclusive_admission_is_retained_through_final_capture(harness): + capture_started = harness.gate() + capture_gate = harness.gate() + + def capture(): + capture_started.set() + assert capture_gate.wait(4) + + harness.toolkit.capture_hook = capture + active = harness.submit("write", "write") + assert capture_started.wait(2) + reader = harness.submit("read", "reader") + harness.wait_for_queued() + assert not harness.event("reader").is_set() + capture_gate.set() + assert not active.result(3).is_error + assert not reader.result(3).is_error + + +def test_close_waits_for_final_capture_before_saving_video(harness, monkeypatch): + capture_started = harness.gate() + capture_gate = harness.gate() + cleaned = Event() + + def capture(): + capture_started.set() + assert capture_gate.wait(4) + + harness.toolkit.capture_hook = capture + harness.toolkit.record_frame([[[0, 0, 0]]]) + monkeypatch.setattr( + harness.toolkit.state, "save", lambda *args, **kwargs: cleaned.set() ) - with ThreadPoolExecutor(max_workers=2) as pool: - active = pool.submit(toolkit.execute_tool, "read", {}) + active = harness.submit("write", "write") + assert capture_started.wait(2) + closing = harness.pool.submit(harness.toolkit.close) + scheduler = harness.toolkit._scheduler + with scheduler._condition: + assert scheduler._condition.wait_for( + lambda: scheduler._state == "closed", timeout=2 + ) + assert not cleaned.is_set() + assert not closing.done() + assert ( + harness.toolkit.execute_tool("read", {"key": "late"}).error + == "Toolkit is closed." + ) + capture_gate.set() + active.result(3) + closing.result(3) + assert cleaned.is_set() + harness.toolkit.cancel_active_and_wait() + harness.toolkit.resume_calls() + assert ( + harness.toolkit.execute_tool("read", {"key": "closed"}).error + == "Toolkit is closed." + ) + + +def test_toolkits_do_not_share_execution_locks(harness, tmp_path): + gate = harness.gate() + harness.block("blocked", gate) + active = harness.submit("write", "blocked") + harness.wait("blocked") + other = ScheduledToolkit(tmp_path / "other", lambda key, ctx: ToolResult()) + try: + assert not other.execute_tool("write", {"key": "independent"}).is_error + finally: + gate.set() + other.close() + assert not active.result(3).is_error + + +def test_cancel_between_admission_and_handler_skips_execution_and_capture( + harness, monkeypatch +): + admitted = harness.gate() + start_gate = harness.gate() + scheduler = harness.toolkit._scheduler + original = scheduler.acquire + + def acquire(tool): + call = original(tool) + admitted.set() + assert start_gate.wait(4) + return call + + monkeypatch.setattr(scheduler, "acquire", acquire) + future = harness.submit("write", "never_start") + assert admitted.wait(2) + stopping = harness.pool.submit(harness.toolkit.cancel_active_and_wait) + with scheduler._condition: + assert scheduler._condition.wait_for( + lambda: scheduler._state == "paused", timeout=2 + ) + start_gate.set() + assert future.result(3).error == "Tool call cancelled." + stopping.result(3) + assert harness.toolkit.captures == 0 + assert not harness.event("never_start").is_set() + + +def test_overlapping_cancellation_requests_wait_for_active_cleanup( + harness, monkeypatch +): + cancelled = harness.gate() + cleanup_gate = harness.gate() + + def loop(ctx): try: - assert entered.wait(5) - rejected = toolkit.execute_tool("read", {}) - assert rejected.error == "another tool operation is still active" - assert len(contexts) == 1 - cancelled = pool.submit(toolkit.cancel_active_and_wait) - assert contexts[0]._cancel_event.wait(5) - assert not cancelled.done() + assert ctx._cancel_event.wait(3) + ctx.check_cancelled() finally: - release.set() - assert active.result(5).error == "Tool call cancelled." - cancelled.result(5) - assert not toolkit.execute_tool("read", {}).is_error - assert len(contexts) == 2 - assert not contexts[1]._cancel_event.is_set() + cancelled.set() + assert cleanup_gate.wait(4) + + harness.hooks["active"] = loop + active = harness.submit("write", "active") + harness.wait("active") + scheduler = harness.toolkit._scheduler + original_wait = scheduler._condition.wait + waiting = set() + both_waiting = Event() + + def wait_done(timeout=None): + waiting.add(get_ident()) + if len(waiting) == 2: + both_waiting.set() + return original_wait(timeout) + + monkeypatch.setattr(scheduler._condition, "wait", wait_done) + first = harness.pool.submit(harness.toolkit.cancel_active_and_wait) + assert cancelled.wait(2) + second = harness.pool.submit(harness.toolkit.cancel_active_and_wait) + assert both_waiting.wait(2) + with pytest.raises(RuntimeError, match="cleanup"): + harness.toolkit.resume_calls() + assert not first.done() and not second.done() + cleanup_gate.set() + assert active.result(3).error == "Tool call cancelled." + first.result(3) + second.result(3) + harness.toolkit.resume_calls() + assert not harness.toolkit.execute_tool("read", {"key": "resumed"}).is_error + + +def test_resume_does_not_revive_cancelled_waiters_or_extend_old_cancellation( + harness, monkeypatch +): + scheduler = harness.toolkit._scheduler + active_gate = harness.gate() + rejection_gate = harness.gate() + new_gate = harness.gate() + old_woken = Event() + old_thread = None + original_wait = scheduler._condition.wait + + def delay_old_waiter(timeout=None): + nonlocal old_thread + if old_thread is None: + old_thread = get_ident() + result = original_wait(timeout) + if get_ident() != old_thread: + return result + scheduler._condition.release() + try: + old_woken.set() + assert rejection_gate.wait(4) + finally: + scheduler._condition.acquire() + return result + + monkeypatch.setattr(scheduler._condition, "wait", delay_old_waiter) + harness.block("active", active_gate) + active = harness.submit("write", "active") + harness.wait("active") + old = harness.submit("write", "old") + harness.wait_for_queued() + stopping = harness.pool.submit(harness.toolkit.cancel_active_and_wait) + assert old_woken.wait(2) + active_gate.set() + active.result(3) + assert not stopping.done() + harness.toolkit.resume_calls() + + def new_call(ctx): + assert new_gate.wait(4) + ctx.check_cancelled() + + harness.hooks["new"] = new_call + new = harness.submit("read", "new") + harness.wait("new") + rejection_gate.set() + assert old.result(3).error == "Tool call cancelled." + assert not harness.event("old").is_set() + stopping.result(3) + assert not new.done() + new_gate.set() + assert not new.result(3).is_error