diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ef34a93..9cb15fa 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -301,12 +301,48 @@ jobs: uv run --no-sync python -m pytest plugins/inspect-robots-agent/tests -q --cov=inspect_robots_agent --cov-report=term-missing --cov-fail-under=0 + # The CaP-X code-as-policy suite needs no GPU, sockets, or CaP-X checkout: + # httpx.MockTransport scripts the LLM plus all three model-server protocols, + # and an in-test joint-space embodiment exercises the complete rollout. + plugin-capx: + name: plugin · capx + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - name: Install uv + uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + cache-dependency-glob: | + **/pyproject.toml + **/uv.lock + - name: Sync workspace (core + plugins, dev extras) + run: uv sync --locked --python 3.11 --all-packages --extra dev + - name: Ruff (plugin, lint) + run: uv run --no-sync ruff check plugins/inspect-robots-capx + - name: Ruff (plugin, format check) + run: uv run --no-sync ruff format --check plugins/inspect-robots-capx + - name: Mypy (plugin, strict) + run: > + uv run --no-sync mypy + --config-file plugins/inspect-robots-capx/pyproject.toml + plugins/inspect-robots-capx/src/inspect_robots_capx + plugins/inspect-robots-capx/tests + - name: Pytest (plugin, with coverage) + # Report-only, like the other plugin jobs. The root 100% gate remains + # scoped to the numpy-only inspect_robots core package. + run: > + uv run --no-sync python -m pytest plugins/inspect-robots-capx/tests -q + --cov=inspect_robots_capx --cov-report=term-missing --cov-fail-under=0 + # Single required status check for branch protection: green iff every job # above is green. Add new jobs to 'needs' here or they won't gate merges. ci-ok: name: ci-ok if: ${{ always() }} - needs: [quality, test, test-rerun, test-extra, core-only-import, plugin-isaacsim, plugin-xpolicylab, plugin-ros, plugin-agent] + needs: [quality, test, test-rerun, test-extra, core-only-import, plugin-isaacsim, plugin-xpolicylab, plugin-ros, plugin-agent, plugin-capx] runs-on: ubuntu-latest steps: - name: Fail unless every needed job succeeded @@ -319,7 +355,7 @@ jobs: alert-red-main: name: alert on red main if: ${{ failure() && github.event_name == 'push' }} - needs: [quality, test, test-rerun, test-extra, core-only-import, plugin-isaacsim, plugin-xpolicylab, plugin-ros, plugin-agent] + needs: [quality, test, test-rerun, test-extra, core-only-import, plugin-isaacsim, plugin-xpolicylab, plugin-ros, plugin-agent, plugin-capx] runs-on: ubuntu-latest permissions: issues: write diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index f23afb5..df21fbd 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -175,3 +175,26 @@ jobs: # The plugin is versioned independently; a core release that doesn't # bump it must not fail on the already-uploaded file. skip-existing: true + + publish-capx: + name: publish inspect-robots-capx + needs: [cut] + if: ${{ !cancelled() && (github.event_name == 'release' || needs.cut.result == 'success') }} + runs-on: ubuntu-latest + environment: pypi-capx + permissions: + id-token: write + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ needs.cut.outputs.tag || github.ref }} + - name: Install uv + uses: astral-sh/setup-uv@v5 + - name: Build sdist + wheel + run: uv build --package inspect-robots-capx + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + # The plugin is versioned independently; a core release that doesn't + # bump it must not fail on the already-uploaded file. + skip-existing: true diff --git a/CLAUDE.md b/CLAUDE.md index 3ace835..b8c231b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -35,7 +35,9 @@ rollout, scores it, and writes an immutable `EvalLog`. Mirrors Inspect AI's (policy adapter speaking the XPolicyLab websocket protocol — 40+ served VLAs, no xpolicylab dep), and `plugins/inspect-robots-agent/` (LLMs as policies via the OpenAI-compatible - wire format — httpx only, no provider SDKs; registered as `agent`). + wire format — httpx only, no provider SDKs; registered as `agent`), and + `plugins/inspect-robots-capx/` (CaP-X code-as-policy over perception and IK + HTTP servers; registered as `capx`). ## Working here diff --git a/README.md b/README.md index b73225a..21cae33 100644 --- a/README.md +++ b/README.md @@ -154,6 +154,27 @@ Read the recorded agent conversation with `inspect-robots view LOG.json`, including the camera frames the model saw (for `--store-frames` runs). +### Generate robot policy code with CaP-X: + +The [inspect-robots-capx](plugins/inspect-robots-capx/) plugin evaluates a +code-as-policy agent in the same policy slot. The LLM writes Python against +SAM3 segmentation, Contact-GraspNet planning, Pyroki IK, and speed-limited +joint-motion helpers. CaP-X model servers run separately, while the persistent +code namespace and action queue run inside the evaluator. + +```bash +uv pip install inspect-robots-capx + +inspect-robots "place the fork on the plate" --policy capx \ + --embodiment \ + -P model=anthropic/claude-fable-5 -P sam3_url=http://gpu-box:8114 +``` + +The v1 adapter requires one `joint_pos` arm with a declared gripper, full +joint-state proprioception, a control rate, and a camera. See the +[plugin README](plugins/inspect-robots-capx/) for model-server bringup, depth +metadata, gripper polarity, and the model-code trust boundary. + ### Run in simulation The same instruction runs on your configured simulator instead of the @@ -264,6 +285,10 @@ adapter shipped from this repo as separate packages: embodiment through tool calls, as a first-class policy. The same `--policy agent` runs ad-hoc instructions and scores on registered tasks next to fine-tuned VLAs. +- **[inspect-robots-capx](plugins/inspect-robots-capx/)**: evaluate CaP-X-style + code-as-policy agents against a joint-space embodiment. Model-generated + Python calls separately served SAM3, Contact-GraspNet, and Pyroki helpers, + then queues approver-checked joint targets behind `--policy capx`. ```bash # Isaac Lab world + a π0 checkpoint served by XPolicyLab, evaluated end to end: diff --git a/plans/0021-capx-policy-plugin.md b/plans/0021-capx-policy-plugin.md new file mode 100644 index 0000000..36eaab2 --- /dev/null +++ b/plans/0021-capx-policy-plugin.md @@ -0,0 +1,450 @@ +# 0021 — `inspect-robots-capx`: code-as-policy agents via CaP-X perception servers + +## 1. Goal + +[CaP-X](https://github.com/capgym/cap-x) (arXiv:2603.22435, ICML 2026) is a +framework for **code-as-policy** manipulation: an LLM writes Python against +perception and motion primitives (SAM3 segmentation, Contact-GraspNet grasp +planning, Pyroki IK, joint-space moves) in a persistent REPL-like namespace, +with a multi-turn execute/observe/regenerate protocol. Frontier models score +30%+ zero-shot on their bench; a 7B coder hits 72% after RL. It is the third +policy class after served VLAs (`inspect-robots-xpolicylab`) and raw LLM +tool-calling (`inspect-robots-agent`), and nobody offers a side-by-side eval +of all three on one task. + +This plan ships `plugins/inspect-robots-capx`: a `Policy` registered as +`capx` in which the LLM writes Python per turn. The code runs with perception +helpers backed by CaP-X's model servers and motion primitives that queue +speed-limited joint targets; the queue becomes the returned `ActionChunk`. + +```bash +# terminal 1 — CaP-X checkout serves the models (GPU box is fine) +uv run capx/serving/launch_sam3_server.py --port 8114 +uv run capx/serving/launch_contact_graspnet_server.py --port 8115 +uv run python -c 'from capx.serving.launch_pyroki_server import main; main(robot="panda_description", port=8116)' +# (launch_pyroki_server.py calls bare main() with no CLI parsing, unlike the +# other two servers, so flags on it are silently ignored) + +# terminal 2 — Inspect Robots drives any joint-space embodiment with it +inspect-robots "pick up the red cube" --policy capx --embodiment \ + -P model=anthropic/claude-fable-5 -P sam3_url=http://gpu-box:8114 +``` + +(The core `cubepick` mock is a 2-D `eef_delta_pos` world and does not fit the +v1 joint-space profile below; the e2e test ships its own joint-space mock +embodiment, §9. Real targets are the YAM / SO-101 / Isaac adapters.) + +One adapter ⇒ CaP-X-style agents become evaluable on any Inspect Robots +embodiment, including ones CaP-X does not support (our YAM and SO-101 +adapters), next to VLAs and the tool-calling agent. + +Non-goals (YAGNI, all deliberate): + +- No CaP-RL / training integration; no CaP-Bench task port (a later plan — + tasks are an independent axis). +- No visual-differencing VLM sidecar, skill library, or ensembling from + CaP-Agent0; the core codegen loop is the paper's load-bearing part. +- No Molmo / OWL-ViT / SAM2 / cuRobo clients — SAM3 + GraspNet + Pyroki are + the reduced-API pipeline the paper's headline results use. +- No dual-arm and no end-effector control mode in v1: the policy requires a + single-arm `joint_pos` action space whose semantics declare a gripper + (`ActionSemantics.gripper != "none"`) and raises a clear `bind()`-time + error otherwise. +- No server *launching* (same doctrine as 0007: lifecycle belongs upstream; + we connect to URLs). + +## 2. Grounding: what CaP-X actually is (verified against source) + +Facts established by reading `capgym/cap-x@main` (2026-07, 7 commits): + +- **Not on PyPI**; research-grade monorepo with heavy deps (robosuite, torch, + omnigibson). The agent loop `exec()`s model code in-process with helper + functions bound into the namespace (`capx/envs/tasks/base.py`); variables + persist across turns. +- **Prompting**: one system line ("generates Python code to directly solve + the task"), a task prompt, then API docs auto-scraped from helper + docstrings via `inspect` (`combined_doc()` in + `capx/integrations/base_api.py`). Replies must be raw Python, no fences. +- **Multi-turn protocol** (from the shipped env_configs): after execution the + model sees the executed code + stdout + stderr and must answer + `REGENERATE` + new code, or `FINISH`. +- **Model servers** are plain FastAPI JSON-over-HTTP (`capx/serving/`): + - SAM3 `/segment`: request `{image_base64: , text_prompt}`; + response `{results: [{mask_base64: , shape: [H, W], + box: [x1, y1, x2, y2], score, label}]}` (masks decode via + `np.frombuffer(...).reshape(shape)`, uint8). + - Contact-GraspNet `/plan`: request `{depth_base64, cam_K_base64, + segmap_base64, segmap_id, local_regions, filter_grasps, + skip_border_objects, z_range, forward_passes, max_retries}` where + `*_base64` are base64-encoded `.npy` payloads (`np.save` round-trip); + response `{grasps_base64, scores_base64, contact_pts_base64}` in the + same `.npy` encoding. Grasp poses are `(K, 4, 4)` in the camera frame. + - Pyroki `/ik`: request `{target_pose_wxyz_xyz: [7 floats], + prev_cfg: [floats] | null}`; response `{joint_positions: [floats]}`. + The server owns the robot model. Both `joint_positions` and `prev_cfg` + are the URDF's **full actuated-joint config**, gripper joints included + (CaP-X strips the gripper entry client-side: + `extract_arm_joints(cfg) = cfg[:-1]` in + `capx/integrations/franka/common.py`, and warm-starts with the full + stored config). +- CaP-X's own clients retry POSTs with backoff (`post_with_retries`) because + model servers cold-start slowly. + +## 3. The two big decisions + +**Speak the wire, don't import the package** (same as 0007): a published +wheel cannot depend on git-only research code, and the protocol surface is +three JSON endpoints plus two base64 codecs (~100 lines with tests). Drift +risk is accepted and mitigated exactly like 0007 (§9). + +**Chunk-per-turn instead of CaP-X's blocking primitives.** In CaP-X, motion +primitives step the simulator mid-code. Inspect Robots inverts control: +`Policy.act(observation)` returns an open-loop `ActionChunk`. Rather than +threads/coroutines to suspend user code at each primitive, motion primitives +**queue** interpolated joint targets and the whole turn's queue is returned +as one chunk. Perception inside a turn sees the turn's initial observation; +post-motion effects are observed next turn. This matches how CaP-X's +reduced-API agent actually behaves (plan grasps from the turn's observation, +then execute; its multi-turn prompt tells the model to verify state *next +turn*) and matches how ActionChunk semantics already work for VLAs. No +concurrency, no partial-code re-entry, testable in-process. + +## 4. Deliverable layout + +```text +plugins/inspect-robots-capx/ +├── pyproject.toml # hatchling; entry point inspect_robots.policies:capx +├── README.md # install, server bringup, arg table, trust model +├── src/inspect_robots_capx/ +│ ├── __init__.py # CapxPolicy, capx_policy factory, __version__ +│ ├── _codec.py # b64 PNG + b64 .npy encode/decode (CaP-X wire codecs) +│ ├── _servers.py # Sam3Client / GraspNetClient / PyrokiClient (httpx) +│ ├── _sandbox.py # per-trial exec namespace, stdout/stderr capture, helper binding +│ ├── _motion.py # joint-target queue: speed-limited interpolation, gripper, hold +│ └── policy.py # CapxPolicy: codegen loop, CaP-X multi-turn protocol +└── tests/ + ├── conftest.py # httpx.MockTransport stubs: LLM + 3 CaP-X servers + ├── test_codec.py # golden wire tests against recorded CaP-X payload shapes + ├── test_servers.py + ├── test_sandbox.py + ├── test_motion.py + └── test_policy.py # incl. e2e vs an in-test joint-space mock embodiment (§9) +plans/0021-capx-policy-plugin.md # this file +.github/workflows/ci.yml # + plugin-capx job, wired into both needs lists +.github/workflows/release.yml # + publish-capx job (environment: pypi-capx) +plugins/inspect-robots-agent/… # export ChatClient/resolve_provider/png_data_url (see §5) +README.md # + code-as-policy section +CLAUDE.md # plugins list mentions the new package +``` + +Dependencies: `inspect-robots>=0.4`, `inspect-robots-agent>=0.12`, `numpy`, +`httpx`. No `capx`, no `requests`, no provider SDKs, no Pillow. No PNG +*reader* exists anywhere in the pipeline: SAM3 requests need PNG *encoding* +(the agent plugin's `_png` writer, re-exported per §5), SAM3 masks return as +raw uint8 bytes, GraspNet speaks `.npy`, and Pyroki speaks JSON floats. + +## 5. Reuse from `inspect-robots-agent`, made explicit + +The capx plugin needs an OpenAI-compatible chat client with provider routing +and PNG data URLs — exactly what `inspect_robots_agent._llm` and `._png` +already implement. Duplicating ~300 maintained lines is worse than a +first-party dependency, so: + +- `inspect-robots-agent` bumps to 0.12.0 (it is at 0.11.0 today) and + re-exports `ChatClient`, `ResponsesClient` (the agent's `wire="responses"` + client; a chat-only capx would inherit `effort="low"` defaults that 400 on + direct OpenAI endpoints with an error message telling users to switch + wires), `AssistantMessage` (the completion type capx must annotate under + mypy strict), `resolve_provider`, `Provider`, `ENV_MODEL` (capx mirrors + the agent's model-default env var), `png_data_url`, and `encode_png` (the + raw PNG writer; SAM3 requests need bare b64 PNG, not a data URL) from its + package `__init__`, extending its `__all__`; the modules stay where they + are. +- `inspect-robots-capx` imports only those names from + `inspect_robots_agent` (never from underscore modules). +- Both packages live in this repo's uv workspace and release together, so + the coupling is CI-checked on every PR. + +## 6. The policy: `CapxPolicy` + +`CapxPolicy(PolicyBase)`, entry point `capx`. Constructor args (all keyword, +recorded via a `CapxPolicyConfig(PolicyConfig)` dataclass like the agent +plugin's): + +| Arg | Default | Meaning | +| --- | --- | --- | +| `model` / `base_url` / `api_key_env` / `temperature` / `effort` / `wire` | agent-plugin defaults | LLM provider routing incl. the chat/responses wire selector, shared via §5 | +| `sam3_url` | `"http://127.0.0.1:8114"` | CaP-X SAM3 server | +| `graspnet_url` | `"http://127.0.0.1:8115"` | CaP-X Contact-GraspNet server | +| `pyroki_url` | `"http://127.0.0.1:8116"` | CaP-X Pyroki IK server | +| `camera` | `None` (sole camera, error if several) | observation camera feeding perception | +| `depth_key` | `"depth"` | `observation.extra` key holding `(H, W)` float depth (see below) | +| `intrinsics_key` / `extrinsics_key` | `"intrinsics"` / `"extrinsics"` | `extra` keys for `(3, 3)` K and `(4, 4)` camera-to-robot-base (the plugin treats world and base as the same frame, as CaP-X does: Pyroki solves in the URDF base frame with no world transform; the prompt and README state this) | +| `max_llm_calls` | `100` | per-trial LLM budget; exhaustion forces give-up | +| `max_code_failures` | `3` | consecutive exec-*error* turns before `RuntimeError`. Per-trial counter, persists across `act()` calls; any turn whose exec raised increments it (even one that queued actions first, which still returns its chunk); any turn that raised nothing resets it. Clean perception-only turns are bounded by the LLM budget | +| `max_speed_frac` | `0.1` | joint-interpolation speed cap, same semantics as the agent plugin | +| `request_timeout_s` | `120` | total wall-clock budget per helper call, retries included; each POST attempt gets `min(remaining, 30)` (CaP-X models are slow) | +| `transcript_echo` | `False` | stderr live echo, same as agent | +| `transport` / `env` | `None` | test injection, same as agent | + +**Depth convention.** Core `Observation.images` are uint8 RGB; depth has no +core slot. The plugin defines (and its README documents) an `extra`-key +convention: embodiments that want grasp planning expose float depth, +intrinsics, and camera-to-world extrinsics under the keys above, each as +**either the array or a zero-arg callable returning it** (helpers call it +if callable). The callable form is the documented recommendation for real +embodiments: the rollout deep-retains each step's `observation.extra` in +`TrialRecord` (only `images` get stripped to the frame store), so a raw +per-step 480x640 float64 depth map would turn trial records into an +in-memory depth video buffer at real control rates; a callable retains +nothing. Helpers that need a missing key raise inside the sandbox with a +message naming the key and the convention, which the model sees as stderr +and can route around (segmentation and IK still work without depth). +Nothing in core changes; a core follow-up (strip declared bulk `extra` keys +in `_store_frames`) is noted in §11 but not required. + +**`bind(embodiment_info)`** adopts the embodiment spaces like the agent +plugin. The v1 profile, enforced with an actionable error naming this plan: + +- `Box` action space, `control_mode="joint_pos"`, at least 2 dims, and + `semantics.gripper != "none"`. The gripper dim is located via the box's + `dim_labels` (a label named `"gripper"`); when labels are absent the last + dim is the documented fallback. Arm dof = the remaining dims; gripper + open/close values come from the box bounds on that dim with a pinned + polarity — **high bound = open, low bound = closed**, mirroring core's + normalized-gripper *state* convention — overridable via a + `gripper_open_is_high: bool = True` constructor arg for + inverted-polarity embodiments. +- a proprioceptive reference resolved from `observation_space.state` the + same way the agent toolset does (`build_toolset`'s state-labels logic): a + single state field whose shape equals the full action dim. The motion + cursor and the FINISH/GIVE_UP hold action read that one field from the + turn's observation; there is no split joint_pos/gripper key convention. +- per-step interpolation deltas derived from `control_hz` and + `max_speed_frac`, both required. `_motion.py` reproduces the agent + toolset's step-limit derivation exactly (the `min(max_speed_frac / hz, + 0.05)`-of-range backstop and the native-dtype range arithmetic of + `DeltaLimitApprover`), so *within a chunk* per-step deltas never exceed + what the CLI's default approver allows — otherwise a low `control_hz` + would get silently `delta_clamped` and the executed trajectory would + diverge from what the sandbox reported. At *turn boundaries* the cursor + re-seeds from observed state while `DeltaLimitApprover` clamps against + the last approved action, so a lagging arm can still get its first + next-turn action clamped; the agent plugin shares this edge, the README + notes it, and the claim is scoped to within-chunk deltas. + +**`reset(scene)`** starts the per-trial conversation: + +- system: code-as-policy instructions modeled on CaP-X's prompt (write raw + Python, no fences; helpers are pre-bound; import numpy explicitly; + variables persist across turns), the helper API docs (§7, a static string + — no runtime docstring scraping), the action-space/embodiment summary and + embodiment `docs` notes, and the call budget. +- user: `Goal: {scene.instruction}`. +- clears the sandbox namespace, the motion queue, and the Pyroki IK + warm-start config (`prev_cfg` is `None` on the first call of each trial, + preserving trial independence); `atexit`-safe `close()` closes the httpx + clients (mirrors 0007's lifecycle care). + +**`act(observation)`** — the codegen loop: + +1. Append the observation message: labeled state text plus camera PNGs, + formatted the same way as the agent plugin but **reimplemented** (~40 + lines; the agent's `_state_lines`/`_observation_content` stay private — + only the §5 re-export list is imported, `png_data_url` included). The + previous turn's execution report is already in `_messages` (step 4), so + the observation message carries only the fresh observation. +2. Ask the LLM. Accept raw Python or a single fenced block (strip fences — + CaP-X asks for raw code, but its own multi-turn prompt then demands + fenced code after `REGENERATE`; be liberal). Recognize the control words + `FINISH` and `GIVE_UP` (bare, before any code): both return a one-action + hold chunk (repeat current joint state) whose action meta carries + `request_stop: True` and a `stop_reason` (the rollout already honors + this, and an `ActionChunk` cannot be empty). +3. Execute the code in the sandbox: helpers bound, stdout/stderr captured + (traceback printed to captured stderr on any exception, CaP-X-style). +4. Append the CaP-X-style execution report (executed code, stdout, stderr, + truncated tail-first to a documented cap so context stays bounded) to + `_messages` **eagerly, before returning** — chunk-level `meta` is a dead + channel (`DefaultController` buffers only `chunk.actions` and the JSON + `EvalLog` persists no chunk meta), so the transcript is the record; a + trial that terminates on this chunk still carries the final turn's + stdout/stderr in `transcript()`. The executed code (untruncated) also + rides the first queued action's `meta["code"]`, which survives into + `StepRecord` for step-level debugging. If the queue has actions: pop it + and return `ActionChunk(actions=queue, control_hz=embodiment control_hz, + inference_latency_s=)`. +5. If the queue is empty (pure-perception turn or exec error): the report + from step 4 is already the feedback message; loop to 2 — bounded by + `max_code_failures` consecutive *error* turns (a clean perception-only + turn resets the failure counter but still loops; the LLM-call budget is + the global bound, exhaustion forces the give-up hold chunk exactly like + the agent plugin's `_forced_give_up`). + +`transcript()` / `transcript_delta()`: sanitized copies with images elided, +identical mechanics to the agent plugin, so `inspect-robots inspect +--transcript` renders capx trials for free. + +**`PolicyInfo`**: name `capx`, adopted spaces and `control_hz` from bind, +same placeholder-before-bind pattern as the agent plugin. + +## 7. The sandbox helpers (what the model's code can call) + +Bound into the exec namespace by `_sandbox.py`; documented verbatim in the +system prompt. All arrays NumPy. `obs` is the current turn's observation +(images dict, state dict, plus depth/intrinsics/extrinsics when the +embodiment provides them). + +- `segment(text: str) -> list[dict]` — SAM3 text-prompt segmentation on the + configured camera's current RGB. Returns `{mask (H, W) bool, box, score, + label}` dicts (CaP-X's client shape). +- `plan_grasp(mask: np.ndarray) -> tuple[np.ndarray, np.ndarray]` — + Contact-GraspNet on current depth + intrinsics with `mask` as the segmap. + Returns `(K, 4, 4)` poses in the **camera frame** and `(K,)` scores; the + prompt documents `extrinsics @ pose` for world frame, mirroring CaP-X's + docstring example. When the server finds no grasps its retry loop returns + flat `np.array([])` payloads; the client normalizes those to + `(0, 4, 4)` / `(0,)` so `K == 0` is an ordinary, checkable result rather + than a shape crash (golden wire test in §9). +- `solve_ik(position: np.ndarray, quaternion_wxyz: np.ndarray) -> + np.ndarray` — Pyroki `/ik`. The wire speaks the URDF's full actuated + config (§2): the client keeps the last full server-returned config as the + warm-start `prev_cfg` (mirroring CaP-X's stored `self.cfg`; `None` on the + first call) and returns the response stripped to the bound arm dof + (leading `arm_dof` entries), raising an actionable in-sandbox error when + `len(joint_positions) < arm_dof` (server robot ≠ embodiment robot, the + §11 risk). The helper doc tells the model it gets arm joints, sized for + `move_to_joints`. +- `move_to_joints(joints: np.ndarray) -> None` — queue speed-limited linear + interpolation from the cursor to `joints` (gripper dim held); advances the + cursor. The cursor starts each turn at the bound proprioceptive state + field (§6) from the turn's observation. +- `open_gripper() / close_gripper() -> None` — queue a gripper ramp at the + cursor's arm pose, interpolated under the same per-step delta limit as arm + moves (so the CLI's default `DeltaLimitApprover` never silently clamps a + gripper jump the sandbox reported as done); advance the cursor's gripper + value. +- `print(...)` — captured stdout is the model's feedback channel (CaP-X + convention). + +Everything else is plain Python: the model may import installed packages +(numpy, scipy if present). **Trust model, stated loudly in README and module +docstring**: the code executes in-process with the evaluator's privileges, +exactly as in CaP-X. This is the policy class under evaluation, not a +sandboxing product; run untrusted models in a container. The safety story +for the *robot* is unchanged: every queued action still passes the rollout's +approver chain. The Clamp + DeltaLimit guardrails are the *CLI* default +(`eval()` called programmatically defaults to `AutoApprover`); the README +says exactly that and shows the programmatic approver wiring. + +## 8. Wire clients: `_codec.py` + `_servers.py` + +- `_codec.py`: `png_b64_encode(rgb)` (reuse agent `_png` writer + b64), + `npy_b64_encode/decode` (`np.save`/`np.load` round-trip, `allow_pickle=False` + on load), `mask_decode(mask_b64, shape)` (`np.frombuffer` reshape). Upstream + commit hash recorded in the module docstring (0007's drift mitigation). +- `_servers.py`: three tiny clients over one shared `httpx.Client` + (injectable transport). POST with retries and backoff (bounded by + `request_timeout_s` wall clock, CaP-X-style cold-start tolerance), JSON + bodies exactly matching §2's schemas. Connection failures raise actionable + errors naming the URL and the CaP-X launch command. Lazy: no request until + the model's code first calls a helper, so `list policies` and compat + checks never touch the network (0007's lazy-connection doctrine). +- GraspNet client hardcodes CaP-X's client-side defaults (`segmap_id=1`, + `local_regions/filter_grasps=True`, `z_range=[0.2, 2.0]`, + `forward_passes=2`, `max_retries=10`) — they are model-tuning knobs, not + user API. + +## 9. Tests (no GPU, no sockets, no capx checkout) + +All network stubbed via `httpx.MockTransport` handlers implementing the §2 +schemas (LLM stub reuses the agent plugin's canned-completion approach): + +- codec: golden round-trips — a float32 depth map through `npy_b64`, a bool + mask through `mask_decode`, structural equality against handcrafted wire + payloads shaped like CaP-X's (guards drift without importing capx), + including the empty-grasp flat-`[]` payload normalizing to + `(0, 4, 4)` / `(0,)`. +- servers: request bodies match the schemas exactly (recorded by the stub); + retries on 503-then-200; actionable ConnectionError text; timeout path; + `/ik` golden test with an `arm_dof + 1`-long full-config response (strip + to arm dof, full config becomes the next `prev_cfg`, short response + raises the actionable error). +- sandbox: namespace persists across turns within a trial and resets across + trials; stdout/stderr capture; traceback lands in stderr; helper raising + (e.g. missing depth) surfaces as stderr not a crash; depth accepted as + both a raw array and a zero-arg callable. +- motion: interpolation respects `max_speed_frac` and `control_hz`, + including a low-`control_hz` case proving the per-step delta stays under + the `DeltaLimitApprover` backstop; cursor chaining across move/gripper + calls; hold chunk shape; gripper values from box bounds. +- policy: bind profile enforcement (rejects ee mode, dual-arm-sized boxes + without gripper, missing control_hz); fence stripping and raw-code paths; + FINISH/GIVE_UP → hold chunk with `request_stop` meta; perception-only + turn loops then returns queued chunk; consecutive-failure and call-budget + bounds; execution report truncation; transcript sanitization; the + execution report (incl. a terminal turn's stdout/stderr) lands in + `transcript()` eagerly, and the first queued action's `meta["code"]` + carries the executed code. +- e2e: `CapxPolicy` vs an in-test joint-space mock embodiment (arm dims + + labeled gripper dim, `joint_pos` semantics, a matching single-field + StateSpec, a small RGB camera — the pattern of the agent plugin's + `_AbsoluteEmbodiment` test double; core `CubePick` is `eef_delta_pos` and + deliberately out of profile) with a scripted LLM transport whose canned + code calls `segment` → `solve_ik` → `move_to_joints` → `close_gripper` and + FINISHes on turn 2; assert the rollout completes and the log carries the + transcript. +- registry: entry point resolves; factory forwards `-P`-style kwargs. + +Coverage: plugin CI runs pytest without the core gate (per repo convention); +aim for full-line coverage of `policy.py`, `_motion.py`, `_sandbox.py`. + +## 10. CI, workspace, release, docs + +- `pyproject.toml` mirrors the agent plugin's (hatchling, + hatch-fancy-pypi-readme, static `version = "0.1.0"`, mypy strict, ruff + line-length 100, `[tool.uv.sources]` workspace pins for both first-party + deps). `uv lock` and commit. +- CI: `plugin-capx` job cloned from `plugin-agent` (ruff, ruff format, mypy + strict, pytest), ubuntu-only; **added to both `needs` lists** (`ci-ok` and + the sibling at ci.yml:322) — CLAUDE.md's rule, and 0007 hit the same edge. +- Release: `publish-capx` job in release.yml (environment `pypi-capx`); + maintainer creates the PyPI trusted-publisher environment before first + release (PR description calls this out; it is a settings action, not code). + The agent-plugin version bump to 0.12.0 rides the same PR. +- Entry point: `[project.entry-points."inspect_robots.policies"] + capx = "inspect_robots_capx.policy:capx_policy"`. +- Docs: plugin README (server bringup from a CaP-X checkout, arg table, + depth-key convention, trust-model warning, troubleshooting cold starts); + root README gets a code-as-policy paragraph next to the agent plugin's; + root CLAUDE.md plugins list gains the package. Repo writing style applies + (no em dashes in prose, no AI-tell patterns). +- Core is untouched: no new core deps, no `__all__`/api-snapshot churn. + +## 11. Risks & mitigations + +| Risk | Mitigation | +| --- | --- | +| Server schema drift (no version field) | Golden wire tests; upstream commit hash in `_codec.py`; schemas isolated in two small modules | +| `exec()` of model output alarms users | Loud trust-model section in README + module docstring; approver chain still gates every action; containerize for untrusted models | +| Embodiments lack depth (SO-101 webcam rigs) | Documented degradation: segmentation + IK still work; `plan_grasp` raises a routable in-sandbox error naming the extra-key convention | +| Raw per-step depth arrays bloat `TrialRecord` memory | Callable form of the extra-key convention is the documented recommendation (§6); possible core follow-up: strip declared bulk `extra` keys in `_store_frames` | +| 5-DOF arms can't reach 6-DOF grasp poses | Out of scope for the plugin (Pyroki returns best-effort IK); README notes the caveat and suggests top-down grasp filtering in the task prompt | +| Model floods context with huge stdout | Execution report truncated to a documented cap, tail-first (errors are at the tail) | +| Coupling to `inspect-robots-agent` internals | Only public re-exports imported (§5); workspace CI breaks loudly on drift | +| Pyroki server robot ≠ embodiment robot | README documents the per-embodiment URDF pairing and the actual invocation (upstream's `__main__` ignores CLI flags, §1); IK results that violate the action box get clamped by approvers and reported next turn | +| Cold-start latency of model servers | Retry-with-backoff bounded by `request_timeout_s`; error text names the launch command | + +## 12. Execution steps (each a commit) + +1. Agent-plugin re-exports + version bump; plugin skeleton (pyproject, + `__init__`, empty modules, README stub); `uv lock`; workspace sync green. +2. `_codec.py` + golden wire tests. +3. `_servers.py` + MockTransport stubs + tests. +4. `_sandbox.py` + `_motion.py` + tests. +5. `policy.py` (bind/reset/act/transcript/close) + tests + the joint-space + mock-embodiment e2e (§9). +6. CI job into both needs lists; release job; root README + CLAUDE.md + + plugin README. +7. Gates: `ruff check`, `ruff format --check`, plugin mypy strict, plugin + pytest, core suite untouched (`uv run pytest --cov` still 100%). diff --git a/plugins/inspect-robots-agent/pyproject.toml b/plugins/inspect-robots-agent/pyproject.toml index e44e568..67cf4a4 100644 --- a/plugins/inspect-robots-agent/pyproject.toml +++ b/plugins/inspect-robots-agent/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "inspect-robots-agent" -version = "0.11.0" +version = "0.12.0" description = "LLM agent policy for Inspect Robots: frontier LLMs (Claude, GPT, anything OpenAI-compatible) drive any registered embodiment through tool calls." dynamic = ["readme"] requires-python = ">=3.10" diff --git a/plugins/inspect-robots-agent/src/inspect_robots_agent/__init__.py b/plugins/inspect-robots-agent/src/inspect_robots_agent/__init__.py index 7402314..49f20bc 100644 --- a/plugins/inspect-robots-agent/src/inspect_robots_agent/__init__.py +++ b/plugins/inspect-robots-agent/src/inspect_robots_agent/__init__.py @@ -23,9 +23,30 @@ from importlib.metadata import version +from inspect_robots_agent._llm import ( + ENV_MODEL, + AssistantMessage, + ChatClient, + Provider, + resolve_provider, +) +from inspect_robots_agent._png import encode_png, png_data_url +from inspect_robots_agent._responses import ResponsesClient from inspect_robots_agent.policy import AgentPolicyConfig, LLMAgentPolicy, agent_policy -__all__ = ["AgentPolicyConfig", "LLMAgentPolicy", "agent_policy"] +__all__ = [ + "ENV_MODEL", + "AgentPolicyConfig", + "AssistantMessage", + "ChatClient", + "LLMAgentPolicy", + "Provider", + "ResponsesClient", + "agent_policy", + "encode_png", + "png_data_url", + "resolve_provider", +] # Derived from package metadata so it can never drift from pyproject again # (0.2.0 and 0.2.1 shipped with a stale hardcoded "0.1.0"). diff --git a/plugins/inspect-robots-agent/tests/test_package.py b/plugins/inspect-robots-agent/tests/test_package.py index 6a42f16..547543f 100644 --- a/plugins/inspect-robots-agent/tests/test_package.py +++ b/plugins/inspect-robots-agent/tests/test_package.py @@ -6,5 +6,18 @@ def test_package_imports_and_exports() -> None: # Pinned so a version bump that misses either side fails loudly (the # 0.1.0 hardcode shipped stale through two releases unnoticed). - assert inspect_robots_agent.__version__ == "0.11.0" + assert inspect_robots_agent.__version__ == "0.12.0" assert callable(inspect_robots_agent.agent_policy) + assert inspect_robots_agent.__all__ == [ + "ENV_MODEL", + "AgentPolicyConfig", + "AssistantMessage", + "ChatClient", + "LLMAgentPolicy", + "Provider", + "ResponsesClient", + "agent_policy", + "encode_png", + "png_data_url", + "resolve_provider", + ] diff --git a/plugins/inspect-robots-capx/README.md b/plugins/inspect-robots-capx/README.md new file mode 100644 index 0000000..e832f48 --- /dev/null +++ b/plugins/inspect-robots-capx/README.md @@ -0,0 +1,238 @@ +# inspect-robots-capx: + +Code-as-policy manipulation for +[Inspect Robots](https://github.com/robocurve/inspect-robots), backed by the +SAM3, Contact-GraspNet, and Pyroki servers from +[CaP-X](https://github.com/capgym/cap-x). The installed policy name is `capx`. + +## Install: + +Install the core and plugin in the environment that drives the embodiment: + +```bash +pip install inspect-robots inspect-robots-capx +``` + +CaP-X is not a Python dependency of this package. Its research checkout and +GPU dependencies stay in the separate server environment. + +## Start the CaP-X servers: + +From an installed CaP-X checkout, start one process for each model: + +```bash +# terminal 1 +uv run capx/serving/launch_sam3_server.py --port 8114 + +# terminal 2 +uv run capx/serving/launch_contact_graspnet_server.py --port 8115 + +# terminal 3 +uv run python -c 'from capx.serving.launch_pyroki_server import main; main(robot="panda_description", port=8116)' +``` + +The Pyroki launcher is different from the other two. Upstream +`launch_pyroki_server.py` calls a bare `main()` and does not parse command-line +flags, so invoking the file with `--robot` or `--port` silently ignores them. +Use the `python -c` form above and select the robot description that matches +the embodiment. The server's full actuated configuration may include a +gripper joint; the client strips the returned vector to the bound arm size and +keeps the full vector only as the next IK warm start. + +## Run a policy: + +The v1 profile needs a joint-space embodiment. The core `cubepick` mock uses +2-D end-effector deltas and is intentionally incompatible. + +```bash +export ANTHROPIC_API_KEY=sk-ant-... + +inspect-robots "pick up the red cube" --policy capx \ + --embodiment \ + -P model=anthropic/claude-fable-5 \ + -P sam3_url=http://gpu-box:8114 \ + -P graspnet_url=http://gpu-box:8115 \ + -P pyroki_url=http://gpu-box:8116 +``` + +Model routing, API key selection, and the `chat` versus `responses` wire follow +the [inspect-robots-agent](../inspect-robots-agent/) plugin. Model strings come +from `-P model=...` or `$INSPECT_ROBOTS_MODEL`. + +## Required embodiment profile: + +`bind()` checks the complete profile before a rollout begins: + +- The action space is a one-dimensional `Box` with at least two dimensions. +- `ActionSemantics.control_mode` is `joint_pos` and `gripper` is not `none`. +- A `dim_labels` entry named `gripper` locates the gripper. If labels are + absent, the final action dimension is the fallback. +- The action box has finite low and high bounds. The selected gripper bound is + high for open and low for closed by default. Set + `gripper_open_is_high=false` for an inverted embodiment. +- Exactly one `StateSpec` field has the full action-vector shape. It is the + proprioceptive reference for motion and stop holds. +- `control_hz` is finite and positive. +- `camera=None` selects the sole declared camera. Set `camera=NAME` when the + embodiment declares several cameras. + +Pyroki solves in its URDF base frame. This plugin treats world and robot base +as the same frame. Match the server's robot model, the embodiment's action +dimensions, and the task's coordinates. + +## Depth and camera metadata: + +Core observations carry RGB images, but do not reserve a depth slot. An +embodiment can opt into grasp planning with entries in `Observation.extra`: + +```python +Observation( + images={"front": rgb}, + state={"joint_pos": joint_config}, + extra={ + "depth": lambda: camera.read_depth(), + "intrinsics": lambda: camera.intrinsics(), + "extrinsics": lambda: camera.camera_to_base(), + }, +) +``` + +Each value may be the array itself or a zero-argument callable returning it. +The callable form is recommended for real embodiments. Trial records retain +`Observation.extra` at every control step, so raw per-step depth arrays can +otherwise become an in-memory depth-video buffer. The defaults expect depth +shape `(H, W)`, intrinsics shape `(3, 3)`, and camera-to-base extrinsics shape +`(4, 4)`. The `depth_key`, `intrinsics_key`, and `extrinsics_key` arguments can +rename them. + +Missing depth does not block segmentation or IK. `plan_grasp()` raises inside +the code namespace with a message naming the missing key, so the model sees the +error on stderr and can choose another route. Contact-GraspNet poses are in the +camera frame; model code transforms them with `extrinsics @ pose` before IK. + +## Execution model: + +Each LLM turn returns raw Python, optionally wrapped in one Markdown code +fence. The namespace persists for the trial, including variables and imports. +The helpers expose the current observation as `obs` and provide: + +- `segment(text)` for SAM3 text-prompt masks +- `plan_grasp(mask)` for camera-frame grasp poses and scores +- `solve_ik(position, quaternion_wxyz)` for arm joints +- `move_to_joints(joints)` for a speed-limited joint target +- `open_gripper()` and `close_gripper()` for speed-limited gripper ramps + +Motion helpers queue full joint targets. The complete queue becomes one open +loop `ActionChunk`, so perception in that Python turn uses the initial +observation. The next policy turn observes the executed result. Within a turn, +the motion cursor chains across every arm and gripper call. At the next turn it +is seeded again from observed proprioception. + +Interpolation uses `max_speed_frac / control_hz` of each action-box range per +step, capped by the same native-dtype 5 percent backstop as the core +`DeltaLimitApprover`. This prevents default guardrails from silently rewriting +steps within a chunk. A lagging arm can still make the first target of the next +turn differ from the last approved action, so the approver may clamp that turn +boundary. + +After execution, the LLM receives its code, captured stdout, and captured +stderr with any traceback. Execution reports keep their final 16,000 +characters when truncation is needed, since exceptions appear at the tail. +The report is appended to the transcript before `act()` returns. The full, +untruncated code is also stored in the first queued action's `meta["code"]`. + +`FINISH` and `GIVE_UP` return a one-action hold with the core policy-stop +metadata. A clean turn with no queued motion continues inside the same +`act()` call. `max_llm_calls` bounds all turns in a trial. +`max_code_failures` counts consecutive Python exceptions across `act()` calls. +A clean Python turn resets it. Code that queues motion and then raises still +returns that motion chunk, while the error remains counted for the next turn. + +## Trust model: + +> [!WARNING] +> Model-generated Python executes in-process with the evaluator's user +> privileges. It can import installed packages, read files, make network +> requests, and mutate process state. This integration follows CaP-X; it is not +> a security sandbox. Run untrusted models inside a container or another +> external isolation boundary. + +Every queued robot action still passes through the rollout approver. CLI runs +install `ClampApprover` plus `DeltaLimitApprover` by default. The Python +`eval()` API defaults to the permissive `AutoApprover`, so programmatic real +robot runs must wire guardrails explicitly: + +```python +from inspect_robots import eval +from inspect_robots.approver import ChainApprover, ClampApprover, DeltaLimitApprover +from inspect_robots_capx import CapxPolicy + +# `embodiment` is your constructed joint-space adapter. +space = embodiment.info.action_space +guardrails = ChainApprover(ClampApprover(space), DeltaLimitApprover(space)) +policy = CapxPolicy(model="anthropic/claude-fable-5") + +eval(task, policy, embodiment, approver=guardrails) +``` + +Guardrails constrain robot actions. They do not constrain what executed Python +can do to the evaluator process. + +## Configuration: + +CLI policy arguments use `-P key=value`. + +| Argument | Default | Contract | +|---|---|---| +| `model` | `$INSPECT_ROBOTS_MODEL` | OpenRouter-style model id or provider-prefixed direct model. | +| `base_url` | provider routing | Custom OpenAI-compatible endpoint. | +| `api_key_env` | `OPENROUTER_API_KEY` with custom URL | Environment variable holding a custom endpoint key. | +| `wire` | `chat` | `chat` or `responses`. | +| `temperature` | omitted | Sampling temperature sent when set. | +| `effort` | `low` | Reasoning effort, or null to omit the field. | +| `sam3_url` | `http://127.0.0.1:8114` | SAM3 server base URL. | +| `graspnet_url` | `http://127.0.0.1:8115` | Contact-GraspNet server base URL. | +| `pyroki_url` | `http://127.0.0.1:8116` | Pyroki server base URL. | +| `camera` | sole camera | RGB camera used by perception. | +| `depth_key` | `depth` | `Observation.extra` depth entry. | +| `intrinsics_key` | `intrinsics` | `Observation.extra` camera intrinsics entry. | +| `extrinsics_key` | `extrinsics` | `Observation.extra` camera-to-base transform entry. | +| `max_llm_calls` | `100` | Total LLM calls per trial. | +| `max_code_failures` | `3` | Consecutive exception turns before policy failure. | +| `max_speed_frac` | `0.1` | Action-range fraction per second before the per-step backstop. | +| `request_timeout_s` | `120` | Total wall-clock budget for one helper HTTP call, retries included. | +| `gripper_open_is_high` | `true` | Use the gripper high bound as open. | +| `transcript_echo` | `false` | Echo code and execution feedback to stderr. | +| `transport` | `None` | Programmatic `httpx` transport injection for tests. | +| `env` | process environment | Programmatic provider-environment override for tests. | + +Each helper request retries transient transport errors, HTTP 429, and server +5xx responses with exponential backoff. The total retry loop stays within +`request_timeout_s`; one attempt receives `min(remaining_budget, 30 seconds)`. + +## Transcripts: + +`CapxPolicy.transcript()` returns the current conversation as an isolated copy. +Camera data URLs are replaced with omission markers, while code, stdout, +stderr, and stop words remain. Inspect a saved run with: + +```bash +inspect-robots inspect LOG.json --transcript +inspect-robots view LOG.json +``` + +## Troubleshooting: + +- A 503 during initial use usually means a model server is still loading. + Requests retry within `request_timeout_s`; increase that value for slow GPU + cold starts. +- A connection error names the failed URL and the exact launch command for that + service. Check host routing and firewall access from the evaluator machine. +- A Pyroki response shorter than the arm size means the server robot and + embodiment do not match. Restart Pyroki with the correct robot description. +- Five-degree-of-freedom arms may not reach arbitrary six-degree-of-freedom + grasp poses. Filter for top-down grasps in the task prompt or use a matching + IK target convention. +- Direct OpenAI reasoning models may require `-P wire=responses`. Chat-only + endpoints should keep `-P wire=chat`; `-P effort=none` can disable reasoning + where function or code-generation requests reject it. diff --git a/plugins/inspect-robots-capx/pyproject.toml b/plugins/inspect-robots-capx/pyproject.toml new file mode 100644 index 0000000..1f5545a --- /dev/null +++ b/plugins/inspect-robots-capx/pyproject.toml @@ -0,0 +1,89 @@ +[build-system] +requires = ["hatchling>=1.18", "hatch-fancy-pypi-readme>=24.1"] +build-backend = "hatchling.build" + +[project] +name = "inspect-robots-capx" +version = "0.1.0" +description = "Code-as-policy manipulation for Inspect Robots using CaP-X perception and IK servers." +dynamic = ["readme"] +requires-python = ">=3.10" +license = "MIT" +authors = [{ name = "Inspect Robots contributors" }] +keywords = ["robotics", "llm", "code-as-policy", "capx", "inspect_robots", "evaluation"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Typing :: Typed", +] +# CaP-X itself is intentionally not a dependency. The research checkout owns +# the model servers; this adapter speaks their small HTTP protocol directly. +dependencies = [ + "inspect-robots>=0.4", + "inspect-robots-agent>=0.12", + "numpy>=1.24", + "httpx>=0.27", +] + +[project.optional-dependencies] +dev = ["pytest>=8.0", "pytest-cov>=5.0", "mypy>=1.11", "ruff>=0.6"] + +[project.entry-points."inspect_robots.policies"] +capx = "inspect_robots_capx.policy:capx_policy" + +[project.urls] +Homepage = "https://github.com/robocurve/inspect-robots/tree/main/plugins/inspect-robots-capx" +Repository = "https://github.com/robocurve/inspect-robots" + +[tool.uv.sources] +inspect-robots = { workspace = true } +inspect-robots-agent = { workspace = true } + +[tool.hatch.build.targets.wheel] +packages = ["src/inspect_robots_capx"] + +[tool.hatch.build.targets.sdist] +include = ["src/inspect_robots_capx", "tests", "README.md"] + +[tool.ruff] +extend = "../../pyproject.toml" + +[tool.ruff.lint.isort] +known-first-party = ["inspect_robots", "inspect_robots_agent", "inspect_robots_capx"] + +[tool.ruff.lint.per-file-ignores] +"tests/**" = ["D1"] + +[tool.mypy] +strict = true + +# PyPI readme: identical to README.md except GitHub-only alert syntax +# (e.g. `> [!NOTE]`) is rewritten to bold blockquotes, which PyPI renders. +[tool.hatch.metadata.hooks.fancy-pypi-readme] +content-type = "text/markdown" + +[[tool.hatch.metadata.hooks.fancy-pypi-readme.fragments]] +path = "README.md" + +[[tool.hatch.metadata.hooks.fancy-pypi-readme.substitutions]] +pattern = '(?m)^> \[!NOTE\][ \t]*$' +replacement = '> **Note:**' + +[[tool.hatch.metadata.hooks.fancy-pypi-readme.substitutions]] +pattern = '(?m)^> \[!TIP\][ \t]*$' +replacement = '> **Tip:**' + +[[tool.hatch.metadata.hooks.fancy-pypi-readme.substitutions]] +pattern = '(?m)^> \[!IMPORTANT\][ \t]*$' +replacement = '> **Important:**' + +[[tool.hatch.metadata.hooks.fancy-pypi-readme.substitutions]] +pattern = '(?m)^> \[!WARNING\][ \t]*$' +replacement = '> **Warning:**' + +[[tool.hatch.metadata.hooks.fancy-pypi-readme.substitutions]] +pattern = '(?m)^> \[!CAUTION\][ \t]*$' +replacement = '> **Caution:**' diff --git a/plugins/inspect-robots-capx/src/inspect_robots_capx/__init__.py b/plugins/inspect-robots-capx/src/inspect_robots_capx/__init__.py new file mode 100644 index 0000000..a4e9a3f --- /dev/null +++ b/plugins/inspect-robots-capx/src/inspect_robots_capx/__init__.py @@ -0,0 +1,11 @@ +"""CaP-X code-as-policy integration for Inspect Robots.""" + +from __future__ import annotations + +from importlib.metadata import version + +from inspect_robots_capx.policy import CapxPolicy, CapxPolicyConfig, capx_policy + +__all__ = ["CapxPolicy", "CapxPolicyConfig", "capx_policy"] + +__version__ = version("inspect-robots-capx") diff --git a/plugins/inspect-robots-capx/src/inspect_robots_capx/_codec.py b/plugins/inspect-robots-capx/src/inspect_robots_capx/_codec.py new file mode 100644 index 0000000..51491d2 --- /dev/null +++ b/plugins/inspect-robots-capx/src/inspect_robots_capx/_codec.py @@ -0,0 +1,63 @@ +"""CaP-X HTTP wire codecs isolated for protocol-drift review. + +The schemas mirror ``capx/serving`` in https://github.com/capgym/cap-x as +validated against upstream commit +``53e9966d7a8e2fa7494676772bccc35280f5c0ed`` (2026-04-09). The protocol has no +version field, so the small golden tests below this module pin the exact bytes +and shapes that were inspected. +""" + +from __future__ import annotations + +import base64 +import io +from typing import Any, cast + +import numpy as np +import numpy.typing as npt + +from inspect_robots_agent import encode_png + + +def png_b64_encode(rgb: npt.NDArray[Any]) -> str: + """Encode one RGB observation as bare base64 PNG for SAM3 requests.""" + return base64.b64encode(encode_png(rgb)).decode("ascii") + + +def npy_b64_encode(array: npt.NDArray[Any]) -> str: + """Encode an array as a base64 ``.npy`` payload without object pickles.""" + buffer = io.BytesIO() + np.save(buffer, np.asarray(array), allow_pickle=False) + return base64.b64encode(buffer.getvalue()).decode("ascii") + + +def npy_b64_decode(payload: str) -> npt.NDArray[Any]: + """Decode a CaP-X base64 ``.npy`` payload while refusing object pickles.""" + raw = base64.b64decode(payload, validate=True) + decoded = np.load(io.BytesIO(raw), allow_pickle=False) + return cast(npt.NDArray[Any], decoded) + + +def mask_decode(payload: str, shape: tuple[int, int]) -> npt.NDArray[np.bool_]: + """Decode a SAM3 raw-byte mask and normalize it to a boolean image.""" + raw = base64.b64decode(payload, validate=True) + return np.frombuffer(raw, dtype=np.uint8).reshape(shape).astype(np.bool_) + + +def grasp_arrays_decode( + grasps_payload: str, + scores_payload: str, +) -> tuple[npt.NDArray[Any], npt.NDArray[Any]]: + """Decode grasp poses/scores and normalize CaP-X's flat empty arrays. + + Contact-GraspNet serializes a no-grasp result as two ``shape == (0,)`` + arrays. Callers always receive the ordinary public shapes ``(K, 4, 4)`` + and ``(K,)``, including ``K == 0``. + """ + grasps = npy_b64_decode(grasps_payload) + scores = npy_b64_decode(scores_payload) + if grasps.size == 0: + grasps = grasps.reshape((0, 4, 4)) + if scores.size == 0: + scores = scores.reshape((0,)) + return grasps, scores diff --git a/plugins/inspect-robots-capx/src/inspect_robots_capx/_motion.py b/plugins/inspect-robots-capx/src/inspect_robots_capx/_motion.py new file mode 100644 index 0000000..8a9a144 --- /dev/null +++ b/plugins/inspect-robots-capx/src/inspect_robots_capx/_motion.py @@ -0,0 +1,212 @@ +"""Joint-space motion synthesis for model-generated policy code. + +The per-step ceiling reproduces ``DeltaLimitApprover``'s native-dtype +``0.05 * (high - low)`` arithmetic. The speed fraction may make the ceiling +tighter, but never looser, so actions within one returned chunk are not +silently rewritten by the CLI's default delta backstop. +""" + +from __future__ import annotations + +import math +from dataclasses import replace +from typing import Any + +import numpy as np +import numpy.typing as npt + +from inspect_robots.spaces import Box +from inspect_robots.types import Action, ActionChunk + +_BACKSTOP_STEP_FRAC = 0.05 +_RELATIVE_HEADROOM = 1e-6 + + +class MotionQueue: + """Queue interpolated full-config targets around a turn-local state cursor.""" + + def __init__( + self, + action_space: Box, + *, + control_hz: float, + max_speed_frac: float, + gripper_index: int, + gripper_open_is_high: bool = True, + ) -> None: + if len(action_space.shape) != 1: + raise ValueError(f"motion needs a 1-D action box, got {action_space.shape}") + if not np.isfinite(control_hz) or control_hz <= 0: + raise ValueError("control_hz must be finite and > 0") + if not np.isfinite(max_speed_frac) or max_speed_frac <= 0: + raise ValueError("max_speed_frac must be finite and > 0") + if not 0 <= gripper_index < action_space.dim: + raise ValueError(f"gripper index {gripper_index} is outside {action_space.dim}-D box") + low = action_space.low + high = action_space.high + if low is None or high is None: + raise ValueError("motion needs finite low and high action bounds") + if not bool(np.all(np.isfinite(low)) and np.all(np.isfinite(high))): + raise ValueError("motion needs finite low and high action bounds") + + low64 = np.asarray(low, dtype=np.float64) + high64 = np.asarray(high, dtype=np.float64) + with np.errstate(over="ignore"): + native_range = np.asarray(high - low, dtype=np.float64) + native_backstop = np.asarray(_BACKSTOP_STEP_FRAC * (high - low), dtype=np.float64) + float64_range = high64 - low64 + if not bool(np.all(np.isfinite(native_range)) and np.all(np.isfinite(float64_range))): + raise ValueError("action-space range (high - low) overflows; bounds are too large") + movable = high64 > low64 + # Interpolants snap to the float grid at the bounds' magnitude; if that + # grid is coarse relative to the backstop (offset boxes like + # [1e16, 1e16 + 2]), emitted steps can exceed the approver's budget. + spacing = np.spacing(np.maximum(np.abs(low64), np.abs(high64))) + if bool(np.any(movable & (spacing > 5e-7 * native_backstop))): + raise ValueError( + "bounds are too coarse at this magnitude for speed-limited " + "interpolation (float spacing exceeds the per-step budget)" + ) + step_frac = min(max_speed_frac / control_hz, _BACKSTOP_STEP_FRAC) + step_limits = np.minimum(step_frac * float64_range, native_backstop) + if bool(np.any(movable & (step_limits <= 0))): + raise ValueError("speed fraction underflows the per-step limit for a movable dimension") + + self.control_hz = control_hz + self._dim = action_space.dim + self._gripper_index = gripper_index + self._arm_indices = tuple(index for index in range(self._dim) if index != gripper_index) + self._low = low64 + self._high = high64 + self._step_limits = step_limits + high_value = float(high64[gripper_index]) + low_value = float(low64[gripper_index]) + self.gripper_open_value = high_value if gripper_open_is_high else low_value + self.gripper_closed_value = low_value if gripper_open_is_high else high_value + self._cursor: npt.NDArray[np.float64] | None = None + self._actions: list[Action] = [] + + @property + def arm_dof(self) -> int: + """Return the joint count expected by ``move_to_joints`` and Pyroki.""" + return len(self._arm_indices) + + @property + def cursor(self) -> npt.NDArray[np.float64] | None: + """Return a defensive copy of the current queued endpoint, if seeded.""" + return None if self._cursor is None else self._cursor.copy() + + def reset(self) -> None: + """Clear queued actions and the observation-derived cursor.""" + self._cursor = None + self._actions.clear() + + def begin_turn(self, state: npt.NDArray[Any]) -> None: + """Discard old queue state and seed this turn from observed proprioception.""" + cursor = np.asarray(state, dtype=np.float64).reshape(-1) + if cursor.shape != (self._dim,): + raise ValueError( + f"proprioceptive reference has shape {cursor.shape}, expected ({self._dim},)" + ) + if not bool(np.all(np.isfinite(cursor))): + raise ValueError("proprioceptive reference contains a non-finite value") + self._cursor = cursor.copy() + self._actions.clear() + + def move_to_joints(self, joints: npt.NDArray[Any]) -> None: + """Queue an arm target while preserving the cursor's gripper value.""" + cursor = self._require_cursor() + arm = np.asarray(joints, dtype=np.float64).reshape(-1) + if arm.shape != (self.arm_dof,): + raise ValueError( + f"move_to_joints expects {self.arm_dof} arm joints, got shape {arm.shape}" + ) + if not bool(np.all(np.isfinite(arm))): + raise ValueError("move_to_joints received a non-finite joint target") + target = cursor.copy() + target[list(self._arm_indices)] = arm + self._queue_target(target) + + def open_gripper(self) -> None: + """Queue a speed-limited ramp to the configured open box bound.""" + self._queue_gripper(self.gripper_open_value) + + def close_gripper(self) -> None: + """Queue a speed-limited ramp to the configured closed box bound.""" + self._queue_gripper(self.gripper_closed_value) + + def has_actions(self) -> bool: + """Report whether model code queued any targets in the current turn.""" + return bool(self._actions) + + def take_chunk( + self, + *, + inference_latency_s: float | None = None, + code: str | None = None, + ) -> ActionChunk: + """Drain the non-empty queue into a chunk and annotate its first action.""" + if not self._actions: + raise ValueError("motion queue is empty") + actions = self._actions + self._actions = [] + if code is not None: + actions[0] = replace(actions[0], meta={**dict(actions[0].meta), "code": code}) + return ActionChunk( + actions=actions, + control_hz=self.control_hz, + inference_latency_s=inference_latency_s, + ) + + def hold_chunk( + self, + state: npt.NDArray[Any], + *, + stop_reason: str, + inference_latency_s: float | None = None, + ) -> ActionChunk: + """Return the required one-action stop chunk at the observed full config.""" + value = np.asarray(state, dtype=np.float64).reshape(-1) + if value.shape != (self._dim,): + raise ValueError(f"hold state has shape {value.shape}, expected ({self._dim},)") + action = Action( + data=value.copy(), + meta={"request_stop": True, "stop_reason": stop_reason}, + ) + return ActionChunk( + actions=[action], + control_hz=self.control_hz, + inference_latency_s=inference_latency_s, + ) + + def _require_cursor(self) -> npt.NDArray[np.float64]: + if self._cursor is None: + raise RuntimeError("motion cursor is unset; provide an observation before moving") + return self._cursor + + def _queue_gripper(self, value: float) -> None: + target = self._require_cursor().copy() + target[self._gripper_index] = value + self._queue_target(target) + + def _queue_target(self, target: npt.NDArray[np.float64]) -> None: + current = self._require_cursor() + delta = target - current + ratios: list[np.float64] = [] + with np.errstate(over="ignore", invalid="ignore", divide="ignore"): + for index, distance in enumerate(np.abs(delta)): + limit = self._step_limits[index] + if distance == 0: + continue + if limit <= 0: + raise ValueError(f"action dimension {index} is fixed and cannot move") + ratios.append(np.divide(distance, limit)) + ratio = max(ratios, default=np.float64(0.0)) + headed_ratio = np.divide(ratio, 1.0 - _RELATIVE_HEADROOM) + if not np.isfinite(headed_ratio): + raise ValueError("motion distance is too large to interpolate safely") + steps = max(1, math.ceil(float(headed_ratio))) + for fraction in np.linspace(1.0 / steps, 1.0, steps): + self._actions.append(Action(data=current + delta * fraction)) + self._actions[-1] = Action(data=target.copy()) + self._cursor = target.copy() diff --git a/plugins/inspect-robots-capx/src/inspect_robots_capx/_sandbox.py b/plugins/inspect-robots-capx/src/inspect_robots_capx/_sandbox.py new file mode 100644 index 0000000..2492a8d --- /dev/null +++ b/plugins/inspect-robots-capx/src/inspect_robots_capx/_sandbox.py @@ -0,0 +1,170 @@ +"""In-process execution boundary for model-generated Python code. + +Model output executes with the evaluator's process privileges. This module is +an integration surface, not a security sandbox; untrusted models require an +external container or equivalent isolation. +""" + +from __future__ import annotations + +import contextlib +import io +import traceback +from dataclasses import dataclass +from typing import Any + +import numpy as np +import numpy.typing as npt + +from inspect_robots.types import Observation +from inspect_robots_capx._motion import MotionQueue +from inspect_robots_capx._servers import CapxServerClients + + +@dataclass(frozen=True) +class ExecutionResult: + """Captured feedback from one model-code turn without propagating its exception.""" + + stdout: str + stderr: str + raised: bool + + +class _TurnView(dict[str, Any]): + """Turn-scoped ``obs`` mapping that resolves zero-arg callable values on access. + + Embodiments may provide bulk ``observation.extra`` entries (depth, + intrinsics, extrinsics) as zero-argument callables to keep trial records + small; resolving them here keeps the documented ``obs[...]`` idioms + working unchanged for both forms. + """ + + def __getitem__(self, key: str) -> Any: + value = super().__getitem__(key) + return value() if callable(value) else value + + def get(self, key: str, default: Any = None) -> Any: + """Return the resolved value for ``key``, or ``default`` when absent. + + Membership decides absence, so a ``KeyError`` raised inside an + embodiment-provided thunk propagates instead of masquerading as a + missing key. + """ + if key not in self: + return default + return self[key] + + +class CodeSandbox: + """Persistent per-trial Python namespace with observation-bound robot helpers.""" + + def __init__( + self, + *, + servers: CapxServerClients, + motion: MotionQueue, + camera: str, + state_key: str, + depth_key: str = "depth", + intrinsics_key: str = "intrinsics", + extrinsics_key: str = "extrinsics", + ) -> None: + self._servers = servers + self._motion = motion + self._camera = camera + self._state_key = state_key + self._depth_key = depth_key + self._intrinsics_key = intrinsics_key + self._extrinsics_key = extrinsics_key + self._observation: Observation | None = None + self._namespace: dict[str, Any] = {} + self.reset() + + def reset(self) -> None: + """Start a fresh trial namespace with only the documented helpers bound.""" + self._observation = None + self._motion.reset() + self._namespace = { + "__builtins__": __builtins__, + "segment": self._segment, + "plan_grasp": self._plan_grasp, + "solve_ik": self._solve_ik, + "move_to_joints": self._motion.move_to_joints, + "open_gripper": self._motion.open_gripper, + "close_gripper": self._motion.close_gripper, + } + + def set_observation(self, observation: Observation) -> None: + """Expose one turn's observation and reseed motion from its full state field.""" + try: + state = observation.state[self._state_key] + except KeyError as exc: + raise ValueError( + f"observation.state is missing bound proprioceptive key {self._state_key!r}" + ) from exc + self._observation = observation + self._motion.begin_turn(state) + obs = _TurnView( + { + "images": observation.images, + "state": observation.state, + } + ) + obs.update(observation.extra) + self._namespace["obs"] = obs + + def execute(self, code: str) -> ExecutionResult: + """Execute one turn and return stdout, stderr, and whether it raised.""" + stdout = io.StringIO() + stderr = io.StringIO() + raised = False + with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr): + try: + exec(code, self._namespace, self._namespace) + except Exception: + raised = True + traceback.print_exc() + return ExecutionResult(stdout=stdout.getvalue(), stderr=stderr.getvalue(), raised=raised) + + def _require_observation(self) -> Observation: + if self._observation is None: + raise RuntimeError("no observation is bound for this code turn") + return self._observation + + def _extra_array(self, key: str) -> npt.NDArray[Any]: + observation = self._require_observation() + if key not in observation.extra: + raise KeyError( + f"observation.extra is missing {key!r}; embodiments must provide " + f"extra[{key!r}] as an array or a zero-argument callable returning one" + ) + value = observation.extra[key] + if callable(value): + value = value() + return np.asarray(value) + + def _segment(self, text: str) -> list[dict[str, Any]]: + observation = self._require_observation() + try: + image = observation.images[self._camera] + except KeyError as exc: + raise KeyError(f"observation has no configured camera {self._camera!r}") from exc + return self._servers.sam3.segment(image, text) + + def _plan_grasp(self, mask: npt.NDArray[Any]) -> tuple[npt.NDArray[Any], npt.NDArray[Any]]: + return self._servers.graspnet.plan( + self._extra_array(self._depth_key), + self._extra_array(self._intrinsics_key), + np.asarray(mask), + ) + + def _solve_ik( + self, + position: npt.NDArray[Any], + quaternion_wxyz: npt.NDArray[Any], + ) -> npt.NDArray[np.float64]: + return self._servers.pyroki.solve_ik( + position, + quaternion_wxyz, + self._motion.arm_dof, + ) diff --git a/plugins/inspect-robots-capx/src/inspect_robots_capx/_servers.py b/plugins/inspect-robots-capx/src/inspect_robots_capx/_servers.py new file mode 100644 index 0000000..3978bb4 --- /dev/null +++ b/plugins/inspect-robots-capx/src/inspect_robots_capx/_servers.py @@ -0,0 +1,301 @@ +"""Lazy HTTP clients for the CaP-X perception and inverse-kinematics servers.""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from typing import Any, cast + +import httpx +import numpy as np +import numpy.typing as npt + +from inspect_robots_capx._codec import ( + grasp_arrays_decode, + mask_decode, + npy_b64_encode, + png_b64_encode, +) + +_ATTEMPT_TIMEOUT_S = 30.0 + +_SAM3_COMMAND = "uv run capx/serving/launch_sam3_server.py --port 8114" +_GRASPNET_COMMAND = "uv run capx/serving/launch_contact_graspnet_server.py --port 8115" +_PYROKI_COMMAND = ( + "uv run python -c 'from capx.serving.launch_pyroki_server import main; " + 'main(robot="panda_description", port=8116)\'' +) + + +class _ServerClient: + """Retrying JSON POST boundary shared by one logical CaP-X service.""" + + def __init__( + self, + url: str, + launch_command: str, + *, + request_timeout_s: float, + http: httpx.Client | None = None, + transport: httpx.BaseTransport | None = None, + backoff_s: float = 1.0, + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, + ) -> None: + if not np.isfinite(request_timeout_s) or request_timeout_s <= 0: + raise ValueError("request_timeout_s must be finite and > 0") + if not np.isfinite(backoff_s) or backoff_s < 0: + raise ValueError("backoff_s must be finite and >= 0") + if http is not None and transport is not None: + raise ValueError("pass either a shared http client or transport, not both") + self._base_url = url.rstrip("/") + self._launch_command = launch_command + self._request_timeout_s = request_timeout_s + self._backoff_s = backoff_s + self._clock = clock + self._sleep = sleep + self._owns_http = http is None + self._http = http or httpx.Client(transport=transport) + + def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any]: + url = self._base_url + path + deadline = self._clock() + self._request_timeout_s + attempt = 0 + last_error = "server did not respond" + timed_out = False + while True: + remaining = deadline - self._clock() + if remaining <= 0: + error_type = TimeoutError if timed_out else ConnectionError + raise error_type(self._guidance(url, last_error)) + try: + response = self._http.post( + url, + json=body, + timeout=min(remaining, _ATTEMPT_TIMEOUT_S), + ) + except httpx.TimeoutException as exc: + timed_out = True + last_error = f"{type(exc).__name__}: {exc}" + except httpx.TransportError as exc: + last_error = f"{type(exc).__name__}: {exc}" + else: + if response.status_code == 200: + payload = response.json() + if not isinstance(payload, dict): + raise ConnectionError( + self._guidance(url, "server returned a non-object JSON payload") + ) + return cast(dict[str, Any], payload) + last_error = f"HTTP {response.status_code}: {response.text[:500]}" + if response.status_code not in (429,) and response.status_code < 500: + raise ConnectionError(self._guidance(url, last_error)) + + remaining = deadline - self._clock() + if remaining <= 0: + continue + delay = min(self._backoff_s * 2**attempt, remaining) + self._sleep(delay) + attempt += 1 + + def _guidance(self, url: str, detail: str) -> str: + return ( + f"CaP-X server request failed at {url}: {detail}.\n" + f"launch it from a CaP-X checkout with: {self._launch_command}" + ) + + def close(self) -> None: + """Release an independently owned HTTP pool, if this client created one.""" + if self._owns_http: + self._http.close() + + +class Sam3Client(_ServerClient): + """Text-prompt segmentation against CaP-X's SAM3 ``/segment`` endpoint.""" + + def __init__( + self, + url: str, + *, + request_timeout_s: float = 120.0, + http: httpx.Client | None = None, + transport: httpx.BaseTransport | None = None, + backoff_s: float = 1.0, + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, + ) -> None: + super().__init__( + url, + _SAM3_COMMAND, + request_timeout_s=request_timeout_s, + http=http, + transport=transport, + backoff_s=backoff_s, + clock=clock, + sleep=sleep, + ) + + def segment(self, rgb: npt.NDArray[Any], text: str) -> list[dict[str, Any]]: + """Return boolean masks and metadata for every SAM3 text match.""" + payload = self._post( + "/segment", + {"image_base64": png_b64_encode(rgb), "text_prompt": text}, + ) + results: list[dict[str, Any]] = [] + for raw in payload["results"]: + shape = tuple(int(value) for value in raw["shape"]) + if len(shape) != 2: + raise ValueError(f"SAM3 returned invalid mask shape {shape!r}; expected (H, W)") + results.append( + { + "mask": mask_decode(raw["mask_base64"], (shape[0], shape[1])), + "box": raw["box"], + "score": raw["score"], + "label": raw["label"], + } + ) + return results + + +class GraspNetClient(_ServerClient): + """Grasp planning against CaP-X's Contact-GraspNet ``/plan`` endpoint.""" + + def __init__( + self, + url: str, + *, + request_timeout_s: float = 120.0, + http: httpx.Client | None = None, + transport: httpx.BaseTransport | None = None, + backoff_s: float = 1.0, + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, + ) -> None: + super().__init__( + url, + _GRASPNET_COMMAND, + request_timeout_s=request_timeout_s, + http=http, + transport=transport, + backoff_s=backoff_s, + clock=clock, + sleep=sleep, + ) + + def plan( + self, + depth: npt.NDArray[Any], + intrinsics: npt.NDArray[Any], + mask: npt.NDArray[Any], + ) -> tuple[npt.NDArray[Any], npt.NDArray[Any]]: + """Return camera-frame grasp poses and scores with stable empty shapes.""" + response = self._post( + "/plan", + { + "depth_base64": npy_b64_encode(np.asarray(depth)), + "cam_K_base64": npy_b64_encode(np.asarray(intrinsics)), + "segmap_base64": npy_b64_encode(np.asarray(mask, dtype=np.uint8)), + "segmap_id": 1, + "local_regions": True, + "filter_grasps": True, + "skip_border_objects": False, + "z_range": [0.2, 2.0], + "forward_passes": 2, + "max_retries": 10, + }, + ) + return grasp_arrays_decode(response["grasps_base64"], response["scores_base64"]) + + +class PyrokiClient(_ServerClient): + """Stateful IK client retaining the server's full configuration as warm-start.""" + + def __init__( + self, + url: str, + *, + request_timeout_s: float = 120.0, + http: httpx.Client | None = None, + transport: httpx.BaseTransport | None = None, + backoff_s: float = 1.0, + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, + ) -> None: + super().__init__( + url, + _PYROKI_COMMAND, + request_timeout_s=request_timeout_s, + http=http, + transport=transport, + backoff_s=backoff_s, + clock=clock, + sleep=sleep, + ) + self._prev_cfg: list[float] | None = None + + def reset(self) -> None: + """Clear the full-config IK warm-start at a trial boundary.""" + self._prev_cfg = None + + def solve_ik( + self, + position: npt.NDArray[Any], + quaternion_wxyz: npt.NDArray[Any], + arm_dof: int, + ) -> npt.NDArray[np.float64]: + """Solve one pose and return the leading embodiment arm joints only.""" + position_array = np.asarray(position, dtype=np.float64).reshape(-1) + quaternion_array = np.asarray(quaternion_wxyz, dtype=np.float64).reshape(-1) + if position_array.shape != (3,) or quaternion_array.shape != (4,): + raise ValueError("solve_ik expects position shape (3,) and quaternion_wxyz shape (4,)") + response = self._post( + "/ik", + { + "target_pose_wxyz_xyz": [ + *quaternion_array.tolist(), + *position_array.tolist(), + ], + "prev_cfg": self._prev_cfg, + }, + ) + full = np.asarray(response["joint_positions"], dtype=np.float64).reshape(-1) + if len(full) < arm_dof: + raise ValueError( + "Pyroki returned " + f"{len(full)} joints but the embodiment arm needs {arm_dof}; " + "launch the Pyroki server with the URDF matching this embodiment" + ) + self._prev_cfg = full.tolist() + return full[:arm_dof].copy() + + +class CapxServerClients: + """The three CaP-X service clients sharing one injectable HTTP connection pool.""" + + def __init__( + self, + *, + sam3_url: str, + graspnet_url: str, + pyroki_url: str, + request_timeout_s: float = 120.0, + transport: httpx.BaseTransport | None = None, + backoff_s: float = 1.0, + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, + ) -> None: + self._http = httpx.Client(transport=transport) + common: dict[str, Any] = { + "request_timeout_s": request_timeout_s, + "http": self._http, + "backoff_s": backoff_s, + "clock": clock, + "sleep": sleep, + } + self.sam3 = Sam3Client(sam3_url, **common) + self.graspnet = GraspNetClient(graspnet_url, **common) + self.pyroki = PyrokiClient(pyroki_url, **common) + + def close(self) -> None: + """Release the single shared HTTP connection pool.""" + self._http.close() diff --git a/plugins/inspect-robots-capx/src/inspect_robots_capx/policy.py b/plugins/inspect-robots-capx/src/inspect_robots_capx/policy.py new file mode 100644 index 0000000..3aafdb2 --- /dev/null +++ b/plugins/inspect-robots-capx/src/inspect_robots_capx/policy.py @@ -0,0 +1,533 @@ +"""Code-generation policy loop for the CaP-X integration. + +Model output is executed in-process with the evaluator's privileges. This is +the policy under evaluation, not a security boundary. Robot actions still pass +through the rollout's approver chain, but untrusted models must run inside an +external container or equivalent isolation. +""" + +from __future__ import annotations + +import atexit +import contextlib +import copy +import os +import re +import sys +import time +from dataclasses import dataclass +from typing import Any + +import httpx +import numpy as np + +from inspect_robots.embodiment import EmbodimentInfo +from inspect_robots.policy import PolicyBase, PolicyConfig, PolicyInfo +from inspect_robots.scene import Scene +from inspect_robots.spaces import Box +from inspect_robots.types import ActionChunk, Observation +from inspect_robots_agent import ( + ENV_MODEL, + AssistantMessage, + ChatClient, + ResponsesClient, + png_data_url, + resolve_provider, +) +from inspect_robots_capx._motion import MotionQueue +from inspect_robots_capx._sandbox import CodeSandbox, ExecutionResult +from inspect_robots_capx._servers import CapxServerClients + +_EFFORT_LEVELS = frozenset({"none", "minimal", "low", "medium", "high", "xhigh", "max"}) +_WIRE_FORMATS = frozenset({"chat", "responses"}) +_EXECUTION_REPORT_CHAR_LIMIT = 16_000 +_REPORT_TRUNCATION_MARKER = "[execution report truncated; tail follows]\n" + +_FENCED_CODE = re.compile(r"^```(?:python)?[ \t]*\n?(.*?)\n?```$", re.DOTALL | re.IGNORECASE) +_FENCED_ANYWHERE = re.compile(r"```(?:python)?[ \t]*\n(.*?)\n?```", re.DOTALL | re.IGNORECASE) +_CONTROL_WORD = re.compile(r"(FINISH|GIVE_UP)[.!]?", re.IGNORECASE) + +_HELPER_DOCS = """Helpers available in the persistent namespace: + +segment(text: str) -> list[dict] + Segment the configured camera with SAM3. Each result has a boolean `mask`, + `box`, `score`, and `label`. +plan_grasp(mask: np.ndarray) -> tuple[np.ndarray, np.ndarray] + Return (K, 4, 4) grasp poses in the camera frame and (K,) scores. K may be + zero. Use `obs["{extrinsics_key}"] @ pose` to transform a pose into the + robot-base frame. World and robot base are treated as the same frame. +solve_ik(position: np.ndarray, quaternion_wxyz: np.ndarray) -> np.ndarray + Return the arm-joint vector expected by move_to_joints. +move_to_joints(joints: np.ndarray) -> None + Queue a speed-limited arm move while holding the gripper. +open_gripper() / close_gripper() -> None + Queue a speed-limited gripper ramp while holding the arm. + +`obs` is the current turn's dict with `images`, `state`, and every +observation.extra entry. Depth (`{depth_key}`), intrinsics (`{intrinsics_key}`), +and extrinsics (`{extrinsics_key}`) may be provided by the embodiment as +arrays or zero-argument callables; `obs[...]` access and the helper functions +resolve callable values automatically. +""" + +_SYSTEM_TEMPLATE = """You generate Python code to directly solve a robot manipulation task. +Write raw Python with no Markdown fences. Helpers are already bound. Import +numpy explicitly when you need it. Variables persist across turns in this +trial. Perception within a turn uses the initial observation; queued motions +execute after your code returns, and you verify their effect next turn. + +After each execution you receive the code, stdout, and stderr. Respond with +FINISH when the goal is complete, GIVE_UP when it cannot be completed, or +REGENERATE followed by corrected Python. You may also return raw corrected +Python directly. You have {budget} LLM calls for the whole trial. + +Embodiment: {name} +Action profile: {action_summary} + +{helper_docs}""" + + +def _sanitize(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Return an image-free deep copy suitable for persistence and visualization.""" + sanitized = copy.deepcopy(messages) + for message in sanitized: + content = message.get("content") + if not isinstance(content, list): + continue + for index, part in enumerate(content): + if isinstance(part, dict) and part.get("type") == "image_url": + content[index] = { + "type": "text", + "text": "[image omitted: streamed camera frame]", + } + return sanitized + + +@dataclass(frozen=True) +class CapxPolicyConfig(PolicyConfig): + """Inference, server, motion, and transcript settings persisted in eval logs.""" + + model: str | None = None + base_url: str | None = None + api_key_env: str | None = None + wire: str = "chat" + effort: str | None = "low" + sam3_url: str = "http://127.0.0.1:8114" + graspnet_url: str = "http://127.0.0.1:8115" + pyroki_url: str = "http://127.0.0.1:8116" + camera: str | None = None + depth_key: str = "depth" + intrinsics_key: str = "intrinsics" + extrinsics_key: str = "extrinsics" + max_llm_calls: int = 100 + max_code_failures: int = 3 + max_speed_frac: float = 0.1 + request_timeout_s: float = 120.0 + gripper_open_is_high: bool = True + transcript_echo: bool = False + + +class CapxPolicy(PolicyBase): + """Runs a persistent CaP-X-style codegen conversation over a bound joint arm.""" + + def __init__( + self, + model: str | None = None, + base_url: str | None = None, + api_key_env: str | None = None, + wire: str = "chat", + max_llm_calls: int = 100, + max_code_failures: int = 3, + temperature: float | None = None, + effort: str | None = "low", + sam3_url: str = "http://127.0.0.1:8114", + graspnet_url: str = "http://127.0.0.1:8115", + pyroki_url: str = "http://127.0.0.1:8116", + camera: str | None = None, + depth_key: str = "depth", + intrinsics_key: str = "intrinsics", + extrinsics_key: str = "extrinsics", + max_speed_frac: float = 0.1, + request_timeout_s: float = 120.0, + gripper_open_is_high: bool = True, + transcript_echo: bool = False, + transport: httpx.BaseTransport | None = None, + env: dict[str, str] | None = None, + ) -> None: + if max_llm_calls < 1: + raise ValueError("max_llm_calls must be >= 1") + if max_code_failures < 1: + raise ValueError("max_code_failures must be >= 1") + if not np.isfinite(max_speed_frac) or max_speed_frac <= 0: + raise ValueError("max_speed_frac must be finite and > 0") + if not np.isfinite(request_timeout_s) or request_timeout_s <= 0: + raise ValueError("request_timeout_s must be finite and > 0") + if effort is not None and effort not in _EFFORT_LEVELS: + raise ValueError( + f"effort must be one of {sorted(_EFFORT_LEVELS)}, or None to omit the field, " + f"got {effort!r}" + ) + if wire not in _WIRE_FORMATS: + raise ValueError(f"wire must be one of {sorted(_WIRE_FORMATS)}, got {wire!r}") + + environ = dict(os.environ) if env is None else env + provider = resolve_provider( + model=model or environ.get(ENV_MODEL), + base_url=base_url, + api_key_env=api_key_env, + env=environ, + ) + if wire == "responses": + self._client: ChatClient | ResponsesClient = ResponsesClient( + provider, transport=transport + ) + else: + self._client = ChatClient(provider, transport=transport) + self._servers = CapxServerClients( + sam3_url=sam3_url, + graspnet_url=graspnet_url, + pyroki_url=pyroki_url, + request_timeout_s=request_timeout_s, + transport=transport, + ) + self._camera_config = camera + self._depth_key = depth_key + self._intrinsics_key = intrinsics_key + self._extrinsics_key = extrinsics_key + self._max_llm_calls = max_llm_calls + self._max_code_failures = max_code_failures + self._temperature = temperature + self._effort = effort + self._max_speed_frac = max_speed_frac + self._gripper_open_is_high = gripper_open_is_high + self._transcript_echo = transcript_echo + self.config = CapxPolicyConfig( + temperature=temperature, + model=provider.model, + base_url=provider.base_url, + api_key_env=api_key_env, + wire=wire, + effort=effort, + sam3_url=sam3_url, + graspnet_url=graspnet_url, + pyroki_url=pyroki_url, + camera=camera, + depth_key=depth_key, + intrinsics_key=intrinsics_key, + extrinsics_key=extrinsics_key, + max_llm_calls=max_llm_calls, + max_code_failures=max_code_failures, + max_speed_frac=max_speed_frac, + request_timeout_s=request_timeout_s, + gripper_open_is_high=gripper_open_is_high, + transcript_echo=transcript_echo, + ) + self.info = PolicyInfo(name="capx", action_space=Box(shape=(1,))) + self._motion: MotionQueue | None = None + self._sandbox: CodeSandbox | None = None + self._state_key: str | None = None + self._state_labels: tuple[str, tuple[str, ...]] | None = None + self._embodiment_name = "(unbound)" + self._embodiment_docs: str | None = None + self._action_summary = "unbound; call bind() before act()" + self._messages: list[dict[str, Any]] = [] + self._delta_cursor = 0 + self._calls_used = 0 + self._consecutive_failures = 0 + self._closed = False + atexit.register(self.close) + + def bind(self, embodiment_info: EmbodimentInfo) -> None: + """Adopt an embodiment that satisfies the plan-0021 single-arm profile.""" + space = embodiment_info.action_space + prefix = "plan 0021 CaP-X v1 profile requires" + if len(space.shape) != 1 or space.dim < 2: + raise ValueError(f"{prefix} a 1-D Box action space with at least 2 dimensions") + semantics = space.semantics + if semantics is None or semantics.control_mode != "joint_pos": + mode = None if semantics is None else semantics.control_mode + raise ValueError(f"{prefix} control_mode='joint_pos', got {mode!r}") + if semantics.gripper == "none": + raise ValueError(f"{prefix} ActionSemantics.gripper to be continuous or binary") + labels = semantics.dim_labels + if labels is None: + gripper_index = space.dim - 1 + else: + matches = [index for index, label in enumerate(labels) if label == "gripper"] + if len(matches) != 1: + raise ValueError( + f"{prefix} exactly one dim_labels entry named 'gripper'; found {matches}" + ) + gripper_index = matches[0] + if embodiment_info.control_hz is None: + raise ValueError(f"{prefix} a declared finite control_hz > 0") + + state_spec = embodiment_info.observation_space.state + if state_spec is None: + raise ValueError(f"{prefix} an observation StateSpec for full joint state") + matching = [field.key for field in state_spec.fields if field.shape == (space.dim,)] + if len(matching) != 1: + raise ValueError( + f"{prefix} exactly one state field with shape ({space.dim},); " + f"found {matching or 'none'}" + ) + state_key = matching[0] + + camera_names = embodiment_info.observation_space.camera_names + if self._camera_config is None: + if len(camera_names) != 1: + raise ValueError( + f"{prefix} one camera when camera=None; declared {sorted(camera_names)}" + ) + camera = next(iter(camera_names)) + else: + camera = self._camera_config + if camera not in camera_names: + raise ValueError( + f"{prefix} the configured camera to be declared; {camera!r} is absent, " + "available: " + f"{sorted(camera_names)}" + ) + + try: + motion = MotionQueue( + space, + control_hz=embodiment_info.control_hz, + max_speed_frac=self._max_speed_frac, + gripper_index=gripper_index, + gripper_open_is_high=self._gripper_open_is_high, + ) + except ValueError as exc: + raise ValueError(f"{prefix} valid bounded motion arithmetic: {exc}") from exc + self._motion = motion + self._sandbox = CodeSandbox( + servers=self._servers, + motion=motion, + camera=camera, + state_key=state_key, + depth_key=self._depth_key, + intrinsics_key=self._intrinsics_key, + extrinsics_key=self._extrinsics_key, + ) + self._state_key = state_key + self._state_labels = (state_key, labels) if labels is not None else None + self._embodiment_name = embodiment_info.name + self._embodiment_docs = getattr(embodiment_info, "docs", None) + arm_labels = [ + labels[index] if labels is not None else str(index) + for index in range(space.dim) + if index != gripper_index + ] + polarity = ( + "high=open, low=closed" if self._gripper_open_is_high else "low=open, high=closed" + ) + self._action_summary = ( + f"joint_pos with {motion.arm_dof} arm joints {arm_labels}, gripper dim " + f"{gripper_index} ({polarity}), control_hz={embodiment_info.control_hz:g}" + ) + self.info = PolicyInfo( + name="capx", + action_space=space, + observation_space=embodiment_info.observation_space, + control_hz=embodiment_info.control_hz, + ) + + def reset(self, scene: Scene) -> None: + """Start a trial with fresh code, transcript, failure, and IK state.""" + system = _SYSTEM_TEMPLATE.format( + budget=self._max_llm_calls, + name=self._embodiment_name, + action_summary=self._action_summary, + helper_docs=_HELPER_DOCS.format( + depth_key=self._depth_key, + intrinsics_key=self._intrinsics_key, + extrinsics_key=self._extrinsics_key, + ), + ) + if self._embodiment_docs is not None and self._embodiment_docs.strip(): + system += "\n\nEmbodiment notes:\n" + self._embodiment_docs.strip() + self._messages = [ + {"role": "system", "content": system}, + {"role": "user", "content": f"Goal: {scene.instruction}"}, + ] + self._delta_cursor = 0 + self._calls_used = 0 + self._consecutive_failures = 0 + if self._sandbox is not None: + self._sandbox.reset() + self._servers.pyroki.reset() + self._echo(f"[capx] goal: {scene.instruction}") + + def transcript(self) -> list[dict[str, Any]] | None: + """Return an image-free deep copy of the current trial conversation.""" + return _sanitize(self._messages) if self._messages else None + + def transcript_delta(self) -> list[dict[str, Any]] | None: + """Return only sanitized messages appended since the previous live-stream read.""" + new = self._messages[self._delta_cursor :] + self._delta_cursor = len(self._messages) + return _sanitize(new) if new else None + + def act(self, observation: Observation) -> ActionChunk: + """Run codegen turns until code queues motion or the model requests a stop.""" + motion = self._motion + sandbox = self._sandbox + state_key = self._state_key + if motion is None or sandbox is None or state_key is None: + raise RuntimeError( + "CapxPolicy.act() before bind(); run it through eval() or call " + "policy.bind(embodiment.info) first" + ) + sandbox.set_observation(observation) + state = observation.state[state_key] + self._messages.append( + { + "role": "user", + "content": _observation_content(observation, self._state_labels), + } + ) + llm_latency = 0.0 + while True: + if self._calls_used >= self._max_llm_calls: + self._echo("[capx] -- LLM call budget exhausted; forcing GIVE_UP") + return motion.hold_chunk( + state, + stop_reason="GIVE_UP", + inference_latency_s=llm_latency, + ) + started = time.monotonic() + message = self._client.complete( + self._messages, + [], + temperature=self._temperature, + reasoning_effort=self._effort, + ) + llm_latency += time.monotonic() - started + self._calls_used += 1 + self._messages.append(message.raw()) + reply = (message.content or "").strip() + if reply: + self._echo(f"[capx] << {reply}") + control = _CONTROL_WORD.fullmatch(reply) + if control is not None: + return motion.hold_chunk( + state, + stop_reason=control.group(1).upper(), + inference_latency_s=llm_latency, + ) + + code = _extract_code(message) + result = sandbox.execute(code) + self._messages.append({"role": "user", "content": _execution_report(code, result)}) + self._echo_execution(result) + if result.raised: + self._consecutive_failures += 1 + else: + self._consecutive_failures = 0 + + if motion.has_actions(): + return motion.take_chunk(inference_latency_s=llm_latency, code=code) + if result.raised and self._consecutive_failures >= self._max_code_failures: + raise RuntimeError( + "model code failed in " + f"{self._consecutive_failures} consecutive turns; see transcript stderr" + ) + + def close(self) -> None: + """Idempotently close LLM and shared CaP-X HTTP connection pools.""" + if self._closed: + return + self._closed = True + self._client.close() + self._servers.close() + with contextlib.suppress(Exception): # defensive during interpreter teardown + atexit.unregister(self.close) + + def _echo(self, text: str) -> None: + if self._transcript_echo: + print(text, file=sys.stderr, flush=True) + + def _echo_execution(self, result: ExecutionResult) -> None: + if not self._transcript_echo: + return + if result.stdout: + self._echo(f"[capx] stdout:\n{result.stdout.rstrip()}") + if result.stderr: + self._echo(f"[capx] stderr:\n{result.stderr.rstrip()}") + + +def _extract_code(message: AssistantMessage) -> str: + """Normalize raw, fenced, prose-wrapped, or ``REGENERATE``-prefixed code. + + A reply that is exactly one fenced block (or a bare snippet) is used as + is; otherwise the first fenced block wins, so surrounding prose does not + burn a failure turn on a ``SyntaxError``. + """ + code = (message.content or "").strip() + first, separator, remainder = code.partition("\n") + if first.strip() == "REGENERATE" and separator: + code = remainder.strip() + fenced = _FENCED_CODE.fullmatch(code) + if fenced is not None: + return fenced.group(1) + embedded = _FENCED_ANYWHERE.search(code) + return embedded.group(1) if embedded is not None else code + + +def _execution_report(code: str, result: ExecutionResult) -> str: + """Build a bounded tail-first CaP-X execution feedback message.""" + report = ( + "Executed code:\n" + f"```python\n{code}\n```\n" + "stdout:\n" + f"```text\n{result.stdout}\n```\n" + "stderr:\n" + f"```text\n{result.stderr}\n```\n" + "Respond with FINISH, GIVE_UP, or REGENERATE followed by Python." + ) + if len(report) <= _EXECUTION_REPORT_CHAR_LIMIT: + return report + tail_size = _EXECUTION_REPORT_CHAR_LIMIT - len(_REPORT_TRUNCATION_MARKER) + return _REPORT_TRUNCATION_MARKER + report[-tail_size:] + + +def _state_lines( + observation: Observation, + state_labels: tuple[str, tuple[str, ...]] | None, +) -> list[str]: + """Render observation state with action labels when the bound vector aligns.""" + lines: list[str] = [] + for key, value in observation.state.items(): + array = np.asarray(value, dtype=np.float64) + rounded = np.round(array, 4).tolist() + if state_labels is not None and key == state_labels[0]: + labels = state_labels[1] + if array.shape == (len(labels),): + labeled = " ".join( + f"{label}={item}" for label, item in zip(labels, rounded, strict=True) + ) + lines.append(f"state[{key}]: {labeled}") + continue + lines.append(f"state[{key}]: {rounded}") + return lines + + +def _observation_content( + observation: Observation, + state_labels: tuple[str, tuple[str, ...]] | None, +) -> list[dict[str, Any]]: + """Encode fresh state text and camera PNGs without repeating execution feedback.""" + lines = ["Current observation."] + if observation.instruction: + lines.append(f"Instruction: {observation.instruction}") + lines.extend(_state_lines(observation, state_labels)) + parts: list[dict[str, Any]] = [{"type": "text", "text": "\n".join(lines)}] + for name, image in observation.images.items(): + parts.append({"type": "text", "text": f"camera {name!r}:"}) + parts.append({"type": "image_url", "image_url": {"url": png_data_url(image)}}) + return parts + + +def capx_policy(**kwargs: Any) -> CapxPolicy: + """Build the registry policy while forwarding CLI ``-P`` keyword arguments.""" + return CapxPolicy(**kwargs) diff --git a/plugins/inspect-robots-capx/src/inspect_robots_capx/py.typed b/plugins/inspect-robots-capx/src/inspect_robots_capx/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/plugins/inspect-robots-capx/tests/conftest.py b/plugins/inspect-robots-capx/tests/conftest.py new file mode 100644 index 0000000..9001d2e --- /dev/null +++ b/plugins/inspect-robots-capx/tests/conftest.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import base64 +import json +from dataclasses import dataclass, field +from typing import Any + +import httpx +import numpy as np +import pytest + +from inspect_robots_capx._codec import npy_b64_encode + + +@dataclass +class CapxStub: + requests: list[tuple[str, dict[str, Any]]] = field(default_factory=list) + + def handler(self, request: httpx.Request) -> httpx.Response: + body: dict[str, Any] = json.loads(request.content) + self.requests.append((request.url.path, body)) + if request.url.path == "/segment": + mask = np.array([[1, 0], [0, 1]], dtype=np.uint8) + return httpx.Response( + 200, + json={ + "results": [ + { + "mask_base64": base64.b64encode(mask.tobytes()).decode("ascii"), + "shape": [2, 2], + "box": [0, 0, 2, 2], + "score": 0.9, + "label": "cube", + } + ] + }, + ) + if request.url.path == "/plan": + return httpx.Response( + 200, + json={ + "grasps_base64": npy_b64_encode(np.eye(4, dtype=np.float32)[None]), + "scores_base64": npy_b64_encode(np.array([0.8], dtype=np.float32)), + "contact_pts_base64": npy_b64_encode( + np.array([[0.0, 0.0, 0.5]], dtype=np.float32) + ), + }, + ) + if request.url.path == "/ik": + return httpx.Response(200, json={"joint_positions": [0.1, -0.2, 0.7]}) + raise AssertionError(f"unexpected request path {request.url.path}") + + +@pytest.fixture() +def capx_stub() -> CapxStub: + return CapxStub() diff --git a/plugins/inspect-robots-capx/tests/test_codec.py b/plugins/inspect-robots-capx/tests/test_codec.py new file mode 100644 index 0000000..876b348 --- /dev/null +++ b/plugins/inspect-robots-capx/tests/test_codec.py @@ -0,0 +1,83 @@ +"""Golden CaP-X wire payloads keep the dependency-free codecs honest.""" + +from __future__ import annotations + +import base64 + +import numpy as np + +from inspect_robots_capx._codec import ( + grasp_arrays_decode, + mask_decode, + npy_b64_decode, + npy_b64_encode, + png_b64_encode, +) + +_DEPTH_WIRE = ( + "k05VTVBZAQB2AHsnZGVzY3InOiAnPGY0JywgJ2ZvcnRyYW5fb3JkZXInOiBGYWxzZSwgJ3NoYXBlJzog" + "KDIsIDIpLCB9ICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAg" + "ICAgIAoAAIA+AADAPwAAAEAAAAAA" +) +_EMPTY_WIRE = ( + "k05VTVBZAQB2AHsnZGVzY3InOiAnPGY4JywgJ2ZvcnRyYW5fb3JkZXInOiBGYWxzZSwgJ3NoYXBlJzog" + "KDAsKSwgfSAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAg" + "ICAgIAo=" +) +_ONE_GRASP_WIRE = ( + "k05VTVBZAQB2AHsnZGVzY3InOiAnPGY0JywgJ2ZvcnRyYW5fb3JkZXInOiBGYWxzZSwgJ3NoYXBlJzog" + "KDEsIDQsIDQpLCB9ICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAg" + "ICAgIAoAAIA/AAAAAAAAAAAAAAAAAAAAAAAAgD8AAAAAAAAAAAAAAAAAAAAAAACAPwAAAAAAAAAAAAAAAA" + "AAAAAAAIA/" +) +_ONE_SCORE_WIRE = ( + "k05VTVBZAQB2AHsnZGVzY3InOiAnPGY0JywgJ2ZvcnRyYW5fb3JkZXInOiBGYWxzZSwgJ3NoYXBlJzog" + "KDEsKSwgfSAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAg" + "ICAgIAoAAEA/" +) + + +def test_float32_depth_matches_recorded_npy_wire() -> None: + depth = np.array([[0.25, 1.5], [2.0, 0.0]], dtype=np.float32) + + assert npy_b64_encode(depth) == _DEPTH_WIRE + decoded = npy_b64_decode(_DEPTH_WIRE) + assert decoded.dtype == np.float32 + assert np.array_equal(decoded, depth) + + +def test_raw_mask_payload_decodes_to_bool_shape() -> None: + expected = np.array([[True, False, True], [False, True, True]]) + payload = base64.b64encode(expected.astype(np.uint8).tobytes()).decode("ascii") + + decoded = mask_decode(payload, (2, 3)) + + assert payload == "AQABAAEB" + assert decoded.dtype == np.bool_ + assert np.array_equal(decoded, expected) + + +def test_recorded_grasp_payloads_keep_capx_shapes() -> None: + grasps, scores = grasp_arrays_decode(_ONE_GRASP_WIRE, _ONE_SCORE_WIRE) + + assert grasps.shape == (1, 4, 4) + assert scores.shape == (1,) + assert np.array_equal(grasps[0], np.eye(4, dtype=np.float32)) + assert np.array_equal(scores, np.array([0.75], dtype=np.float32)) + + +def test_flat_empty_grasp_payloads_normalize_to_public_shapes() -> None: + raw = npy_b64_decode(_EMPTY_WIRE) + assert raw.shape == (0,) + + grasps, scores = grasp_arrays_decode(_EMPTY_WIRE, _EMPTY_WIRE) + + assert grasps.shape == (0, 4, 4) + assert scores.shape == (0,) + + +def test_png_encoding_is_bare_base64_not_a_data_url() -> None: + encoded = png_b64_encode(np.zeros((1, 2, 3), dtype=np.uint8)) + + assert not encoded.startswith("data:") + assert base64.b64decode(encoded).startswith(b"\x89PNG\r\n\x1a\n") diff --git a/plugins/inspect-robots-capx/tests/test_motion.py b/plugins/inspect-robots-capx/tests/test_motion.py new file mode 100644 index 0000000..0c58396 --- /dev/null +++ b/plugins/inspect-robots-capx/tests/test_motion.py @@ -0,0 +1,135 @@ +"""Motion helpers preserve cursor semantics and the core delta backstop.""" + +from __future__ import annotations + +from itertools import pairwise + +import numpy as np +import pytest + +from inspect_robots.approver import DeltaLimitApprover +from inspect_robots.spaces import ActionSemantics, Box +from inspect_robots.types import Action +from inspect_robots_capx._motion import MotionQueue + + +def _space() -> Box: + return Box( + shape=(3,), + low=np.array([-1.0, -2.0, 0.0]), + high=np.array([1.0, 2.0, 1.0]), + semantics=ActionSemantics( + "joint_pos", + gripper="continuous", + dim_labels=("shoulder", "elbow", "gripper"), + ), + ) + + +def _motion( + *, control_hz: float = 10.0, max_speed_frac: float = 0.1, open_high: bool = True +) -> MotionQueue: + return MotionQueue( + _space(), + control_hz=control_hz, + max_speed_frac=max_speed_frac, + gripper_index=2, + gripper_open_is_high=open_high, + ) + + +def test_speed_fraction_and_control_rate_set_per_step_interpolation() -> None: + motion = _motion(control_hz=10.0, max_speed_frac=0.1) + start = np.array([0.0, 0.0, 1.0]) + motion.begin_turn(start) + + motion.move_to_joints(np.array([0.2, -0.4])) + chunk = motion.take_chunk() + + points = [start, *(action.data for action in chunk.actions)] + deltas = np.abs(np.diff(np.stack(points), axis=0)) + assert len(chunk.actions) == 11 + assert np.all(deltas <= np.array([0.02, 0.04, 0.01])) + assert np.array_equal(chunk.actions[-1].data, np.array([0.2, -0.4, 1.0])) + assert chunk.control_hz == 10.0 + + +def test_low_control_rate_never_exceeds_delta_limit_approver_backstop() -> None: + space = _space() + motion = MotionQueue( + space, + control_hz=1.0, + max_speed_frac=0.5, + gripper_index=2, + ) + start = np.array([-1.0, -2.0, 1.0]) + motion.begin_turn(start) + motion.move_to_joints(np.array([1.0, 2.0])) + chunk = motion.take_chunk() + + approver = DeltaLimitApprover(space) + store: dict[str, object] = {} + approver.review(Action(data=start), store) + for action in chunk.actions: + assert approver.review(action, store) is action + + points = [start, *(action.data for action in chunk.actions)] + per_step = np.abs(np.diff(np.stack(points), axis=0)) + native_backstop = 0.05 * (space.high - space.low) # type: ignore[operator] + assert np.all(per_step <= native_backstop) + + +def test_cursor_chains_across_arm_and_gripper_calls() -> None: + motion = _motion() + motion.begin_turn(np.array([0.0, 0.0, 1.0])) + + motion.move_to_joints(np.array([0.1, 0.2])) + arm_endpoint = motion.cursor + motion.close_gripper() + closed_endpoint = motion.cursor + motion.open_gripper() + open_endpoint = motion.cursor + chunk = motion.take_chunk(code="move_to_joints(...)") + + assert arm_endpoint is not None and np.array_equal(arm_endpoint, [0.1, 0.2, 1.0]) + assert closed_endpoint is not None and np.array_equal(closed_endpoint, [0.1, 0.2, 0.0]) + assert open_endpoint is not None and np.array_equal(open_endpoint, [0.1, 0.2, 1.0]) + assert np.array_equal(chunk.actions[-1].data, open_endpoint) + assert chunk.actions[0].meta["code"] == "move_to_joints(...)" + assert all( + np.max(np.abs(right.data - left.data)) <= 0.04 for left, right in pairwise(chunk.actions) + ) + + +def test_hold_chunk_has_one_full_state_action_and_stop_metadata() -> None: + motion = _motion() + state = np.array([0.3, -0.2, 0.8]) + + chunk = motion.hold_chunk(state, stop_reason="FINISH", inference_latency_s=0.25) + + assert len(chunk.actions) == 1 + assert np.array_equal(chunk.actions[0].data, state) + assert chunk.actions[0].meta == {"request_stop": True, "stop_reason": "FINISH"} + assert chunk.inference_latency_s == 0.25 + + +def test_gripper_values_come_from_box_bounds_with_configurable_polarity() -> None: + normal = _motion(open_high=True) + inverted = _motion(open_high=False) + + assert normal.gripper_open_value == 1.0 + assert normal.gripper_closed_value == 0.0 + assert inverted.gripper_open_value == 0.0 + assert inverted.gripper_closed_value == 1.0 + + +def test_offset_bounds_with_coarse_float_grid_are_rejected() -> None: + space = Box( + shape=(2,), + low=np.array([1e16, 0.0]), + high=np.array([1e16 + 2.0, 1.0]), + semantics=ActionSemantics("joint_pos", gripper="continuous", dim_labels=("j0", "gripper")), + ) + + with pytest.raises(ValueError, match="too coarse"): + MotionQueue(space, control_hz=10.0, max_speed_frac=0.1, gripper_index=1) diff --git a/plugins/inspect-robots-capx/tests/test_policy.py b/plugins/inspect-robots-capx/tests/test_policy.py new file mode 100644 index 0000000..614ffd5 --- /dev/null +++ b/plugins/inspect-robots-capx/tests/test_policy.py @@ -0,0 +1,465 @@ +"""CapxPolicy profile enforcement, codegen protocol, registry, and rollout integration.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import httpx +import numpy as np +import pytest + +from conftest import CapxStub +from inspect_robots import eval as ir_eval +from inspect_robots.embodiment import EmbodimentInfo +from inspect_robots.scene import Scene +from inspect_robots.scorer import success_at_end +from inspect_robots.spaces import ( + ActionSemantics, + Box, + CameraSpec, + ObservationSpace, + StateField, + StateSpec, +) +from inspect_robots.task import Task +from inspect_robots.types import Action, Observation, StepResult +from inspect_robots_capx import CapxPolicy, CapxPolicyConfig, capx_policy +from inspect_robots_capx.policy import ( + _EXECUTION_REPORT_CHAR_LIMIT, + _REPORT_TRUNCATION_MARKER, +) + + +def _completion(content: str) -> dict[str, Any]: + return {"choices": [{"message": {"role": "assistant", "content": content}}]} + + +class _ScriptedTransport: + def __init__(self, responses: list[dict[str, Any]]) -> None: + self.responses = list(responses) + self.llm_requests: list[dict[str, Any]] = [] + self.capx = CapxStub() + + def __call__(self, request: httpx.Request) -> httpx.Response: + if request.url.path in {"/v1/chat/completions", "/v1/responses"}: + self.llm_requests.append(json.loads(request.content)) + if not self.responses: + raise AssertionError("scripted LLM response queue exhausted") + return httpx.Response(200, json=self.responses.pop(0)) + return self.capx.handler(request) + + +def _info( + *, + mode: str = "joint_pos", + gripper: str = "continuous", + labels: tuple[str, ...] | None = ("j0", "j1", "gripper"), + control_hz: float | None = 10.0, + cameras: tuple[CameraSpec, ...] = (CameraSpec("front", 2, 2),), + state: StateSpec | None = StateSpec(fields=(StateField("joint_pos", (3,)),)), +) -> EmbodimentInfo: + return EmbodimentInfo( + name="joint-test", + action_space=Box( + shape=(3,), + low=np.array([-1.0, -1.0, 0.0]), + high=np.array([1.0, 1.0, 1.0]), + semantics=ActionSemantics( + mode, # type: ignore[arg-type] + gripper=gripper, # type: ignore[arg-type] + dim_labels=labels, + ), + ), + observation_space=ObservationSpace(cameras=cameras, state=state), + control_hz=control_hz, + is_simulated=True, + docs="Joint state includes the gripper as its final labeled dimension.", + ) + + +def _observation(*, image: bool = False) -> Observation: + images = {"front": np.zeros((2, 2, 3), dtype=np.uint8)} if image else {} + return Observation( + images=images, + state={"joint_pos": np.array([0.0, 0.0, 1.0])}, + instruction="pick the cube", + ) + + +def _policy(script: _ScriptedTransport, **kwargs: Any) -> CapxPolicy: + return CapxPolicy( + model="test/model", + base_url="http://llm.test/v1", + sam3_url="http://sam.test", + graspnet_url="http://grasp.test", + pyroki_url="http://ik.test", + transport=httpx.MockTransport(script), + env={}, + **kwargs, + ) + + +def _bound_policy(script: _ScriptedTransport, **kwargs: Any) -> CapxPolicy: + policy = _policy(script, **kwargs) + policy.bind(_info()) + policy.reset(Scene(id="s0", instruction="pick the cube")) + return policy + + +@pytest.mark.parametrize( + "info", + [ + _info(mode="eef_abs_pose"), + _info(gripper="none"), + _info(control_hz=None), + _info(labels=("left_j0", "right_j0", "right_gripper")), + ], + ids=["ee_mode", "no_gripper_semantics", "missing_hz", "no_gripper_label"], +) +def test_bind_rejects_out_of_profile_embodiments(info: EmbodimentInfo) -> None: + policy = _policy(_ScriptedTransport([])) + + with pytest.raises(ValueError, match="plan 0021 CaP-X v1 profile"): + policy.bind(info) + + +def test_bind_requires_unambiguous_state_and_camera() -> None: + policy = _policy(_ScriptedTransport([])) + ambiguous_state = StateSpec(fields=(StateField("a", (3,)), StateField("b", (3,)))) + with pytest.raises(ValueError, match="exactly one state field"): + policy.bind(_info(state=ambiguous_state)) + + with pytest.raises(ValueError, match="one camera when camera=None"): + policy.bind(_info(cameras=(CameraSpec("front", 2, 2), CameraSpec("wrist", 2, 2)))) + + +def test_bind_uses_last_dimension_fallback_when_labels_are_absent() -> None: + script = _ScriptedTransport([_completion("close_gripper()")]) + policy = _policy(script) + policy.bind(_info(labels=None)) + policy.reset(Scene(id="s0", instruction="close")) + + chunk = policy.act(_observation()) + + assert np.array_equal(chunk.actions[-1].data, np.array([0.0, 0.0, 0.0])) + + +@pytest.mark.parametrize( + ("reply", "expected_code"), + [ + ( + "import numpy as np\nmove_to_joints(np.array([0.1, -0.1]))", + "import numpy as np\nmove_to_joints(np.array([0.1, -0.1]))", + ), + ( + "```python\nimport numpy as np\nmove_to_joints(np.array([0.1, -0.1]))\n```", + "import numpy as np\nmove_to_joints(np.array([0.1, -0.1]))", + ), + ( + "REGENERATE\n```python\nimport numpy as np\nmove_to_joints(np.array([0.1, -0.1]))\n```", + "import numpy as np\nmove_to_joints(np.array([0.1, -0.1]))", + ), + ], + ids=["raw", "fenced", "regenerate_fenced"], +) +def test_raw_and_fenced_code_paths(reply: str, expected_code: str) -> None: + policy = _bound_policy(_ScriptedTransport([_completion(reply)])) + + chunk = policy.act(_observation()) + + assert chunk.actions[0].meta["code"] == expected_code + assert np.array_equal(chunk.actions[-1].data, np.array([0.1, -0.1, 1.0])) + + +@pytest.mark.parametrize("word", ["FINISH", "GIVE_UP"]) +def test_control_words_return_one_action_stop_hold(word: str) -> None: + policy = _bound_policy(_ScriptedTransport([_completion(word)])) + + chunk = policy.act(_observation()) + + assert len(chunk.actions) == 1 + assert np.array_equal(chunk.actions[0].data, np.array([0.0, 0.0, 1.0])) + assert chunk.actions[0].meta == {"request_stop": True, "stop_reason": word} + + +@pytest.mark.parametrize( + ("reply", "reason"), + [("FINISH.", "FINISH"), ("GIVE_UP!", "GIVE_UP"), ("finish", "FINISH")], +) +def test_control_words_tolerate_punctuation_and_case(reply: str, reason: str) -> None: + policy = _bound_policy(_ScriptedTransport([_completion(reply)])) + + chunk = policy.act(_observation()) + + assert chunk.actions[0].meta == {"request_stop": True, "stop_reason": reason} + + +def test_prose_wrapped_fenced_code_executes_the_fenced_block() -> None: + script = _ScriptedTransport( + [ + _completion( + "Sure, here is the corrected code:\n" + "```python\nimport numpy as np\nmove_to_joints(np.array([0.1, -0.1]))\n```\n" + "Let me know how it goes." + ), + ] + ) + policy = _bound_policy(script) + + chunk = policy.act(_observation()) + + assert np.array_equal(chunk.actions[-1].data, np.array([0.1, -0.1, 1.0])) + + +def test_perception_only_turn_feeds_report_back_then_returns_motion() -> None: + script = _ScriptedTransport( + [ + _completion("print('need another look')"), + _completion("import numpy as np\nmove_to_joints(np.array([0.2, 0.0]))"), + ] + ) + policy = _bound_policy(script) + + chunk = policy.act(_observation()) + + assert len(script.llm_requests) == 2 + second_messages = script.llm_requests[1]["messages"] + assert "need another look" in second_messages[-1]["content"] + assert np.array_equal(chunk.actions[-1].data, np.array([0.2, 0.0, 1.0])) + + +def test_failure_counter_persists_across_act_even_when_error_queued_actions() -> None: + script = _ScriptedTransport( + [ + _completion( + "import numpy as np\nmove_to_joints(np.array([0.1, 0.0]))\n" + "raise RuntimeError('after queue')" + ), + _completion("raise ValueError('next turn')"), + ] + ) + policy = _bound_policy(script, max_code_failures=2) + + first = policy.act(_observation()) + assert first.actions + + with pytest.raises(RuntimeError, match="2 consecutive turns"): + policy.act(_observation()) + transcript = policy.transcript() + assert transcript is not None and "next turn" in transcript[-1]["content"] + + +def test_clean_perception_turn_resets_consecutive_failure_counter() -> None: + script = _ScriptedTransport( + [ + _completion( + "import numpy as np\nmove_to_joints(np.array([0.1, 0.0]))\nraise ValueError('one')" + ), + _completion("print('clean turn')"), + _completion("raise ValueError('one again')"), + _completion("FINISH"), + ] + ) + policy = _bound_policy(script, max_code_failures=2) + + policy.act(_observation()) + stopped = policy.act(_observation()) + + assert stopped.actions[0].meta["stop_reason"] == "FINISH" + + +def test_call_budget_forces_give_up_after_clean_empty_turn() -> None: + policy = _bound_policy( + _ScriptedTransport([_completion("print('nothing queued')")]), + max_llm_calls=1, + ) + + chunk = policy.act(_observation()) + + assert chunk.actions[0].meta == {"request_stop": True, "stop_reason": "GIVE_UP"} + + +def test_execution_report_is_eager_and_code_lands_on_first_action() -> None: + code = "import numpy as np\nprint('terminal feedback')\nmove_to_joints(np.array([0.1, 0.0]))" + policy = _bound_policy(_ScriptedTransport([_completion(code)])) + + chunk = policy.act(_observation()) + transcript = policy.transcript() + + assert transcript is not None + assert "terminal feedback" in transcript[-1]["content"] + assert transcript[-1]["role"] == "user" + assert chunk.actions[0].meta["code"] == code + + +def test_execution_report_is_tail_first_truncated_to_documented_cap() -> None: + script = _ScriptedTransport( + [ + _completion("print('x' * 20000)"), + _completion("import numpy as np\nmove_to_joints(np.array([0.1, 0.0]))"), + ] + ) + policy = _bound_policy(script) + + policy.act(_observation()) + transcript = policy.transcript() + assert transcript is not None + reports = [ + message["content"] + for message in transcript + if isinstance(message.get("content"), str) + and message["content"].startswith(_REPORT_TRUNCATION_MARKER) + ] + assert len(reports) == 1 + assert len(reports[0]) == _EXECUTION_REPORT_CHAR_LIMIT + assert reports[0].endswith("Respond with FINISH, GIVE_UP, or REGENERATE followed by Python.") + + +def test_transcript_and_delta_are_deep_sanitized_and_reset_rewinds_cursor() -> None: + policy = _bound_policy(_ScriptedTransport([_completion("FINISH")])) + + policy.act(_observation(image=True)) + first = policy.transcript_delta() + assert first is not None + assert "data:image" not in json.dumps(first) + assert policy.transcript_delta() is None + + first[0]["content"] = "mutated" + full = policy.transcript() + assert full is not None and full[0]["content"] != "mutated" + + policy.reset(Scene(id="s1", instruction="again")) + reset_delta = policy.transcript_delta() + assert reset_delta is not None + assert [message["role"] for message in reset_delta] == ["system", "user"] + + +def test_reset_clears_pyroki_full_config_warm_start() -> None: + code = ( + "import numpy as np\n" + "q = solve_ik(np.array([0.4, 0.0, 0.2]), np.array([1.0, 0.0, 0.0, 0.0]))\n" + "move_to_joints(q)" + ) + script = _ScriptedTransport([_completion(code), _completion(code)]) + policy = _bound_policy(script) + + policy.act(_observation()) + policy.reset(Scene(id="s1", instruction="repeat")) + policy.act(_observation()) + + ik_bodies = [body for path, body in script.capx.requests if path == "/ik"] + assert [body["prev_cfg"] for body in ik_bodies] == [None, None] + + +def test_config_and_close_lifecycle() -> None: + policy = _policy( + _ScriptedTransport([]), + max_llm_calls=7, + max_code_failures=4, + gripper_open_is_high=False, + transcript_echo=True, + ) + + assert isinstance(policy.config, CapxPolicyConfig) + assert policy.config.max_llm_calls == 7 + assert policy.config.max_code_failures == 4 + assert policy.config.gripper_open_is_high is False + policy.close() + policy.close() + assert policy._servers._http.is_closed + + +def test_unbound_act_is_actionable() -> None: + policy = _policy(_ScriptedTransport([])) + with pytest.raises(RuntimeError, match="before bind"): + policy.act(_observation()) + + +class _JointEmbodiment: + def __init__(self) -> None: + self._q = np.array([0.0, 0.0, 1.0]) + self.info = _info(control_hz=2.0) + + def _observation(self, instruction: str | None = None) -> Observation: + return Observation( + images={"front": np.zeros((2, 2, 3), dtype=np.uint8)}, + state={"joint_pos": self._q.copy()}, + instruction=instruction, + extra={ + "depth": lambda: np.ones((2, 2), dtype=np.float32), + "intrinsics": np.eye(3), + "extrinsics": np.eye(4), + }, + ) + + def reset(self, scene: Scene, *, seed: int | None = None) -> Observation: + self._q = np.array([0.0, 0.0, 1.0]) + return self._observation(scene.instruction) + + def step(self, action: Action) -> StepResult: + self._q = np.asarray(action.data, dtype=np.float64).copy() + return StepResult(observation=self._observation()) + + def close(self) -> None: + return None + + +def test_end_to_end_joint_embodiment_rollout_carries_transcript(tmp_path: Path) -> None: + code = ( + "import numpy as np\n" + "objects = segment('red cube')\n" + "q = solve_ik(np.array([0.4, 0.0, 0.2]), np.array([1.0, 0.0, 0.0, 0.0]))\n" + "move_to_joints(q)\n" + "close_gripper()" + ) + script = _ScriptedTransport([_completion(code), _completion("FINISH")]) + policy = _policy(script) + task = Task( + name="capx-e2e", + scenes=[Scene(id="s0", instruction="pick the red cube")], + scorer=success_at_end(), + max_steps=40, + ) + + logs = ir_eval(task, policy, _JointEmbodiment(), log_dir=str(tmp_path)) + + assert logs[0].status == "success" + sample = logs[0].samples[0] + assert sample.status == "success" + (transcript,) = sample.policy_transcripts + assert transcript is not None + serialized = json.dumps(transcript) + assert "segment('red cube')" in serialized + assert "FINISH" in serialized + assert "data:image" not in serialized + policy.close() + + +def test_registry_entry_point_resolves_and_factory_forwards_kwargs() -> None: + from inspect_robots.registry import resolve + + direct = capx_policy( + model="test/model", + base_url="http://llm.test/v1", + max_llm_calls=7, + env={}, + ) + resolved = resolve( + "policy", + "capx", + model="test/model", + base_url="http://llm.test/v1", + max_llm_calls=9, + env={}, + ) + + assert isinstance(direct, CapxPolicy) + assert isinstance(direct.config, CapxPolicyConfig) + assert direct.config.max_llm_calls == 7 + assert isinstance(resolved, CapxPolicy) + assert isinstance(resolved.config, CapxPolicyConfig) + assert resolved.config.max_llm_calls == 9 + direct.close() + resolved.close() diff --git a/plugins/inspect-robots-capx/tests/test_sandbox.py b/plugins/inspect-robots-capx/tests/test_sandbox.py new file mode 100644 index 0000000..173e428 --- /dev/null +++ b/plugins/inspect-robots-capx/tests/test_sandbox.py @@ -0,0 +1,173 @@ +"""The execution namespace persists per trial and turns helper failures into feedback.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import httpx +import numpy as np +import pytest + +from conftest import CapxStub +from inspect_robots.spaces import ActionSemantics, Box +from inspect_robots.types import Observation +from inspect_robots_capx._codec import npy_b64_decode +from inspect_robots_capx._motion import MotionQueue +from inspect_robots_capx._sandbox import CodeSandbox +from inspect_robots_capx._servers import CapxServerClients + + +def _sandbox(capx_stub: CapxStub) -> tuple[CodeSandbox, MotionQueue, CapxServerClients]: + space = Box( + shape=(3,), + low=np.array([-1.0, -1.0, 0.0]), + high=np.array([1.0, 1.0, 1.0]), + semantics=ActionSemantics( + "joint_pos", + gripper="continuous", + dim_labels=("j0", "j1", "gripper"), + ), + ) + motion = MotionQueue(space, control_hz=10.0, max_speed_frac=0.1, gripper_index=2) + servers = CapxServerClients( + sam3_url="http://sam.test", + graspnet_url="http://grasp.test", + pyroki_url="http://ik.test", + transport=httpx.MockTransport(capx_stub.handler), + ) + sandbox = CodeSandbox( + servers=servers, + motion=motion, + camera="front", + state_key="joint_pos", + ) + return sandbox, motion, servers + + +def _observation(**extra: Any) -> Observation: + return Observation( + images={"front": np.zeros((2, 2, 3), dtype=np.uint8)}, + state={"joint_pos": np.array([0.0, 0.0, 1.0])}, + extra=extra, + ) + + +def test_namespace_persists_within_trial_and_resets_across_trials(capx_stub: CapxStub) -> None: + sandbox, _, _ = _sandbox(capx_stub) + sandbox.set_observation(_observation()) + + first = sandbox.execute("counter = 1") + second = sandbox.execute("counter += 1\nprint(counter)") + sandbox.reset() + sandbox.set_observation(_observation()) + after_reset = sandbox.execute("print(counter)") + + assert first.raised is False + assert second.stdout == "2\n" + assert second.stderr == "" + assert after_reset.raised is True + assert "NameError" in after_reset.stderr + + +def test_stdout_stderr_and_traceback_are_captured(capx_stub: CapxStub) -> None: + sandbox, _, _ = _sandbox(capx_stub) + sandbox.set_observation(_observation()) + + result = sandbox.execute( + "import sys\nprint('out')\nprint('warning', file=sys.stderr)\nraise ValueError('boom')" + ) + + assert result.stdout == "out\n" + assert "warning\n" in result.stderr + assert "Traceback (most recent call last)" in result.stderr + assert "ValueError: boom" in result.stderr + assert result.raised is True + + +def test_helper_error_becomes_stderr_instead_of_crashing(capx_stub: CapxStub) -> None: + sandbox, motion, _ = _sandbox(capx_stub) + sandbox.set_observation(_observation()) + + result = sandbox.execute("plan_grasp(segment('cube')[0]['mask'])") + + assert result.raised is True + assert "observation.extra is missing 'depth'" in result.stderr + assert "zero-argument callable" in result.stderr + assert motion.has_actions() is False + + +@pytest.mark.parametrize( + "depth_factory", + [ + lambda depth: depth, + lambda depth: lambda: depth, + ], + ids=["raw_array", "zero_arg_callable"], +) +def test_depth_accepts_raw_array_and_zero_arg_callable( + capx_stub: CapxStub, + depth_factory: Callable[[np.ndarray], Any], +) -> None: + sandbox, _, _ = _sandbox(capx_stub) + depth = np.array([[0.5, 0.6], [0.7, 0.8]], dtype=np.float32) + sandbox.set_observation( + _observation(depth=depth_factory(depth), intrinsics=np.eye(3), extrinsics=np.eye(4)) + ) + + result = sandbox.execute("mask = segment('cube')[0]['mask']\nposes, scores = plan_grasp(mask)") + + assert result.raised is False + plan_body = next(body for path, body in capx_stub.requests if path == "/plan") + assert np.array_equal(npy_b64_decode(plan_body["depth_base64"]), depth) + + +def test_motion_helpers_share_one_cursor_and_queue(capx_stub: CapxStub) -> None: + sandbox, motion, _ = _sandbox(capx_stub) + sandbox.set_observation(_observation()) + + result = sandbox.execute("move_to_joints(np.array([0.1, -0.1]))") + assert result.raised is True + assert "NameError: name 'np' is not defined" in result.stderr + + recovered = sandbox.execute( + "import numpy as np\nmove_to_joints(np.array([0.1, -0.1]))\nclose_gripper()" + ) + chunk = motion.take_chunk() + + assert recovered.raised is False + assert np.array_equal(chunk.actions[-1].data, np.array([0.1, -0.1, 0.0])) + + sandbox.execute("move_to_joints(np.array([0.2, 0.2]))") + assert motion.has_actions() is True + sandbox.reset() + assert motion.has_actions() is False + assert motion.cursor is None + + +def test_obs_access_resolves_callable_values(capx_stub: CapxStub) -> None: + sandbox, _, _ = _sandbox(capx_stub) + sandbox.set_observation(_observation(extrinsics=lambda: np.eye(4))) + + result = sandbox.execute( + "import numpy as np\n" + "pose = obs['extrinsics'] @ np.eye(4)\n" + "assert pose.shape == (4, 4)\n" + "assert obs.get('extrinsics').shape == (4, 4)\n" + "assert obs.get('absent') is None" + ) + + assert result.raised is False, result.stderr + + +def test_obs_get_propagates_errors_raised_inside_thunks(capx_stub: CapxStub) -> None: + def broken() -> Any: + raise KeyError("embodiment thunk bug") + + sandbox, _, _ = _sandbox(capx_stub) + sandbox.set_observation(_observation(depth=broken)) + + result = sandbox.execute("value = obs.get('depth')") + + assert result.raised is True + assert "embodiment thunk bug" in result.stderr diff --git a/plugins/inspect-robots-capx/tests/test_servers.py b/plugins/inspect-robots-capx/tests/test_servers.py new file mode 100644 index 0000000..d310566 --- /dev/null +++ b/plugins/inspect-robots-capx/tests/test_servers.py @@ -0,0 +1,222 @@ +"""CaP-X clients speak the recorded schemas with bounded retry behavior.""" + +from __future__ import annotations + +import json +from typing import Any + +import httpx +import numpy as np +import pytest + +from conftest import CapxStub +from inspect_robots_capx._codec import npy_b64_decode +from inspect_robots_capx._servers import ( + CapxServerClients, + GraspNetClient, + PyrokiClient, + Sam3Client, +) + + +class _FakeTime: + def __init__(self) -> None: + self.now = 10.0 + + def monotonic(self) -> float: + return self.now + + def sleep(self, seconds: float) -> None: + self.now += seconds + + +def test_clients_are_lazy_and_share_one_injected_http_pool(capx_stub: CapxStub) -> None: + servers = CapxServerClients( + sam3_url="http://sam.test", + graspnet_url="http://grasp.test", + pyroki_url="http://ik.test", + transport=httpx.MockTransport(capx_stub.handler), + ) + + assert capx_stub.requests == [] + assert servers.sam3._http is servers.graspnet._http is servers.pyroki._http + + servers.close() + assert servers._http.is_closed + + +def test_sam3_request_and_response_match_wire_schema(capx_stub: CapxStub) -> None: + client = Sam3Client("http://sam.test", transport=httpx.MockTransport(capx_stub.handler)) + + results = client.segment(np.zeros((2, 2, 3), dtype=np.uint8), "red cube") + + path, body = capx_stub.requests[0] + assert path == "/segment" + assert set(body) == {"image_base64", "text_prompt"} + assert body["text_prompt"] == "red cube" + assert results[0]["mask"].dtype == np.bool_ + assert np.array_equal(results[0]["mask"], np.eye(2, dtype=np.bool_)) + assert results[0]["box"] == [0, 0, 2, 2] + client.close() + + +def test_graspnet_request_pins_defaults_and_npy_arrays(capx_stub: CapxStub) -> None: + client = GraspNetClient("http://grasp.test", transport=httpx.MockTransport(capx_stub.handler)) + depth = np.array([[0.5, 0.6], [0.7, 0.8]], dtype=np.float32) + intrinsics = np.eye(3, dtype=np.float64) + mask = np.array([[True, False], [False, True]]) + + grasps, scores = client.plan(depth, intrinsics, mask) + + path, body = capx_stub.requests[0] + assert path == "/plan" + assert set(body) == { + "depth_base64", + "cam_K_base64", + "segmap_base64", + "segmap_id", + "local_regions", + "filter_grasps", + "skip_border_objects", + "z_range", + "forward_passes", + "max_retries", + } + assert np.array_equal(npy_b64_decode(body["depth_base64"]), depth) + assert np.array_equal(npy_b64_decode(body["cam_K_base64"]), intrinsics) + assert np.array_equal(npy_b64_decode(body["segmap_base64"]), mask.astype(np.uint8)) + tuning = { + key: value + for key, value in body.items() + if key not in {"depth_base64", "cam_K_base64", "segmap_base64"} + } + assert tuning == { + "segmap_id": 1, + "local_regions": True, + "filter_grasps": True, + "skip_border_objects": False, + "z_range": [0.2, 2.0], + "forward_passes": 2, + "max_retries": 10, + } + assert grasps.shape == (1, 4, 4) + assert np.array_equal(scores, np.array([0.8], dtype=np.float32)) + + +def test_transient_503_retries_then_succeeds() -> None: + calls = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + if calls == 1: + return httpx.Response(503, text="warming") + return httpx.Response(200, json={"results": []}) + + client = Sam3Client( + "http://sam.test", + transport=httpx.MockTransport(handler), + backoff_s=0.0, + ) + + assert client.segment(np.zeros((1, 1, 3), dtype=np.uint8), "cube") == [] + assert calls == 2 + + +def test_connection_error_names_url_and_launch_command() -> None: + fake_time = _FakeTime() + + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("offline", request=request) + + client = Sam3Client( + "http://gpu-box:8114", + request_timeout_s=0.25, + transport=httpx.MockTransport(handler), + clock=fake_time.monotonic, + sleep=fake_time.sleep, + ) + + with pytest.raises(ConnectionError) as excinfo: + client.segment(np.zeros((1, 1, 3), dtype=np.uint8), "cube") + message = str(excinfo.value) + assert "http://gpu-box:8114/segment" in message + assert "launch_sam3_server.py --port 8114" in message + + +def test_total_timeout_budget_caps_each_attempt() -> None: + fake_time = _FakeTime() + attempt_timeouts: list[float] = [] + + def handler(request: httpx.Request) -> httpx.Response: + timeout = request.extensions["timeout"] + attempt_timeouts.append(float(timeout["read"])) + raise httpx.ReadTimeout("still loading", request=request) + + client = Sam3Client( + "http://sam.test", + request_timeout_s=1.5, + transport=httpx.MockTransport(handler), + clock=fake_time.monotonic, + sleep=fake_time.sleep, + ) + + with pytest.raises(TimeoutError, match="still loading"): + client.segment(np.zeros((1, 1, 3), dtype=np.uint8), "cube") + assert attempt_timeouts == [1.5, 0.5] + assert fake_time.now == pytest.approx(11.5) + + +def test_attempt_timeout_is_never_above_thirty_seconds() -> None: + seen: list[dict[str, Any]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request.extensions["timeout"]) + return httpx.Response(200, json={"results": []}) + + client = Sam3Client( + "http://sam.test", + request_timeout_s=120.0, + transport=httpx.MockTransport(handler), + ) + client.segment(np.zeros((1, 1, 3), dtype=np.uint8), "cube") + + assert seen[0]["read"] == 30.0 + + +def test_pyroki_strips_arm_and_reuses_full_config_as_warm_start() -> None: + bodies: list[dict[str, Any]] = [] + responses = [[0.1, -0.2, 0.7], [0.2, -0.1, 0.6]] + + def handler(request: httpx.Request) -> httpx.Response: + bodies.append(json.loads(request.content)) + return httpx.Response(200, json={"joint_positions": responses.pop(0)}) + + client = PyrokiClient("http://ik.test", transport=httpx.MockTransport(handler)) + position = np.array([0.4, 0.1, 0.2]) + quaternion = np.array([1.0, 0.0, 0.0, 0.0]) + + first = client.solve_ik(position, quaternion, arm_dof=2) + second = client.solve_ik(position, quaternion, arm_dof=2) + + assert bodies[0] == { + "target_pose_wxyz_xyz": [1.0, 0.0, 0.0, 0.0, 0.4, 0.1, 0.2], + "prev_cfg": None, + } + assert bodies[1]["prev_cfg"] == [0.1, -0.2, 0.7] + assert np.array_equal(first, np.array([0.1, -0.2])) + assert np.array_equal(second, np.array([0.2, -0.1])) + + client.reset() + assert client._prev_cfg is None + + +def test_pyroki_short_response_is_actionable_and_not_warm_started() -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"joint_positions": [0.1]}) + + client = PyrokiClient("http://ik.test", transport=httpx.MockTransport(handler)) + + with pytest.raises(ValueError, match="URDF matching this embodiment"): + client.solve_ik(np.zeros(3), np.array([1.0, 0.0, 0.0, 0.0]), arm_dof=2) + assert client._prev_cfg is None diff --git a/uv.lock b/uv.lock index 2b7e389..685f5af 100644 --- a/uv.lock +++ b/uv.lock @@ -11,6 +11,7 @@ resolution-markers = [ members = [ "inspect-robots", "inspect-robots-agent", + "inspect-robots-capx", "inspect-robots-isaacsim", "inspect-robots-ros", "inspect-robots-xpolicylab", @@ -358,7 +359,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -520,7 +521,7 @@ provides-extras = ["all", "dev", "docs", "rerun", "viz"] [[package]] name = "inspect-robots-agent" -version = "0.11.0" +version = "0.12.0" source = { editable = "plugins/inspect-robots-agent" } dependencies = [ { name = "httpx" }, @@ -549,6 +550,39 @@ requires-dist = [ ] provides-extras = ["dev"] +[[package]] +name = "inspect-robots-capx" +version = "0.1.0" +source = { editable = "plugins/inspect-robots-capx" } +dependencies = [ + { name = "httpx" }, + { name = "inspect-robots" }, + { name = "inspect-robots-agent" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] + +[package.optional-dependencies] +dev = [ + { name = "mypy" }, + { name = "pytest" }, + { name = "pytest-cov" }, + { name = "ruff" }, +] + +[package.metadata] +requires-dist = [ + { name = "httpx", specifier = ">=0.27" }, + { name = "inspect-robots", editable = "." }, + { name = "inspect-robots-agent", editable = "plugins/inspect-robots-agent" }, + { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.11" }, + { name = "numpy", specifier = ">=1.24" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, + { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=5.0" }, + { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.6" }, +] +provides-extras = ["dev"] + [[package]] name = "inspect-robots-isaacsim" version = "0.1.1"