diff --git a/CONTEXT.md b/CONTEXT.md index d97571b990..37617d1187 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -1,3 +1,57 @@ +# Grasp Planning + +This context defines the language at the boundary between upstream perception, grasp generation, and robot motion planning. + +## Language + +**Segmented Object Cloud**: +A target-only 3D point cloud supplied by upstream perception in the manipulation planning frame. +_Avoid_: Scene cloud, detection cloud, raw camera cloud + +**Feasible Grasp Sequence**: +A connected, collision-free trajectory from the robot's current state through any required safety lift, pre-grasp, grasp, and retreat. Each segment begins at the preceding segment's endpoint. Validation is a no-motion dry run; execution replans each segment from fresh measured state. +_Avoid_: Reachable grasp, feasible pose, independent IK success + +**Safety Lift**: +An optional shared preparation segment planned before evaluating grasp candidates. If required and unplannable, the pick aborts once in `PREPARE`; the failure is not attributed to every candidate. +_Avoid_: Pre-grasp failure, candidate rejection + +**Retreat Feasibility (MVP)**: +A connected grasp-to-retreat plan that collision-checks the robot and gripper against the non-target scene. The selected target remains excluded; attached-object geometry and held-object clearance are not modeled. +_Avoid_: Payload-safe retreat, attached-object validation + +**Live-Scene Validation (MVP)**: +Each segment of a feasible grasp sequence is checked against the latest available planning scene. The MVP does not snapshot the scene or freeze non-target obstacle updates across the sequence. Execution replans against freshly measured state and scene data. +_Avoid_: Atomic scene validation, frozen-scene guarantee + +**Gripper Geometry During Validation (MVP)**: +Arm-path collision checks use the gripper configuration currently represented in the planning scene. Dry-run validation does not model the open-to-closed gripper transition or claim separate clearance guarantees for each finger configuration. +_Avoid_: Coordinated arm-gripper plan, validated finger sweep + +**Candidate Rejection Reason (MVP)**: +Candidate rejection is reported by failed sequence stage: `pre_grasp_infeasible`, `grasp_infeasible`, or `retreat_infeasible`. Detailed IK and planner outcomes remain diagnostic logs rather than public skill-result categories. +_Avoid_: Backend-specific public failure codes + +**Pipeline Demo**: +A no-hardware contributor command that runs a recorded Segmented Object Cloud through real grasp proposal and connected motion validation, then saves candidate outcomes and all planned segments. It stops before trajectory execution. +_Avoid_: Grasp-only demo, hardware pick demo + +**Visualization Layer**: +A display-only, named collection of visual elements owned by exactly one producer. Publishing replaces its contents, while clearing leaves the layer registered and preserves viewer-owned visibility; the layer cannot affect collision checking or other planning behavior. +_Avoid_: Collision layer, shared scene state, visualization object + +**Visual Element**: +A backend-neutral drawable contained in a Visualization Layer, initially a point cloud or line set. It carries no collision or planning authority. +_Avoid_: Grasp visualization command, Viser handle, collision object + +**Visualization Layer Group**: +A viewer-only grouping of independently replaceable and toggleable Visualization Layers that share a name prefix, such as `grasp/object-cloud` and `grasp/proposals`. +_Avoid_: Compound layer, element-level visibility + +**Accepted Collision Projection**: +A display-only representation published after the planning world accepts a collision-object change. Its presence, absence, or rendering failure never changes collision checking. +_Avoid_: Collision authority, visualization obstacle + # Manipulation Planning This context describes requests for planning robot motion through joint and Cartesian spaces. diff --git a/bin/setup-graspgenx-env b/bin/setup-graspgenx-env new file mode 100644 index 0000000000..712b341ec9 --- /dev/null +++ b/bin/setup-graspgenx-env @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +set -euo pipefail + +PYTHON=.venv/bin/python +uv pip install --python "$PYTHON" --extra-index-url https://download.pytorch.org/whl/cu128 \ + --index-strategy unsafe-best-match \ + torch==2.7.1+cu128 torchvision==0.22.1+cu128 +uv pip install --python "$PYTHON" \ + torch-geometric h5py hydra-core matplotlib numpy==1.26.4 webdataset scikit-learn scipy \ + tensorboard transformers tensordict diffusers==0.11.1 timm==1.0.15 \ + huggingface-hub==0.25.2 PyOpenGL==3.1.5 addict yapf==0.40.1 tensorboardx \ + sharedarray yourdfpy==0.0.56 urdfpy imageio viser tqdm pyyaml edgetam-dimos "networkx>=3.3" +uv pip install --python "$PYTHON" --no-deps "git+https://github.com/NVlabs/GraspGenX.git" diff --git a/data/.lfs/graspgenx_ycb_banana_scene.tar.gz b/data/.lfs/graspgenx_ycb_banana_scene.tar.gz new file mode 100644 index 0000000000..88ee847b1d --- /dev/null +++ b/data/.lfs/graspgenx_ycb_banana_scene.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:54a41be0741c10959e17d8aa79db9bec9a755da6385dd33c31e78189081fff8d +size 38410 diff --git a/dimos/agents/capabilities.py b/dimos/agents/capabilities.py index 0f64d0af65..be505e6eb0 100644 --- a/dimos/agents/capabilities.py +++ b/dimos/agents/capabilities.py @@ -35,6 +35,7 @@ from typing import NamedTuple CAP_MOVEMENT = "movement" +CAP_PERCEPTION = "perception" class _Hold(NamedTuple): diff --git a/dimos/agents/mcp/mcp_client.py b/dimos/agents/mcp/mcp_client.py index 859b15451b..4031db8b54 100644 --- a/dimos/agents/mcp/mcp_client.py +++ b/dimos/agents/mcp/mcp_client.py @@ -21,7 +21,7 @@ from langchain.agents import create_agent from langchain.chat_models import init_chat_model -from langchain_core.messages import HumanMessage +from langchain_core.messages import HumanMessage, ToolMessage from langchain_core.messages.base import BaseMessage from langchain_core.tools import StructuredTool from langchain_openai import ChatOpenAI @@ -69,6 +69,7 @@ class McpClient(Module): agent: Out[BaseMessage] human_input: In[str] agent_idle: Out[bool] + agent_cancel: In[bool] _lock: RLock _state_graph: CompiledStateGraph[Any, Any, Any, Any] | None @@ -77,6 +78,7 @@ class McpClient(Module): _history: list[BaseMessage] _thread: Thread _stop_event: Event + _cancel_event: Event _http_client: requests.Session _seq_ids: SequentialIds _tool_stream_cleanup: Callable[[], None] | None @@ -94,6 +96,7 @@ def __init__(self, **kwargs: Any) -> None: daemon=True, ) self._stop_event = Event() + self._cancel_event = Event() self._http_client = requests.Session() self._seq_ids = SequentialIds() self._tool_stream_cleanup = None @@ -217,6 +220,12 @@ def _on_human_input(string: str) -> None: self.register_disposable(Disposable(self.human_input.subscribe(_on_human_input))) + def _on_agent_cancel(_cancel: bool) -> None: + self._cancel_event.set() + self.agent_idle.publish(True) + + self.register_disposable(Disposable(self.agent_cancel.subscribe(_on_agent_cancel))) + # Subscribe directly over LCM rather than through the server's GET # /mcp SSE channel. HTTP would add a startup race: the first few # updates of a short-lived stream can fire before the SSE connection @@ -330,7 +339,13 @@ def _thread_loop(self) -> None: with self._lock: if not self._state_graph: raise ValueError("No state graph initialized") - self._process_message(self._state_graph, message) + self._cancel_event.clear() + try: + self._process_message(self._state_graph, message) + except Exception: + self._close_cancelled_tool_calls() + logger.exception("Agent turn failed") + self.agent_idle.publish(True) def _process_message( self, state_graph: CompiledStateGraph[Any, Any, Any, Any], message: BaseMessage @@ -341,15 +356,39 @@ def _process_message( self.agent.publish(message) for update in state_graph.stream({"messages": self._history}, stream_mode="updates"): + if self._cancel_event.is_set(): + self._close_cancelled_tool_calls() + break for node_output in update.values(): for msg in node_output.get("messages", []): self._history.append(msg) pretty_print_langchain_message(msg) self.agent.publish(msg) - if self._message_queue.empty(): + if self._cancel_event.is_set(): + self._close_cancelled_tool_calls() + break + + if self._cancel_event.is_set() or self._message_queue.empty(): self.agent_idle.publish(True) + def _close_cancelled_tool_calls(self) -> None: + """Give every retained tool call a result before sending history to the model again.""" + pending: dict[str, str] = {} + for message in self._history: + for tool_call in getattr(message, "tool_calls", []): + if call_id := tool_call.get("id"): + pending[call_id] = tool_call.get("name", "tool") + if isinstance(message, ToolMessage): + pending.pop(message.tool_call_id, None) + for call_id, name in pending.items(): + self._history.append( + ToolMessage( + content=f"{name} was cancelled before a result was available.", + tool_call_id=call_id, + ) + ) + def _append_image_to_history( mcp_client: McpClient, func_name: str, uuid_: str, result: Any diff --git a/dimos/agents/mcp/mcp_server.py b/dimos/agents/mcp/mcp_server.py index 61d7572af8..113b22634f 100644 --- a/dimos/agents/mcp/mcp_server.py +++ b/dimos/agents/mcp/mcp_server.py @@ -33,7 +33,7 @@ from dimos.agents.capabilities import CapabilityRegistry from dimos.agents.mcp import tool_stream from dimos.core.core import rpc -from dimos.core.module import Module +from dimos.core.module import Module, ModuleConfig from dimos.core.rpc_client import RpcCall, RPCClient from dimos.core.transport_factory import make_transport from dimos.utils.logging_config import setup_logger @@ -52,6 +52,14 @@ # `_can_wait` in `_handle_tools_call`). DEFAULT_CAP_ACQUIRE_TIMEOUT = 30.0 # seconds + +class McpServerConfig(ModuleConfig): + """Configuration for the MCP HTTP server.""" + + allowed_skills: list[str] | None = None + """Optional names of skills exposed through MCP; None exposes every deployed skill.""" + + app = FastAPI() app.add_middleware( CORSMiddleware, @@ -80,6 +88,27 @@ def _jsonrpc_error(req_id: Any, code: int, message: str) -> dict[str, Any]: return {"jsonrpc": "2.0", "id": req_id, "error": {"code": code, "message": message}} +def _filter_skills(skills: list[SkillInfo], allowed_skills: list[str] | None) -> list[SkillInfo]: + """Keep only explicitly exposed skills when an MCP allowlist is configured.""" + if allowed_skills is None: + return skills + allowed = set(allowed_skills) + return [skill_info for skill_info in skills if skill_info.func_name in allowed] + + +def _select_module_skills( + modules: list[RPCClient], allowed_skills: list[str] | None +) -> list[tuple[RPCClient, SkillInfo]]: + """Return exposed skills together with their deployed module RPC address.""" + allowed = set(allowed_skills) if allowed_skills is not None else None + return [ + (module, skill_info) + for module in modules + for skill_info in (module.get_skills() or []) + if allowed is None or skill_info.func_name in allowed + ] + + def _handle_initialize(req_id: Any) -> dict[str, Any]: return _jsonrpc_result( req_id, @@ -345,6 +374,7 @@ async def event_generator() -> AsyncGenerator[str, None]: class McpServer(Module): + config: McpServerConfig _uvicorn_server: uvicorn.Server | None = None _serve_future: concurrent.futures.Future[None] | None = None _tool_stream_cleanup: Callable[[], None] | None = None @@ -381,15 +411,14 @@ def stop(self) -> None: def on_system_modules(self, modules: list[RPCClient]) -> None: # TODO: this is a bit hacky, also not thread-safe assert self.rpc is not None - app.state.skills = [ - skill_info for module in modules for skill_info in (module.get_skills() or []) - ] + module_skills = _select_module_skills(modules, self.config.allowed_skills) + app.state.skills = [skill_info for _, skill_info in module_skills] app.state.skills_by_name = {s.func_name: s for s in app.state.skills} app.state.rpc_calls = { skill_info.func_name: RpcCall( - None, self.rpc, skill_info.func_name, skill_info.class_name, [] + None, self.rpc, skill_info.func_name, module.remote_name, [] ) - for skill_info in app.state.skills + for module, skill_info in module_skills } @skill diff --git a/dimos/agents/mcp/test_mcp_client_unit.py b/dimos/agents/mcp/test_mcp_client_unit.py index a49df130ff..b2208d11b3 100644 --- a/dimos/agents/mcp/test_mcp_client_unit.py +++ b/dimos/agents/mcp/test_mcp_client_unit.py @@ -18,7 +18,7 @@ from threading import RLock from unittest.mock import MagicMock, create_autospec, patch -from langchain_core.messages import HumanMessage +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage from langchain_core.messages.base import BaseMessage from langchain_openai import ChatOpenAI import pytest @@ -163,6 +163,28 @@ def test_tool_stream_notification_becomes_human_message(mcp_client: McpClient) - assert "Person follow stopped: lost track." in str(msg.content) +def test_cancelled_turn_closes_unresolved_tool_calls(mcp_client: McpClient) -> None: + mcp_client._history = [ + AIMessage( + content="", + tool_calls=[ + {"name": "move_to_pose", "args": {}, "id": "call-pending", "type": "tool_call"} + ], + ), + AIMessage( + content="", + tool_calls=[{"name": "scan", "args": {}, "id": "call-complete", "type": "tool_call"}], + ), + ToolMessage(content="Detected 2 object(s)", tool_call_id="call-complete"), + ] + + mcp_client._close_cancelled_tool_calls() + + assert isinstance(mcp_client._history[-1], ToolMessage) + assert mcp_client._history[-1].tool_call_id == "call-pending" + assert "cancelled" in str(mcp_client._history[-1].content) + + def test_tool_stream_ignores_unrelated_frames(mcp_client: McpClient) -> None: """Unknown methods and empty bodies are dropped on the floor.""" diff --git a/dimos/agents/mcp/test_mcp_server.py b/dimos/agents/mcp/test_mcp_server.py index 0e2a74925b..e5a6851d84 100644 --- a/dimos/agents/mcp/test_mcp_server.py +++ b/dimos/agents/mcp/test_mcp_server.py @@ -20,7 +20,7 @@ from unittest.mock import MagicMock from dimos.agents.capabilities import CapabilityRegistry -from dimos.agents.mcp.mcp_server import app, handle_request +from dimos.agents.mcp.mcp_server import _filter_skills, _select_module_skills, app, handle_request from dimos.core.module import SkillInfo @@ -39,6 +39,26 @@ def _make_rpc_calls( return rpc_calls +def test_filter_skills_respects_allowlist() -> None: + schema = json.dumps({"type": "object", "properties": {}}) + skills = [ + SkillInfo(class_name="TestSkills", func_name="safe", args_schema=schema), + SkillInfo(class_name="TestSkills", func_name="unsafe", args_schema=schema), + ] + + assert [skill.func_name for skill in _filter_skills(skills, ["safe"])] == ["safe"] + assert _filter_skills(skills, None) == skills + + +def test_select_module_skills_retains_deployed_remote_name() -> None: + schema = json.dumps({"type": "object", "properties": {}}) + skill = SkillInfo(class_name="PickNPlaceModule", func_name="scan", args_schema=schema) + module = MagicMock(remote_name="pnp") + module.get_skills.return_value = [skill] + + assert _select_module_skills([module], ["scan"]) == [(module, skill)] + + def test_mcp_module_request_flow() -> None: schema = json.dumps( { diff --git a/dimos/agents/test_utils.py b/dimos/agents/test_utils.py new file mode 100644 index 0000000000..923afec9c4 --- /dev/null +++ b/dimos/agents/test_utils.py @@ -0,0 +1,25 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dimos.agents.utils import _display_message_content + + +def test_display_message_content_omits_encrypted_reasoning() -> None: + content = [ + {"type": "reasoning", "encrypted_content": "secret"}, + {"type": "function_call", "name": "go_home"}, + {"type": "text", "text": "Reached home."}, + ] + + assert _display_message_content(content) == "Reached home." diff --git a/dimos/agents/utils.py b/dimos/agents/utils.py index 5084c65b1f..475606707c 100644 --- a/dimos/agents/utils.py +++ b/dimos/agents/utils.py @@ -51,7 +51,7 @@ def pretty_print_langchain_message(msg: BaseMessage) -> None: time_str = f"{GRAY}{timestamp}{RESET} " type_str = f"{type_color}{msg_type:<{TYPE_WIDTH}}{RESET}" - content = _try_to_remove_url_data(d.get("content", "")) + content = _display_message_content(d.get("content", "")) tool_calls = d.get("tool_calls", []) # 12 chars for timestamp + 1 space + TYPE_WIDTH + 1 space @@ -96,16 +96,14 @@ def _log_message(msg_type: str, content: object, tool_calls: list[dict[str, Any] logger.info("Agent message", **kw) -def _try_to_remove_url_data(content: Any) -> Any: +def _display_message_content(content: Any) -> Any: + """Keep only user-visible text from OpenAI Responses content blocks.""" if not isinstance(content, list): return content - - ret = [] - + text_parts = [] for item in content: - if isinstance(item, dict) and item.get("type") == "image_url": - ret.append({**item, "image_url": ""}) - else: - ret.append(item) - - return ret + if isinstance(item, dict) and item.get("type") == "text": + text = item.get("text") + if isinstance(text, str): + text_parts.append(text) + return "\n".join(text_parts) diff --git a/dimos/cli/agentspy/agentspy.py b/dimos/cli/agentspy/agentspy.py index 8d844d8304..8484357d14 100644 --- a/dimos/cli/agentspy/agentspy.py +++ b/dimos/cli/agentspy/agentspy.py @@ -31,6 +31,7 @@ from textual.binding import Binding from textual.widgets import Footer, RichLog +from dimos.agents.utils import _display_message_content from dimos.cli import theme from dimos.core.transport_factory import apply_transport_arg, make_transport @@ -126,14 +127,14 @@ def format_message_content(msg: AnyMessage) -> str: for tc in msg.tool_calls: args_str = str(tc.get("args", {})) tool_info.append(f"{tc.get('name')}({args_str})") - content = msg.content or "" + content = _display_message_content(msg.content) if content and tool_info: return f"{content}\n[Tool Calls: {', '.join(tool_info)}]" elif tool_info: return f"[Tool Calls: {', '.join(tool_info)}]" return content # type: ignore[return-value] else: - return str(msg.content) if hasattr(msg, "content") else str(msg) + return _display_message_content(msg.content) if hasattr(msg, "content") else str(msg) class AgentSpyApp(App): # type: ignore[type-arg] diff --git a/dimos/cli/agentspy/test_agentspy.py b/dimos/cli/agentspy/test_agentspy.py new file mode 100644 index 0000000000..0e8e63f402 --- /dev/null +++ b/dimos/cli/agentspy/test_agentspy.py @@ -0,0 +1,28 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from langchain_core.messages import AIMessage + +from dimos.cli.agentspy.agentspy import format_message_content + + +def test_format_message_content_omits_responses_reasoning_blocks() -> None: + message = AIMessage( + content=[ + {"type": "reasoning", "encrypted_content": "secret"}, + {"type": "text", "text": "Placed the block."}, + ] + ) + + assert format_message_content(message) == "Placed the block." diff --git a/dimos/cli/human/humancli.py b/dimos/cli/human/humancli.py index 7fe5c9bbe3..360e68ad7c 100644 --- a/dimos/cli/human/humancli.py +++ b/dimos/cli/human/humancli.py @@ -85,6 +85,19 @@ def _split_tool_message(content: Any) -> tuple[str, str] | None: return content[len(TOOL_MSG_PREFIX) : end], content[end + 1 :].lstrip() +def _content_text(content: Any) -> str: + """Extract displayable text from string or OpenAI Responses content blocks.""" + if isinstance(content, str): + return content + if isinstance(content, list): + return "\n".join( + text + for item in content + if isinstance(item, dict) and isinstance((text := item.get("text")), str) + ) + return str(content) if content is not None else "" + + class ToolPanel: """Live state for one streaming tool's box. @@ -264,6 +277,7 @@ class HumanCLIApp(App): # type: ignore[type-arg] Binding("q", "quit", "Quit", show=False), Binding("ctrl+c", "quit", "Quit"), Binding("ctrl+l", "clear", "Clear chat"), + Binding("escape", "stop_agent", "Stop agent"), ] def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def] @@ -271,6 +285,7 @@ def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def] self._human_transport = make_transport("/human_input") self._agent_transport = make_transport("/agent") self._agent_idle = make_transport("/agent_idle") + self._agent_cancel = make_transport("/agent_cancel") self.chat_log: RichLog | None = None self.input_widget: Input | None = None self._subscription_thread: threading.Thread | None = None @@ -296,7 +311,7 @@ def compose(self) -> ComposeResult: yield Container(id="tool-panels") - self.input_widget = Input(placeholder="Type a message...") + self.input_widget = Input(placeholder="Type a message... Esc stops the current turn") yield self.input_widget def on_mount(self) -> None: @@ -360,7 +375,7 @@ def receive_msg(msg) -> None: # type: ignore[no-untyped-def] theme.YELLOW, ) elif isinstance(msg, AIMessage): - content = msg.content or "" + content = _content_text(msg.content) tool_calls = getattr(msg, "tool_calls", None) or msg.additional_kwargs.get( "tool_calls", [] ) @@ -634,6 +649,7 @@ def on_input_submitted(self, event: Input.Submitted) -> None: /exit - Exit the application /quit - Exit the application +Press Esc to stop the current agent turn. Tool calls are displayed in cyan with ▶ prefix""" self._add_system_message(help_text) return @@ -650,6 +666,16 @@ def action_clear(self) -> None: self._tool_call_anchors.clear() self.chat_log.clear() # type: ignore[union-attr] + def action_stop_agent(self) -> None: + """Request cancellation of the active agent turn and return to input.""" + if self._agent_is_idle: + return + self._agent_cancel.publish(True) + self._set_agent_idle(True) + self._add_system_message("Agent turn cancelled. You can enter a new message.") + if self.input_widget is not None: + self.input_widget.focus() + def action_quit(self) -> None: # type: ignore[override] """Quit the application.""" self._running = False diff --git a/dimos/cli/human/test_humancli.py b/dimos/cli/human/test_humancli.py new file mode 100644 index 0000000000..c398126633 --- /dev/null +++ b/dimos/cli/human/test_humancli.py @@ -0,0 +1,24 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dimos.cli.human.humancli import _content_text + + +def test_content_text_extracts_responses_api_text_blocks() -> None: + content = [ + {"type": "reasoning", "content": []}, + {"type": "text", "text": "The block is blue."}, + ] + + assert _content_text(content) == "The block is blue." diff --git a/dimos/control/coordinator.py b/dimos/control/coordinator.py index eeed8dead0..059d5c0699 100644 --- a/dimos/control/coordinator.py +++ b/dimos/control/coordinator.py @@ -913,7 +913,7 @@ def set_gripper_position(self, hardware_id: str, position: float) -> bool: if isinstance(hw, ConnectedTwistBase): logger.warning(f"Hardware '{hardware_id}' is a twist base, no gripper support") return False - return hw.adapter.write_gripper_position(position) + return hw.set_gripper_position(position) @rpc def get_gripper_position(self, hardware_id: str) -> float | None: diff --git a/dimos/control/hardware_interface.py b/dimos/control/hardware_interface.py index 3a7c74f430..bc3c5fb8ff 100644 --- a/dimos/control/hardware_interface.py +++ b/dimos/control/hardware_interface.py @@ -136,6 +136,17 @@ def read_state(self) -> dict[JointName, JointState]: return result + def set_gripper_position(self, position: float) -> bool: + """Command the gripper and preserve that command during arm trajectories.""" + if not self._gripper_joints: + return False + if not self._initialized: + self._initialize_last_commanded() + normalized_position = self._physical_to_normalized(position) + for joint_name in self._gripper_joints: + self._last_commanded[joint_name] = normalized_position + return self._adapter.write_gripper_position(position) + def write_command(self, commands: dict[str, float], mode: ControlMode) -> bool: """Write commands - allows partial joint sets, holds last for missing. diff --git a/dimos/core/test_transport_factory.py b/dimos/core/test_transport_factory.py index 419deb44f1..8e7a18c336 100644 --- a/dimos/core/test_transport_factory.py +++ b/dimos/core/test_transport_factory.py @@ -85,6 +85,7 @@ def test_default_zenoh_qos_agent_channels_never_drop() -> None: assert default_zenoh_qos("/human_input") == QOS_NEVER_DROP assert default_zenoh_qos("/agent") == QOS_NEVER_DROP assert default_zenoh_qos("/agent_idle") == QOS_NEVER_DROP + assert default_zenoh_qos("/agent_cancel") == QOS_NEVER_DROP def test_default_zenoh_qos_everything_else_uses_zenoh_defaults() -> None: diff --git a/dimos/core/transport_factory.py b/dimos/core/transport_factory.py index 2a69c270ec..04d9735d8d 100644 --- a/dimos/core/transport_factory.py +++ b/dimos/core/transport_factory.py @@ -58,7 +58,7 @@ def transport_topic(name: str, g: GlobalConfig = global_config) -> str: _LATEST_WINS_TYPES = ("sensor_msgs.Image", "sensor_msgs.PointCloud2") # Low-rate channels where a drop loses something that never comes back: a whole # turn of agent/human conversation, or a one-shot robot action verb. -_NEVER_DROP_CHANNELS = ("human_input", "agent", "agent_idle", "command") +_NEVER_DROP_CHANNELS = ("human_input", "agent", "agent_idle", "agent_cancel", "command") def default_zenoh_qos(name: str, msg_type: type | None = None) -> ZenohQoS | None: diff --git a/dimos/hardware/manipulators/xarm/adapter.py b/dimos/hardware/manipulators/xarm/adapter.py index c91292051c..73089c9271 100644 --- a/dimos/hardware/manipulators/xarm/adapter.py +++ b/dimos/hardware/manipulators/xarm/adapter.py @@ -128,7 +128,6 @@ def set_control_mode(self, mode: ControlMode) -> bool: """ if not self._arm: return False - mode_map = { ControlMode.POSITION: _XARM_MODE_POSITION, ControlMode.SERVO_POSITION: _XARM_MODE_SERVO_CARTESIAN, # Mode 1 for high-freq @@ -231,7 +230,6 @@ def activate(self) -> bool: """Enable motion and move the arm to its initial joint pose.""" if not self._arm: return False - self._prepare_for_position_motion() if not self._move_to_initial_pose(): return False @@ -241,12 +239,14 @@ def deactivate(self) -> bool: """Move the arm to its initial joint pose and enter stopped state.""" if not self._arm: return False - self._prepare_for_position_motion() homed = self._move_to_initial_pose() + gripper_opened = True + if self._gripper_enabled: + gripper_opened = self._arm.set_gripper_position(0.85 * M_TO_MM, wait=True) == 0 self._arm.motion_enable(enable=False) code: int = self._arm.set_state(4) - return homed and code == 0 + return homed and gripper_opened and code == 0 def _move_to_initial_pose(self) -> bool: if not self._arm: diff --git a/dimos/hardware/manipulators/xarm/test_adapter.py b/dimos/hardware/manipulators/xarm/test_adapter.py index c3e21d58bc..67a268ad81 100644 --- a/dimos/hardware/manipulators/xarm/test_adapter.py +++ b/dimos/hardware/manipulators/xarm/test_adapter.py @@ -82,6 +82,14 @@ def set_servo_angle_j(self, angles: list[float], *, speed: float, mvacc: float) self.actions.append(("set_servo_angle_j", list(angles), speed, mvacc)) return 0 + def set_gripper_enable(self, enable: bool) -> int: + self.actions.append(("set_gripper_enable", enable)) + return 0 + + def set_gripper_position(self, position: float, *, wait: bool) -> int: + self.actions.append(("set_gripper_position", position, wait)) + return 0 + @pytest.fixture def xarm_adapter_module(monkeypatch: pytest.MonkeyPatch) -> Iterator[ModuleType]: @@ -130,3 +138,18 @@ def test_joint_position_commands_use_degrees_for_xarm_sdk( arm = _FakeXArmSdk.instances[-1] assert arm.servo_joint_commands[-1] == pytest.approx([90.0, -45.0, 180.0]) + + +def test_deactivate_opens_an_enabled_gripper(xarm_adapter_module: ModuleType) -> None: + adapter = xarm_adapter_module.XArmAdapter(address="192.0.2.10", dof=6) + assert adapter.connect() + assert adapter.write_gripper_position(0.0) + + assert adapter.deactivate() + + arm = _FakeXArmSdk.instances[-1] + assert arm.actions[-3:] == [ + ("set_gripper_position", 850.0, True), + ("motion_enable", False), + ("set_state", 4), + ] diff --git a/dimos/manipulation/README.md b/dimos/manipulation/README.md new file mode 100644 index 0000000000..5717c1988a --- /dev/null +++ b/dimos/manipulation/README.md @@ -0,0 +1,128 @@ +# Pick And Place + +This directory contains the configurable xArm6 `picknplace` operator pipeline. +It uses the wrist-mounted RealSense and object-scene registration in `link_base`. + +## Setup + +GraspGenX runs in the main worktree `.venv` so it shares the live DimOS +pipeline. Its CUDA requirements differ from the repository lockfile; install +them once from the worktree root: + +```bash +bash bin/setup-graspgenx-env +``` + +The setup installs Torch 2.7.1 CUDA 12.8, which supports the RTX 5070's +`sm_120` architecture, along with GraspGenX and its inference dependencies. +Use `uv run --no-sync` afterwards. Plain `uv run` reconciles the environment to +the lockfile's Torch 2.6 and removes the GPU architecture support required by +GraspGenX. + +The first GraspGenX startup downloads the pinned model checkpoint to the +Hugging Face cache and loads it onto the GPU. Later starts reuse that cache. +The setup also installs `edgetam-dimos`, which provides the `sam2` runtime used +by the EdgeTAM blueprint. + +## Run + +Start the default YOLO-E and OBB-center-grasp pipeline: + +```bash +uv run --no-sync dimos run picknplace --daemon +``` + +Use text-prompted Moondream detection, EdgeTAM segmentation, and an OBB-center grasp: + +```bash +uv run --no-sync dimos run picknplace --daemon \ + -o osr.det=moondream -o osr.seg=edgetam -o pnp.grasp=obb_center +``` + +Use the same perception stack with GraspGenX: + +```bash +uv run --no-sync dimos run picknplace --daemon \ + -o osr.det=moondream -o osr.seg=edgetam -o pnp.grasp=graspgenx +``` + +`osr.det` accepts `yoloe` or `moondream`; `osr.seg` accepts `yolo` or `edgetam`. +Moondream requires EdgeTAM because it produces detection boxes rather than masks. +`pnp.grasp` accepts `obb_center` or `graspgenx`. GraspGenX loads only when selected. + +Then connect the console: + +```bash +uv run --no-sync python -m dimos.manipulation.pnpconsole +``` + +Stop a running pipeline with: + +```bash +uv run --no-sync dimos stop +``` + +## Operator Flow + +The console intentionally keeps planning and execution separate: + +1. Select `1` to scan the current scene. +2. Select `2` to inspect object number, name, and confidence. +3. Select `3` and choose an object. The GraspGenX blueprint prints its top + proposals and displays the selected grasp. Viser shows the selected object + cloud in amber, the grasp TCP axes in red, and the pre-grasp TCP axes in green. + The top ten proposals are filtered through collision-aware xArm IK; after + table calibration, candidates intersecting the table are omitted. +4. Select `4` to plan and preview the approach. Each Viser preview plays once + at a slow two-second duration. +5. Execute the approach only after inspecting the proposal and preview. +6. Select `6` to plan and preview descent, then select `7` to execute it. +7. Close the gripper with `8`, then select `9` to plan and preview ascent. +8. Select `10` to execute the ascent, `11` to open, and `13` to return home. +9. After a scene scan, select `14` to estimate and preview the tabletop. Once + the blue Viser outline matches the table, enter a collision clearance in + millimeters. The recommended clearance is 10 mm; enter `0` for no extra clearance. + The manual action installs the collision slab at the measured tabletop position for + all subsequent IK and trajectory plans. The pick-and-place blueprints do not install + a fixed floor slab. +10. After executing the approach, select `15` to collision-plan and execute the + descent, close the gripper, and execute the ascent without previews. It stops at + the first failed stage. + +Do not execute a learned grasp without checking its pose, the 100 mm pre-grasp +pose, the point-cloud/overlay visualization, and the collision-free preview. + +## Grasp Geometry + +`PickNPlaceModule.get_goal_pose()` stores the top ranked GraspGenX candidate as +the TCP goal in the candidate point cloud's frame. Its pre-grasp is computed as: + +```text +pre_grasp_position = grasp_position - grasp_orientation * (0, 0, 0.100 m) +``` + +GraspGenX local `+Z` is the final approach direction, so the pre-grasp retreats +along local `-Z`. It is not a world-Z lift: an angled or side grasp receives an +equally angled or sideward pre-grasp. Descent and ascent use Cartesian paths +between the current TCP pose and the selected grasp or pre-grasp target. + +The `picknplace-graspgenx` blueprint uses the xArm 85 mm gripper sweep-volume +and calibrated base-to-TCP transform. The TCP is rolled 90 degrees around the +GraspGenX approach axis so its closing jaws are perpendicular to a bottle's +length. Candidate score order comes from GraspGenX; no additional ranking is +applied by the operator pipeline. + +## Implementation Guide + +- `blueprints.py`: robot, camera, OBB, and GraspGenX blueprint composition. +- `picknplace.py`: scan request, target selection, OBB fallback, learned grasp + selection, and tool-axis pre-grasp calculation. +- `pnpconsole.py`: explicit operator stages and manual gripper/home controls. +- `grasping/grasp_gen_x.py`: import-safe proposal adapter and candidate contract. +- `grasping/grasp_gen_x_runtime.py`: in-process checkpoint load and GPU inference. +- `visualization/pose_overlay.py` and `visualization/rerun.py`: selected-object + cloud, image, and grasp overlays. + +The current scan is a single wrist-camera view. Automatic multi-view scanning, +EdgeTAM segmentation, and fused object clouds are planned follow-up work. Until +then, select targets with a complete enough visible point cloud for grasping. diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index 5c81178271..a5ad7ee5b4 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -17,6 +17,23 @@ Robot-owned manipulation blueprints now live under ``dimos.robot.manipulators``. """ +import math + +from dimos.agents.mcp.mcp_client import McpClient +from dimos.agents.mcp.mcp_server import McpServer +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.global_config import global_config +from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera +from dimos.manipulation.grasping.grasp_gen_x import GraspGenXModule +from dimos.manipulation.manipulation_module import ManipulationModule +from dimos.manipulation.picknplace import PickNPlaceModule +from dimos.manipulation.visualization.rerun import picknplace_rerun_config +from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.perception.experimental.object_scene_registration import ObjectSceneRegistrationModule +from dimos.robot.manipulators.common.blueprints import coordinator, trajectory_task from dimos.robot.manipulators.xarm.blueprints.agentic import ( xarm7_planner_coordinator_agent as xarm7_planner_coordinator_agent, xarm_perception_agent as xarm_perception_agent, @@ -29,3 +46,99 @@ from dimos.robot.manipulators.xarm.blueprints.simulation import ( xarm_perception_sim as xarm_perception_sim, ) +from dimos.robot.manipulators.xarm.config import make_xarm6_model_config, xarm6_hardware +from dimos.robot.manipulators.xarm.grasp_config import make_xarm_graspgenx_config +from dimos.visualization.vis_module import vis_module + +PICKNPLACE_CAMERA_TRANSFORM = Transform( + translation=Vector3(0.06693724, -0.0309563, 0.00691482), + rotation=Quaternion(0.70513398, 0.00535696, 0.70897578, -0.01052180), +) + +BOX_FILLING_SYSTEM_PROMPT = """You are operating an xArm box-filling workspace with RGB-D perception. + +Your recurring task is to collect requested blocks from the table and drop them into the measured white box. The available tools are the live interface to the robot, planner, gripper, and scene. Use their results as authoritative, make multiple calls when needed, and only report physical actions after a tool confirms success. + +For a collection task: go home to observe, use ``scan_objects`` with separate simple noun phrases such as ``["colored wooden block", "white box"]``, estimate and install the table collision with no added margin, and measure the white box with ``install_open_box``. Use ``get_object_geometry`` to identify blocks whose centers are inside the measured box opening; those blocks are complete and must be ignored. Select only outside blocks, then call ``pick_selected``. If it succeeds, call ``place_selected`` to drop it into the remembered box. Repeat for other outside blocks. If pickup verification fails, rescan and select before another attempt. + +When the user says put, place, or drop an object in the box, use ``place_selected``. It is a depth-derived drop: it computes the box-rim and held-object clearance itself, releases above the rim, and does not lower the end effector into the box. Do not substitute manually chosen poses or individual gripper commands for pick or drop sequences. +""" + +_picknplace_xarm6_hardware = xarm6_hardware("arm", gripper=True) +_picknplace_xarm6_model = make_xarm6_model_config( + name="arm", + add_gripper=True, + tf_extra_links=["link_base", "link6"], + home_joints=[0.0, math.radians(-40.0), math.radians(-50.0), 0.0, math.radians(90.0), 0.0], +) +_picknplace_xarm6_model.max_velocity = 0.25 +_picknplace_xarm6_model.max_acceleration = 0.5 +_xarm_graspgenx = make_xarm_graspgenx_config() + + +picknplace = autoconnect( + coordinator( + hardware=[_picknplace_xarm6_hardware], + tasks=[trajectory_task(_picknplace_xarm6_hardware)], + ), + ManipulationModule.blueprint( + robots=[_picknplace_xarm6_model], + visualization=ViserVisualizationConfig(port=8095), + planning_timeout=10.0, + ), + RealSenseCamera.blueprint( + width=848, + height=480, + fps=15, + camera_name="camera", + base_frame_id="link6", + base_transform=PICKNPLACE_CAMERA_TRANSFORM, + enable_depth=True, + align_depth_to_color=True, + enable_pointcloud=False, + ), + ObjectSceneRegistrationModule.blueprint( + instance_name="osr", + target_frame="link_base", + register_objects=False, + detect_on_request=True, + detector_confidence=0.4, + object_voxel_downsample=0.001, + ), + PickNPlaceModule.blueprint(instance_name="pnp", align_grasp_yaw=True), + GraspGenXModule.blueprint( + instance_name="ggx", + load_on_start=False, + **_xarm_graspgenx.model_dump( + exclude={"rpc_transport", "tf_transport", "g", "instance_name", "load_on_start"} + ), + ), + vis_module( + global_config.viewer, + rerun_config=picknplace_rerun_config(), + ), +).global_config(rerun_open="web") + +picknplace_agent = autoconnect( + picknplace, + McpServer.blueprint( + allowed_skills=[ + "describe_scene", + "scan_objects", + "estimate_table", + "select_object", + "pick_selected", + "place_selected", + "get_object_geometry", + "install_open_box", + "set_table_collision", + "get_robot_state", + "reset", + "move_to_pose", + "close_gripper", + "open_gripper", + "go_home", + ] + ), + McpClient.blueprint(system_prompt=BOX_FILLING_SYSTEM_PROMPT), +) diff --git a/dimos/manipulation/candidate_filter_spec.py b/dimos/manipulation/candidate_filter_spec.py new file mode 100644 index 0000000000..da6ca47c99 --- /dev/null +++ b/dimos/manipulation/candidate_filter_spec.py @@ -0,0 +1,32 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Planner contract used to reject unsafe learned grasp proposals.""" + +from typing import Protocol + +from dimos.manipulation.planning.spec.models import IKResult, RobotName +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.spec.utils import Spec + + +class GraspCandidateFilterSpec(Spec, Protocol): + def inverse_kinematics_single( + self, + pose: Pose, + robot_name: RobotName | None = None, + seed: JointState | None = None, + check_collision: bool = True, + ) -> IKResult: ... diff --git a/dimos/manipulation/demo_grasp_pipeline/__main__.py b/dimos/manipulation/demo_grasp_pipeline/__main__.py new file mode 100644 index 0000000000..9711eb6577 --- /dev/null +++ b/dimos/manipulation/demo_grasp_pipeline/__main__.py @@ -0,0 +1,91 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run the offline GraspGenX-to-connected-motion-planning pipeline.""" + +import argparse +from collections.abc import Sequence +from pathlib import Path +import sys + +from .demo import run_contributor_demo + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("grasp-pipeline-demo"), + help="Directory for summary.json, plans.json, and the selected-grasp PNG.", + ) + parser.add_argument( + "--max-candidates", + type=int, + default=20, + help="Maximum number of ranked proposals to plan.", + ) + parser.add_argument( + "--workspace-center", + nargs=3, + type=float, + metavar=("X", "Y", "Z"), + default=(0.45, 0.0, 0.25), + help="Target-cloud centroid in the synthetic xArm world, in metres.", + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parser().parse_args(argv) + try: + result = run_contributor_demo( + output_dir=args.output_dir, + max_candidates=args.max_candidates, + workspace_center=args.workspace_center, + ) + except Exception as exc: + cause: BaseException | None = exc + while cause is not None: + if isinstance(cause, ModuleNotFoundError) and cause.name == "graspgenx": + print( + "GraspGenX is not installed. Run this demo with " + "`uv run --extra graspgenx python -m " + "dimos.manipulation.demo_grasp_pipeline ...` " + f"({cause})", + file=sys.stderr, + flush=True, + ) + return 2 + cause = cause.__cause__ + raise + for outcome in result.outcomes: + detail = f" rejection={outcome.rejection}" if outcome.rejection else "" + print( + f"candidate rank={outcome.rank} score={outcome.score:.6f} " + f"status={outcome.status}{detail}", + flush=True, + ) + status = "selected" if result.success else result.failure_reason + print( + f"grasp-pipeline-demo status={status} candidates={result.candidate_count} " + f"summary={result.summary_path} plans={result.plans_path} " + f"visualization={result.image_path}", + flush=True, + ) + return 0 if result.success else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dimos/manipulation/demo_grasp_pipeline/demo.py b/dimos/manipulation/demo_grasp_pipeline/demo.py new file mode 100644 index 0000000000..3f29aee877 --- /dev/null +++ b/dimos/manipulation/demo_grasp_pipeline/demo.py @@ -0,0 +1,443 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run grasp proposal and connected motion planning without robot execution.""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Callable, Sequence +from dataclasses import asdict, dataclass +import json +from pathlib import Path +from typing import Any + +import numpy as np + +from dimos.manipulation.demo_graspgenx.fixture import load_demo_clouds +from dimos.manipulation.demo_graspgenx.render import SweepVolumeLike, render_grasp_image +from dimos.manipulation.grasping.grasp_gen_spec import GraspGenSpec +from dimos.manipulation.grasping.grasp_gen_x import ( + IDENTITY_TRANSFORM, + GraspGenXModule, + RigidTransform, +) +from dimos.manipulation.manipulation_module import ConnectedPoseSequenceResult +from dimos.manipulation.pick_and_place_module import PickAndPlaceModule +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.std_msgs.Header import Header +from dimos.robot.manipulators.xarm.config import make_xarm7_sim_robot_config +from dimos.robot.manipulators.xarm.grasp_config import make_xarm_graspgenx_config + +CloudLoader = Callable[[], tuple[PointCloud2, PointCloud2]] +Renderer = Callable[ + [Path, PointCloud2, PointCloud2, GraspCandidateArray, SweepVolumeLike, int], Path +] + +_REJECTION_BY_INDEX = ( + "pre_grasp_infeasible", + "grasp_infeasible", + "retreat_infeasible", +) + + +def render_selected_grasp_image( + output_path: Path, + scene: PointCloud2, + object_cloud: PointCloud2, + candidates: GraspCandidateArray, + gripper: SweepVolumeLike, + rank: int, +) -> Path: + """Render one selected grasp while preserving its original proposal rank.""" + return render_grasp_image( + output_path, + scene, + object_cloud, + candidates, + gripper, + ranks=(rank,), + title=f"Selected connected grasp — proposal rank #{rank}", + ) + + +@dataclass(frozen=True) +class CandidateOutcome: + """Planning outcome for one ranked proposal.""" + + rank: int + score: float + status: str + rejection: str | None = None + + +@dataclass(frozen=True) +class PipelineDemoResult: + """Artifacts and candidate outcomes produced by one offline run.""" + + success: bool + output_dir: Path + summary_path: Path + plans_path: Path + image_path: Path | None + candidate_count: int + selected_rank: int | None + selected_score: float | None + outcomes: tuple[CandidateOutcome, ...] + failure_reason: str | None = None + + +def _relocate_clouds( + scene: PointCloud2, + object_cloud: PointCloud2, + workspace_center: Sequence[float], +) -> tuple[PointCloud2, PointCloud2]: + object_points = object_cloud.points_f32() + if not len(object_points): + raise ValueError("recorded target cloud is empty") + center = np.asarray(workspace_center, dtype=np.float32) + if center.shape != (3,) or not np.all(np.isfinite(center)): + raise ValueError("workspace center must contain three finite values") + translation = center - np.mean(object_points, axis=0) + relocated_scene = PointCloud2.from_numpy( + scene.points_f32() + translation, + frame_id=scene.frame_id, + timestamp=scene.ts, + ) + relocated_object = PointCloud2.from_numpy( + object_points + translation, + frame_id=object_cloud.frame_id, + timestamp=object_cloud.ts, + ) + return relocated_scene, relocated_object + + +def _joint_state_dict(state: JointState) -> dict[str, Any]: + return { + "names": list(state.name), + "positions": [float(value) for value in state.position], + } + + +def _pose_dict(pose: Pose) -> dict[str, list[float]]: + return { + "position": [ + float(pose.position.x), + float(pose.position.y), + float(pose.position.z), + ], + "orientation_xyzw": [ + float(pose.orientation.x), + float(pose.orientation.y), + float(pose.orientation.z), + float(pose.orientation.w), + ], + } + + +def _grasp_frame_candidate( + candidate: GraspCandidate, + grasp_frame_to_tcp: RigidTransform, +) -> GraspCandidate: + """Convert a planned TCP candidate back to its sweep-geometry frame.""" + world_to_tcp = np.eye(4, dtype=float) + world_to_tcp[:3, :3] = candidate.pose.orientation.to_rotation_matrix() + world_to_tcp[:3, 3] = np.asarray(candidate.pose.position.as_tuple, dtype=float) + world_to_grasp = world_to_tcp @ np.linalg.inv(np.asarray(grasp_frame_to_tcp, dtype=float)) + grasp_pose = Pose( + Vector3(world_to_grasp[:3, 3]), + Quaternion.from_rotation_matrix(world_to_grasp[:3, :3]), + ) + return GraspCandidate(grasp_pose, candidate.score) + + +def _segment_dicts( + names: Sequence[str], + result: ConnectedPoseSequenceResult, +) -> list[dict[str, Any]]: + return [ + { + "name": name, + "waypoints": [_joint_state_dict(state) for state in path], + } + for name, path in zip(names[: len(result.paths)], result.paths, strict=True) + ] + + +def _write_artifacts( + *, + output_dir: Path, + success: bool, + frame: str, + scene_points: int, + object_points: int, + candidate_count: int, + outcomes: Sequence[CandidateOutcome], + selected_rank: int | None, + selected: GraspCandidate | None, + segments: Sequence[dict[str, Any]], + image_path: Path | None, + failure_reason: str | None, +) -> tuple[Path, Path]: + output_dir.mkdir(parents=True, exist_ok=True) + rejection_counts = Counter( + outcome.rejection for outcome in outcomes if outcome.rejection is not None + ) + summary: dict[str, Any] = { + "success": success, + "failure_reason": failure_reason, + "frame": frame, + "scene_points": scene_points, + "object_points": object_points, + "candidate_count": candidate_count, + "checked_count": len(outcomes), + "selected": ( + { + "rank": selected_rank, + "score": float(selected.score), + "pose": _pose_dict(selected.pose), + } + if selected is not None + else None + ), + "candidate_outcomes": [asdict(outcome) for outcome in outcomes], + "rejection_counts": dict(sorted(rejection_counts.items())), + "execution_performed": False, + "artifacts": { + "plans": "plans.json", + "visualization": image_path.name if image_path is not None else None, + }, + } + plans = { + "frame": frame, + "execution_performed": False, + "segments": list(segments), + } + summary_path = output_dir / "summary.json" + plans_path = output_dir / "plans.json" + summary_path.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") + plans_path.write_text(json.dumps(plans, indent=2) + "\n", encoding="utf-8") + return summary_path, plans_path + + +def run_demo( + proposer: GraspGenSpec, + planner: PickAndPlaceModule, + output_dir: Path, + *, + gripper: SweepVolumeLike, + max_candidates: int = 20, + workspace_center: Sequence[float] = (0.45, 0.0, 0.25), + sequence_start: JointState | None = None, + grasp_frame_to_tcp: RigidTransform = IDENTITY_TRANSFORM, + cloud_loader: CloudLoader = load_demo_clouds, + renderer: Renderer = render_selected_grasp_image, +) -> PipelineDemoResult: + """Run proposal and connected planning, save artifacts, and never execute.""" + if max_candidates <= 0: + raise ValueError("max_candidates must be positive") + + scene, object_cloud = _relocate_clouds(*cloud_loader(), workspace_center) + proposals = proposer.propose_grasps(object_cloud) + if ( + proposals.header.frame_id != object_cloud.frame_id + or proposals.header.timestamp != object_cloud.ts + ): + raise ValueError("proposer changed the target cloud frame or timestamp") + ranked = sorted(proposals.candidates, key=lambda candidate: candidate.score, reverse=True) + robot = planner._get_robot("arm") + if robot is None: + raise ValueError("offline planner does not contain the xArm planning model") + pre_grasp_offset = float(robot[2].pre_grasp_offset) + approach = Vector3(planner.config.grasp_approach_vector) + + segments: list[dict[str, Any]] = [] + lift_pose = planner._safety_lift_pose("arm") + if lift_pose is not None: + lift_result = planner._plan_connected_pose_sequence( + (lift_pose,), + "arm", + sequence_start, + ) + segments.extend(_segment_dicts(("safety_lift",), lift_result)) + if lift_result.failed_index is not None: + summary_path, plans_path = _write_artifacts( + output_dir=output_dir, + success=False, + frame=proposals.header.frame_id, + scene_points=len(scene), + object_points=len(object_cloud), + candidate_count=len(ranked), + outcomes=(), + selected_rank=None, + selected=None, + segments=segments, + image_path=None, + failure_reason="safety_lift_infeasible", + ) + return PipelineDemoResult( + False, + output_dir, + summary_path, + plans_path, + None, + len(ranked), + None, + None, + (), + "safety_lift_infeasible", + ) + sequence_start = lift_result.endpoint + + outcomes: list[CandidateOutcome] = [] + selected: GraspCandidate | None = None + selected_rank: int | None = None + candidate_segments: list[dict[str, Any]] = [] + for rank, candidate in enumerate(ranked[:max_candidates], start=1): + if not planner._valid_candidate(candidate): + outcomes.append(CandidateOutcome(rank, candidate.score, "rejected", "invalid")) + continue + pre_grasp = planner._compute_pre_grasp_pose( + candidate.pose, + pre_grasp_offset, + approach, + ) + retreat = planner._compute_pre_grasp_pose( + candidate.pose, + pre_grasp_offset, + approach, + ) + result = planner._plan_connected_pose_sequence( + (pre_grasp, candidate.pose, retreat), + "arm", + sequence_start, + ) + if result.failed_index is not None: + outcomes.append( + CandidateOutcome( + rank, + candidate.score, + "rejected", + _REJECTION_BY_INDEX[result.failed_index], + ) + ) + continue + selected = candidate + selected_rank = rank + candidate_segments = _segment_dicts(("pre_grasp", "grasp", "retreat"), result) + outcomes.append(CandidateOutcome(rank, candidate.score, "selected")) + break + + image_path: Path | None = None + if selected is not None: + assert selected_rank is not None + output_dir.mkdir(parents=True, exist_ok=True) + image_path = renderer( + output_dir / "selected-grasp.png", + scene, + object_cloud, + GraspCandidateArray( + Header(proposals.header.timestamp, proposals.header.frame_id), + [_grasp_frame_candidate(selected, grasp_frame_to_tcp)], + ), + gripper, + selected_rank, + ) + segments.extend(candidate_segments) + + summary_path, plans_path = _write_artifacts( + output_dir=output_dir, + success=selected is not None, + frame=proposals.header.frame_id, + scene_points=len(scene), + object_points=len(object_cloud), + candidate_count=len(ranked), + outcomes=outcomes, + selected_rank=selected_rank, + selected=selected, + segments=segments, + image_path=image_path, + failure_reason=None if selected is not None else "no_complete_candidate", + ) + return PipelineDemoResult( + selected is not None, + output_dir, + summary_path, + plans_path, + image_path, + len(ranked), + selected_rank, + float(selected.score) if selected is not None else None, + tuple(outcomes), + None if selected is not None else "no_complete_candidate", + ) + + +def run_contributor_demo( + *, + output_dir: Path, + max_candidates: int = 20, + workspace_center: Sequence[float] = (0.45, 0.0, 0.25), +) -> PipelineDemoResult: + """Build real GraspGenX and xArm planning modules for one offline run.""" + grasp_config = make_xarm_graspgenx_config() + robot_config = make_xarm7_sim_robot_config() + proposer = GraspGenXModule( + **grasp_config.model_dump(exclude={"rpc_transport", "tf_transport", "g"}) + ) + planner = PickAndPlaceModule( + robots=[robot_config], + planning_timeout=10.0, + visualization={"backend": "none"}, + floor_z=None, + ) + # This standalone module has no blueprint streams to subscribe to. + standalone_planner: Any = planner + standalone_planner.coordinator_joint_state = None + standalone_planner.objects = None + proposer_started = False + planner_started = False + try: + proposer.start() + proposer_started = True + planner.start() + planner_started = True + if robot_config.home_joints is None: + raise ValueError("xArm demo configuration has no home joint state") + synthetic_start = JointState( + name=list(robot_config.get_coordinator_joint_names()), + position=list(robot_config.home_joints), + ) + planner._on_joint_state(synthetic_start) + return run_demo( + proposer, + planner, + output_dir, + gripper=grasp_config.gripper, + max_candidates=max_candidates, + workspace_center=workspace_center, + sequence_start=synthetic_start, + grasp_frame_to_tcp=grasp_config.grasp_frame_to_tcp, + ) + finally: + if planner_started: + planner.stop() + if proposer_started: + proposer.stop() diff --git a/dimos/manipulation/demo_grasp_pipeline/test_demo.py b/dimos/manipulation/demo_grasp_pipeline/test_demo.py new file mode 100644 index 0000000000..40f64ef4d0 --- /dev/null +++ b/dimos/manipulation/demo_grasp_pipeline/test_demo.py @@ -0,0 +1,346 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hermetic coverage for the offline proposal-to-motion-planning demo.""" + +from __future__ import annotations + +import json +from pathlib import Path +import subprocess +import sys +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +from pytest_mock import MockerFixture + +from dimos.manipulation.manipulation_module import ConnectedPoseSequenceResult +from dimos.manipulation.pick_and_place_module import PickAndPlaceModule +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.std_msgs.Header import Header +from dimos.robot.manipulators.xarm.grasp_config import make_xarm_graspgenx_config + +from . import __main__ +from .demo import PipelineDemoResult, run_demo + + +def _clouds() -> tuple[PointCloud2, PointCloud2]: + points = np.asarray( + [ + [1.0, 2.0, 3.0], + [1.1, 2.0, 3.0], + [1.0, 2.1, 3.0], + ], + dtype=np.float32, + ) + return ( + PointCloud2.from_numpy(points, frame_id="world", timestamp=42.0), + PointCloud2.from_numpy(points[:2], frame_id="world", timestamp=42.0), + ) + + +def _candidates() -> list[GraspCandidate]: + return [ + GraspCandidate( + Pose( + { + "position": [0.45, 0.0, 0.25], + "orientation": [0.0, 0.0, 0.0, 1.0], + } + ), + score, + ) + for score in (0.9, 0.8) + ] + + +def _path(start: float, goal: float) -> tuple[JointState, JointState]: + return ( + JointState(name=["arm/joint1"], position=[start]), + JointState(name=["arm/joint1"], position=[goal]), + ) + + +def _planner(mocker: MockerFixture) -> Any: + planner = mocker.Mock(spec=PickAndPlaceModule) + planner.config = SimpleNamespace(grasp_approach_vector=(0.0, 0.0, -1.0)) + planner._get_robot.return_value = ("arm", "robot", SimpleNamespace(pre_grasp_offset=0.05), None) + planner._safety_lift_pose.return_value = None + planner._valid_candidate.return_value = True + planner._compute_pre_grasp_pose.side_effect = lambda pose, _offset, _approach: pose + return planner + + +def _proposer(mocker: MockerFixture) -> Any: + proposer = mocker.Mock() + proposer.propose_grasps.return_value = GraspCandidateArray( + Header(42.0, "world"), + _candidates(), + ) + return proposer + + +def test_demo_selects_first_complete_sequence_and_writes_artifacts( + mocker: MockerFixture, + tmp_path: Path, +) -> None: + planner = _planner(mocker) + proposer = _proposer(mocker) + planner._plan_connected_pose_sequence.side_effect = [ + ConnectedPoseSequenceResult(1, None, (_path(0.0, 0.1),)), + ConnectedPoseSequenceResult( + None, + JointState(name=["arm/joint1"], position=[0.3]), + (_path(0.0, 0.1), _path(0.1, 0.2), _path(0.2, 0.3)), + ), + ] + + def write_image(path: Path, *_args: Any) -> Path: + path.write_bytes(b"png") + return path + + renderer = mocker.Mock(side_effect=write_image) + output = tmp_path / "pipeline" + + result = run_demo( + proposer, + planner, + output, + gripper=make_xarm_graspgenx_config().gripper, + grasp_frame_to_tcp=( + (1.0, 0.0, 0.0, 0.0), + (0.0, 1.0, 0.0, 0.0), + (0.0, 0.0, 1.0, 0.1), + (0.0, 0.0, 0.0, 1.0), + ), + cloud_loader=_clouds, + renderer=renderer, + ) + + assert result.success is True + assert result.selected_rank == 2 + assert result.selected_score == 0.8 + assert [outcome.rejection for outcome in result.outcomes] == [ + "grasp_infeasible", + None, + ] + summary = json.loads(result.summary_path.read_text(encoding="utf-8")) + assert summary["selected"]["rank"] == 2 + assert summary["rejection_counts"] == {"grasp_infeasible": 1} + assert summary["execution_performed"] is False + plans = json.loads(result.plans_path.read_text(encoding="utf-8")) + assert [segment["name"] for segment in plans["segments"]] == [ + "pre_grasp", + "grasp", + "retreat", + ] + assert plans["segments"][-1]["waypoints"][-1]["positions"] == [0.3] + selected_array = renderer.call_args.args[3] + assert selected_array.candidates[0].pose.position.z == pytest.approx(0.15) + assert summary["selected"]["pose"]["position"][2] == pytest.approx(0.25) + assert renderer.call_args.args[5] == 2 + assert result.image_path is not None + assert result.image_path.read_bytes() == b"png" + planner.execute.assert_not_called() + planner.set_gripper.assert_not_called() + + +def test_demo_reuses_an_explicit_synthetic_start_for_every_candidate( + mocker: MockerFixture, + tmp_path: Path, +) -> None: + planner = _planner(mocker) + proposer = _proposer(mocker) + start = JointState(name=["arm/joint1"], position=[0.0]) + planner._plan_connected_pose_sequence.side_effect = [ + ConnectedPoseSequenceResult(0, None, ()), + ConnectedPoseSequenceResult(0, None, ()), + ] + + run_demo( + proposer, + planner, + tmp_path, + gripper=make_xarm_graspgenx_config().gripper, + sequence_start=start, + cloud_loader=_clouds, + renderer=mocker.Mock(), + ) + + assert all( + call.args[2] is start for call in planner._plan_connected_pose_sequence.call_args_list + ) + + +def test_demo_records_shared_lift_and_exhausted_candidates( + mocker: MockerFixture, + tmp_path: Path, +) -> None: + planner = _planner(mocker) + proposer = _proposer(mocker) + planner._safety_lift_pose.return_value = _candidates()[0].pose + lift_endpoint = JointState(name=["arm/joint1"], position=[0.1]) + planner._plan_connected_pose_sequence.side_effect = [ + ConnectedPoseSequenceResult(None, lift_endpoint, (_path(0.0, 0.1),)), + ConnectedPoseSequenceResult(0, None, ()), + ConnectedPoseSequenceResult(2, None, (_path(0.1, 0.2), _path(0.2, 0.3))), + ] + renderer = mocker.Mock() + + result = run_demo( + proposer, + planner, + tmp_path, + gripper=make_xarm_graspgenx_config().gripper, + cloud_loader=_clouds, + renderer=renderer, + ) + + assert result.success is False + assert result.failure_reason == "no_complete_candidate" + assert result.selected_rank is None + summary = json.loads(result.summary_path.read_text(encoding="utf-8")) + assert summary["selected"] is None + assert summary["rejection_counts"] == { + "pre_grasp_infeasible": 1, + "retreat_infeasible": 1, + } + plans = json.loads(result.plans_path.read_text(encoding="utf-8")) + assert [segment["name"] for segment in plans["segments"]] == ["safety_lift"] + candidate_calls = planner._plan_connected_pose_sequence.call_args_list[1:] + assert all(call.args[2] is lift_endpoint for call in candidate_calls) + renderer.assert_not_called() + planner.execute.assert_not_called() + planner.set_gripper.assert_not_called() + + +def test_demo_records_shared_lift_failure_without_screening_candidates( + mocker: MockerFixture, + tmp_path: Path, +) -> None: + planner = _planner(mocker) + proposer = _proposer(mocker) + planner._safety_lift_pose.return_value = _candidates()[0].pose + planner._plan_connected_pose_sequence.return_value = ConnectedPoseSequenceResult( + 0, + None, + (), + ) + renderer = mocker.Mock() + + result = run_demo( + proposer, + planner, + tmp_path, + gripper=make_xarm_graspgenx_config().gripper, + cloud_loader=_clouds, + renderer=renderer, + ) + + assert result.success is False + assert result.failure_reason == "safety_lift_infeasible" + summary = json.loads(result.summary_path.read_text(encoding="utf-8")) + assert summary["failure_reason"] == "safety_lift_infeasible" + assert summary["checked_count"] == 0 + assert json.loads(result.plans_path.read_text(encoding="utf-8"))["segments"] == [] + assert planner._plan_connected_pose_sequence.call_count == 1 + renderer.assert_not_called() + planner.execute.assert_not_called() + planner.set_gripper.assert_not_called() + + +def test_demo_import_does_not_load_optional_graspgenx_runtime() -> None: + result = subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys; " + "import dimos.manipulation.demo_grasp_pipeline.demo; " + "assert 'dimos.manipulation.grasping.grasp_gen_x_runtime' not in sys.modules" + ), + ], + check=False, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stderr + + +def test_cli_explains_how_to_install_missing_graspgenx( + mocker: MockerFixture, + capsys: pytest.CaptureFixture[str], +) -> None: + missing = ModuleNotFoundError("No module named 'graspgenx'", name="graspgenx") + failure = RuntimeError("failed to initialize GraspGenX") + failure.__cause__ = missing + mocker.patch.object(__main__, "run_contributor_demo", side_effect=failure) + + assert __main__.main([]) == 2 + output = capsys.readouterr().err + assert "uv run --extra graspgenx" in output + assert "No module named 'graspgenx'" in output + + +@pytest.mark.parametrize(("success", "exit_code"), [(True, 0), (False, 1)]) +def test_cli_forwards_options_and_reports_result( + mocker: MockerFixture, + tmp_path: Path, + success: bool, + exit_code: int, +) -> None: + run = mocker.patch.object( + __main__, + "run_contributor_demo", + return_value=PipelineDemoResult( + success, + tmp_path, + tmp_path / "summary.json", + tmp_path / "plans.json", + None, + 2, + 1 if success else None, + 0.9 if success else None, + (), + None if success else "no_complete_candidate", + ), + ) + + assert ( + __main__.main( + [ + "--output-dir", + str(tmp_path), + "--max-candidates", + "7", + "--workspace-center", + "0.4", + "0.1", + "0.3", + ] + ) + == exit_code + ) + run.assert_called_once_with( + output_dir=tmp_path, + max_candidates=7, + workspace_center=[0.4, 0.1, 0.3], + ) diff --git a/dimos/manipulation/demo_grasp_visualization/__main__.py b/dimos/manipulation/demo_grasp_visualization/__main__.py new file mode 100644 index 0000000000..1bf78f88f3 --- /dev/null +++ b/dimos/manipulation/demo_grasp_visualization/__main__.py @@ -0,0 +1,70 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Visualize the banana point cloud and GraspGenX proposals in Viser.""" + +import argparse +from collections.abc import Sequence + +from .demo import DEFAULT_MAX_CANDIDATES, run_contributor_demo + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("must be positive") + return parsed + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--max-candidates", + type=_positive_int, + default=DEFAULT_MAX_CANDIDATES, + help="Maximum score-ranked grasp wireframes to display.", + ) + return parser + + +def _install_hint(error: BaseException) -> str | None: + current: BaseException | None = error + while current is not None: + if isinstance(current, ModuleNotFoundError): + return ( + "Install the visualization and GraspGenX dependencies, then retry: " + "`uv sync --extra manipulation --extra graspgenx`" + ) + current = current.__cause__ + return None + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parser().parse_args(argv) + try: + run_contributor_demo(max_candidates=args.max_candidates) + except KeyboardInterrupt: + print("grasp-visualization-demo stopped", flush=True) + return 0 + except Exception as error: + hint = _install_hint(error) + if hint is None: + raise + print(f"grasp-visualization-demo failed: {error}\n{hint}", flush=True) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dimos/manipulation/demo_grasp_visualization/demo.py b/dimos/manipulation/demo_grasp_visualization/demo.py new file mode 100644 index 0000000000..c6597697fc --- /dev/null +++ b/dimos/manipulation/demo_grasp_visualization/demo.py @@ -0,0 +1,177 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Build display-only layers from the banana fixture and real grasp proposals.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from threading import Event +from typing import Protocol + +import numpy as np +from numpy.typing import NDArray + +from dimos.manipulation.demo_graspgenx.demo import _validate +from dimos.manipulation.demo_graspgenx.fixture import load_demo_clouds +from dimos.manipulation.demo_graspgenx.render import gripper_wireframe_geometry +from dimos.manipulation.grasping.grasp_gen_spec import GraspGenSpec +from dimos.manipulation.grasping.grasp_gen_x import ( + GraspGenXConfig, + GraspGenXModule, + RigidTransform, + SweepVolumeGripperConfig, +) +from dimos.manipulation.visualization.layers import ( + LineSetElement, + PointCloudElement, + VisualizationLayer, +) +from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig +from dimos.manipulation.visualization.viser.visualizer import ViserManipulationVisualizer +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.robot.manipulators.xarm.grasp_config import make_xarm_graspgenx_config + +DEFAULT_MAX_CANDIDATES = 20 +OBJECT_CLOUD_LAYER_ID = "grasp/object-cloud" +PROPOSAL_LAYER_ID = "grasp/proposals" + +CloudLoader = Callable[[], tuple[PointCloud2, PointCloud2]] +Waiter = Callable[[], None] + + +class LayerVisualizer(Protocol): + def set_layer(self, layer: VisualizationLayer) -> None: ... + + def get_visualization_url(self) -> str | None: ... + + def close(self) -> None: ... + + +@dataclass(frozen=True) +class GraspVisualizationDemoResult: + """Summary of layers submitted by one interactive visualization run.""" + + candidate_count: int + displayed_count: int + frame_id: str + visualization_url: str | None + + +def _rank_color(rank_index: int, count: int) -> NDArray[np.uint8]: + fraction = 0.0 if count <= 1 else rank_index / (count - 1) + start = np.asarray([0, 220, 80], dtype=float) + end = np.asarray([255, 140, 0], dtype=float) + return np.asarray( + np.rint(start + fraction * (end - start)), + dtype=np.uint8, + ) + + +def run_demo( + proposer: GraspGenSpec, + visualizer: LayerVisualizer, + *, + gripper: SweepVolumeGripperConfig, + grasp_frame_to_tcp: RigidTransform, + max_candidates: int = DEFAULT_MAX_CANDIDATES, + cloud_loader: CloudLoader = load_demo_clouds, +) -> GraspVisualizationDemoResult: + """Publish banana object-cloud and ranked gripper-wireframe layers.""" + if max_candidates <= 0: + raise ValueError("max_candidates must be positive") + _, object_cloud = cloud_loader() + proposals = proposer.propose_grasps(object_cloud) + _validate(proposals, object_cloud) + ranked = sorted(proposals.candidates, key=lambda candidate: candidate.score, reverse=True) + displayed = ranked[:max_candidates] + + points, colors = object_cloud.as_numpy() + cloud_layer = VisualizationLayer( + OBJECT_CLOUD_LAYER_ID, + object_cloud.frame_id, + (PointCloudElement("object", points, colors),), + ) + wireframes = [] + for index, candidate in enumerate(displayed): + vertices, edges = gripper_wireframe_geometry( + candidate, + gripper, + grasp_frame_to_tcp, + ) + wireframes.append( + LineSetElement( + f"rank-{index + 1}", + vertices, + edges, + colors=_rank_color(index, len(displayed)), + line_width=2.5, + ) + ) + proposal_layer = VisualizationLayer( + PROPOSAL_LAYER_ID, + proposals.header.frame_id, + tuple(wireframes), + ) + visualizer.set_layer(cloud_layer) + visualizer.set_layer(proposal_layer) + return GraspVisualizationDemoResult( + candidate_count=len(ranked), + displayed_count=len(displayed), + frame_id=proposals.header.frame_id, + visualization_url=visualizer.get_visualization_url(), + ) + + +def wait_until_interrupted() -> None: + """Keep the interactive Viser process alive until Ctrl-C.""" + Event().wait() + + +def run_contributor_demo( + *, + max_candidates: int = DEFAULT_MAX_CANDIDATES, + config: GraspGenXConfig | None = None, + waiter: Waiter = wait_until_interrupted, +) -> GraspVisualizationDemoResult: + """Run real GraspGenX, publish layers, and own all interactive resources.""" + active_config = config if config is not None else make_xarm_graspgenx_config() + proposer = GraspGenXModule( + **active_config.model_dump(exclude={"rpc_transport", "tf_transport", "g"}) + ) + visualizer = ViserManipulationVisualizer(config=ViserVisualizationConfig(panel_enabled=False)) + proposer_started = False + try: + proposer.start() + proposer_started = True + result = run_demo( + proposer, + visualizer, + gripper=active_config.gripper, + grasp_frame_to_tcp=active_config.grasp_frame_to_tcp, + max_candidates=max_candidates, + ) + print( + "grasp-visualization-demo " + f"candidates={result.candidate_count} displayed={result.displayed_count} " + f"url={result.visualization_url}", + flush=True, + ) + waiter() + return result + finally: + visualizer.close() + if proposer_started: + proposer.stop() diff --git a/dimos/manipulation/demo_grasp_visualization/test_demo.py b/dimos/manipulation/demo_grasp_visualization/test_demo.py new file mode 100644 index 0000000000..69940e0aff --- /dev/null +++ b/dimos/manipulation/demo_grasp_visualization/test_demo.py @@ -0,0 +1,203 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hermetic tests for the interactive grasp visualization demo.""" + +from __future__ import annotations + +import numpy as np +import open3d as o3d +import pytest +from pytest_mock import MockerFixture + +from dimos.manipulation.demo_grasp_visualization import __main__ +import dimos.manipulation.demo_grasp_visualization.demo as demo_module +from dimos.manipulation.demo_grasp_visualization.demo import ( + GraspVisualizationDemoResult, + run_contributor_demo, + run_demo, +) +from dimos.manipulation.demo_graspgenx.demo import deployment_config +from dimos.manipulation.visualization.layers import ( + LineSetElement, + PointCloudElement, + VisualizationLayer, +) +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.std_msgs.Header import Header + + +class Proposer: + def __init__(self, candidates: list[GraspCandidate]) -> None: + self.candidates = candidates + self.calls = 0 + + def propose_grasps(self, cloud: PointCloud2) -> GraspCandidateArray: + self.calls += 1 + return GraspCandidateArray(Header(float(cloud.ts), cloud.frame_id), self.candidates) + + +class Visualizer: + def __init__(self) -> None: + self.layers: list[VisualizationLayer] = [] + self.closed = False + + def set_layer(self, layer: VisualizationLayer) -> None: + self.layers.append(layer) + + def get_visualization_url(self) -> str: + return "http://localhost:8095" + + def close(self) -> None: + self.closed = True + + +def candidate(x: float, score: float) -> GraspCandidate: + return GraspCandidate( + Pose( + { + "position": [x, 0.0, 0.2], + "orientation": [0.0, 0.0, 0.0, 1.0], + } + ), + score, + ) + + +def clouds() -> tuple[PointCloud2, PointCloud2]: + points = np.asarray([[0.0, 0.0, 0.0], [0.1, 0.2, 0.3]], dtype=np.float32) + colors = np.asarray([[1.0, 0.5, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32) + cloud = PointCloud2.from_numpy(points, frame_id="world", timestamp=42.0) + cloud.pointcloud.colors = o3d.utility.Vector3dVector(colors) + return cloud, cloud + + +def test_demo_publishes_colored_cloud_and_limited_ranked_wireframes() -> None: + proposer = Proposer([candidate(1.0, 0.9), candidate(2.0, 0.7), candidate(3.0, 0.5)]) + visualizer = Visualizer() + config = deployment_config() + + result = run_demo( + proposer, + visualizer, + gripper=config.gripper, + grasp_frame_to_tcp=config.grasp_frame_to_tcp, + max_candidates=2, + cloud_loader=clouds, + ) + + assert result == GraspVisualizationDemoResult(3, 2, "world", "http://localhost:8095") + assert proposer.calls == 1 + assert [layer.id for layer in visualizer.layers] == [ + "grasp/object-cloud", + "grasp/proposals", + ] + cloud_element = visualizer.layers[0].elements[0] + assert isinstance(cloud_element, PointCloudElement) + np.testing.assert_array_equal( + cloud_element.colors, + np.asarray([[255, 128, 0], [0, 255, 0]], dtype=np.uint8), + ) + proposals = visualizer.layers[1].elements + assert [element.id for element in proposals] == ["rank-1", "rank-2"] + assert all(isinstance(element, LineSetElement) for element in proposals) + np.testing.assert_array_equal(proposals[0].colors, [0, 220, 80]) + np.testing.assert_array_equal(proposals[1].colors, [255, 140, 0]) + + +def test_demo_applies_grasp_frame_to_tcp_to_wireframe() -> None: + proposer = Proposer([candidate(1.0, 0.9)]) + visualizer = Visualizer() + config = deployment_config() + transform = np.eye(4) + transform[0, 3] = 0.2 + + run_demo( + proposer, + visualizer, + gripper=config.gripper, + grasp_frame_to_tcp=tuple(tuple(float(value) for value in row) for row in transform), # type: ignore[arg-type] + cloud_loader=clouds, + ) + + proposal = visualizer.layers[1].elements[0] + assert isinstance(proposal, LineSetElement) + np.testing.assert_allclose(proposal.vertices[1], [0.8, 0.0, 0.2], atol=1e-6) + + +def test_demo_rejects_non_positive_candidate_limit() -> None: + with pytest.raises(ValueError, match="positive"): + run_demo( + Proposer([candidate(1.0, 0.9)]), + Visualizer(), + gripper=deployment_config().gripper, + grasp_frame_to_tcp=deployment_config().grasp_frame_to_tcp, + max_candidates=0, + cloud_loader=clouds, + ) + + +def test_contributor_closes_resources_when_waiter_fails( + mocker: MockerFixture, +) -> None: + config = deployment_config() + proposer = mocker.patch.object(demo_module, "GraspGenXModule").return_value + visualizer = mocker.patch.object(demo_module, "ViserManipulationVisualizer").return_value + mocker.patch.object( + demo_module, + "run_demo", + return_value=GraspVisualizationDemoResult(3, 2, "world", "http://localhost:8095"), + ) + + with pytest.raises(RuntimeError, match="stop waiting"): + run_contributor_demo( + config=config, + waiter=mocker.Mock(side_effect=RuntimeError("stop waiting")), + ) + + proposer.start.assert_called_once_with() + proposer.stop.assert_called_once_with() + visualizer.close.assert_called_once_with() + + +def test_entrypoint_passes_limit_and_handles_interrupt( + mocker: MockerFixture, capsys: pytest.CaptureFixture[str] +) -> None: + run = mocker.patch.object(__main__, "run_contributor_demo", side_effect=KeyboardInterrupt) + + assert __main__.main(["--max-candidates", "7"]) == 0 + run.assert_called_once_with(max_candidates=7) + assert "stopped" in capsys.readouterr().out + + +def test_entrypoint_reports_optional_dependency_hint( + mocker: MockerFixture, capsys: pytest.CaptureFixture[str] +) -> None: + error = RuntimeError("failed to start") + error.__cause__ = ModuleNotFoundError("graspgenx") + mocker.patch.object(__main__, "run_contributor_demo", side_effect=error) + + assert __main__.main([]) == 1 + output = capsys.readouterr().out + assert "uv sync --extra manipulation --extra graspgenx" in output + + +def test_entrypoint_rejects_non_positive_limit() -> None: + with pytest.raises(SystemExit) as raised: + __main__.main(["--max-candidates", "0"]) + + assert raised.value.code == 2 diff --git a/dimos/manipulation/demo_graspgenx/__main__.py b/dimos/manipulation/demo_graspgenx/__main__.py new file mode 100644 index 0000000000..68e59e9fdd --- /dev/null +++ b/dimos/manipulation/demo_graspgenx/__main__.py @@ -0,0 +1,48 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run GraspGenX once and save an annotated point-cloud PNG.""" + +import argparse +from collections.abc import Sequence +import os +from pathlib import Path + +from .demo import run_contributor_demo + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output", + type=Path, + default=Path(os.environ.get("DIMOS_GRASPGENX_OUTPUT", "graspgenx-ycb-demo.png")), + help="Destination PNG path.", + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parser().parse_args(argv) + result = run_contributor_demo(output_path=args.output) + print( + f"graspgenx-ycb-demo complete candidates={result.candidate_count} " + f"image={result.image_path}", + flush=True, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dimos/manipulation/demo_graspgenx/demo.py b/dimos/manipulation/demo_graspgenx/demo.py new file mode 100644 index 0000000000..17490485c2 --- /dev/null +++ b/dimos/manipulation/demo_graspgenx/demo.py @@ -0,0 +1,150 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""One-shot GraspGenX inference and static image generation.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import torch + +from dimos.manipulation.grasping.grasp_gen_spec import GraspGenSpec +from dimos.manipulation.grasping.grasp_gen_x import ( + GRASPGENX_MODEL_REPO, + GRASPGENX_MODEL_REVISION, + GraspGenXConfig, + GraspGenXModule, + SweepVolumeGripperConfig, +) +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + +from .fixture import load_demo_clouds +from .render import SweepVolumeLike, render_grasp_image + + +@dataclass(frozen=True) +class DemoResult: + image_path: Path + scene_points: int + object_points: int + candidate_count: int + best_score: float + frame: str + + +def deployment_config() -> GraspGenXConfig: + """Build the fixed sweep-volume deployment without downloading the checkpoint.""" + return GraspGenXConfig( + gripper=SweepVolumeGripperConfig( + extents_open=(0.08, 0.045, 0.04), + offset_open=(0.0, 0.0, 0.135), + extents_half_open=(0.04, 0.045, 0.035), + offset_half_open=(0.0, 0.0, 0.118), + fingertip_depth=0.15, + family="revolute_3f", + ), + max_candidates=100, + ) + + +def _cuda_context() -> dict[str, object]: + available = bool(torch.cuda.is_available()) + return {"available": available, "device": torch.cuda.get_device_name(0) if available else "cpu"} + + +def _validate(result: GraspCandidateArray, object_cloud: PointCloud2) -> None: + if not result.candidates: + raise ValueError("proposer returned no grasp candidates") + if result.header.frame_id != "world": + raise ValueError("proposer returned a result outside world frame") + if ( + result.header.frame_id != object_cloud.frame_id + or result.header.timestamp != object_cloud.ts + ): + raise ValueError("proposer changed the object point-cloud frame or timestamp") + + scores = np.asarray([candidate.score for candidate in result.candidates], dtype=float) + if not np.all(np.isfinite(scores)) or np.any(scores[:-1] < scores[1:]): + raise ValueError("grasp scores must be finite and descending") + for candidate in result.candidates: + p, q = candidate.pose.position, candidate.pose.orientation + values = np.asarray([p.x, p.y, p.z, q.x, q.y, q.z, q.w], dtype=float) + if not np.all(np.isfinite(values)): + raise ValueError("proposer returned a non-finite TCP pose") + + +def run_demo( + proposer: GraspGenSpec, + output_path: Path, + *, + gripper: SweepVolumeLike, + renderer: Callable[ + [Path, PointCloud2, PointCloud2, GraspCandidateArray, SweepVolumeLike], Path + ] = render_grasp_image, +) -> DemoResult: + """Load one scene, run inference once, render one PNG, and return.""" + if not callable(getattr(proposer, "propose_grasps", None)): + raise TypeError("proposer must implement GraspGenSpec.propose_grasps") + + scene, object_cloud = load_demo_clouds() + result = proposer.propose_grasps(object_cloud) + _validate(result, object_cloud) + image_path = renderer(output_path, scene, object_cloud, result, gripper) + best_score = float(result.candidates[0].score) + print( + "graspgenx-ycb-demo " + f"scene_points={len(scene)} object_points={len(object_cloud)} " + f"candidates={len(result)} best_score={best_score:.6f} " + f"image={image_path}", + flush=True, + ) + return DemoResult( + image_path=image_path, + scene_points=len(scene), + object_points=len(object_cloud), + candidate_count=len(result), + best_score=best_score, + frame=result.header.frame_id, + ) + + +def run_contributor_demo( + *, + output_path: Path, + config: GraspGenXConfig | None = None, +) -> DemoResult: + """Run the real adapter once and save its annotated point-cloud image.""" + active_config = config if config is not None else deployment_config() + cuda = _cuda_context() + print( + "graspgenx-ycb-demo " + f"checkpoint=hf://{GRASPGENX_MODEL_REPO}@{GRASPGENX_MODEL_REVISION} " + f"cuda={cuda['available']} device={cuda['device']}", + flush=True, + ) + + module_args = active_config.model_dump( + exclude={"rpc_transport", "tf_transport", "g"}, + ) + adapter = GraspGenXModule(**module_args) + try: + adapter.start() + return run_demo(adapter, output_path, gripper=active_config.gripper) + finally: + adapter.stop() diff --git a/dimos/manipulation/demo_graspgenx/fixture.py b/dimos/manipulation/demo_graspgenx/fixture.py new file mode 100644 index 0000000000..c57266b765 --- /dev/null +++ b/dimos/manipulation/demo_graspgenx/fixture.py @@ -0,0 +1,124 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Deterministic YCB scene input loaded through the repository data system.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +import re +from typing import TypedDict + +import numpy as np + +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.utils.data import get_data + +DATASET = "graspgenx_ycb_banana_scene" +SCENE_FILE = "scene.npz" +METADATA_FILE = "scene.json" +BANANA_LABEL = 0 + + +class SceneMetadata(TypedDict, total=False): + counts: dict[str, int] + labels: dict[str, dict[str, str]] + frame: str + timestamp: float + final_sha256: str + source: dict[str, object] + + +def _dataset_dir() -> Path: + return get_data(DATASET) + + +def _validate_record( + path: Path, + points: np.ndarray, + labels: np.ndarray, + timestamp: float, + metadata: SceneMetadata, +) -> None: + expected_sha = metadata.get("final_sha256") + if not isinstance(expected_sha, str) or not re.fullmatch(r"[0-9a-f]{64}", expected_sha): + raise ValueError("scene fixture metadata must contain a valid final SHA-256") + if hashlib.sha256(path.read_bytes()).hexdigest() != expected_sha: + raise ValueError("scene fixture bytes do not match provenance SHA-256") + + expected_counts = {"banana": 3500, "table": 256, "distractor": 48, "total": 3804} + if metadata.get("counts") != expected_counts or expected_counts["total"] != len(points): + raise ValueError("scene fixture counts do not match bytes") + if (metadata.get("labels") or {}).get("encoding") != { + "0": "banana", + "1": "table", + "2": "distractor", + }: + raise ValueError("scene fixture label encoding metadata is invalid") + + actual_counts = dict(zip(*np.unique(labels, return_counts=True), strict=True)) + if {str(label): int(count) for label, count in actual_counts.items()} != { + "0": 3500, + "1": 256, + "2": 48, + }: + raise ValueError("scene fixture label values do not match metadata") + if metadata.get("frame") != "world" or timestamp != float( + metadata.get("timestamp", float("nan")) + ): + raise ValueError("scene fixture frame/timestamp metadata is invalid") + + source = metadata.get("source") + if not isinstance(source, dict) or source.get("path") != ( + "assets/sample_data/object_mesh/banana.obj" + ): + raise ValueError("scene fixture source OBJ provenance is incomplete") + source_sha = source.get("source_obj_sha256") + if not isinstance(source_sha, str) or not re.fullmatch(r"[0-9a-f]{64}", source_sha): + raise ValueError("scene fixture source OBJ must have its 64-character SHA-256") + + +def load_scene_record( + dataset_dir: Path | None = None, +) -> tuple[np.ndarray, np.ndarray, SceneMetadata]: + """Load stored XYZ points, semantic labels, and provenance metadata.""" + root = dataset_dir if dataset_dir is not None else _dataset_dir() + scene_path = root / SCENE_FILE + metadata_path = root / METADATA_FILE + with np.load(scene_path, allow_pickle=False) as data: + points = np.asarray(data["points"], dtype=np.float32) + labels = np.asarray(data["labels"], dtype=np.uint8) + timestamp = float(np.asarray(data["timestamp"]).item()) + if points.shape != (len(labels), 3) or not np.all(np.isfinite(points)): + raise ValueError("scene fixture has invalid points or labels") + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + _validate_record(scene_path, points, labels, timestamp, metadata) + return points, labels, metadata + + +def load_demo_clouds( + dataset_dir: Path | None = None, +) -> tuple[PointCloud2, PointCloud2]: + """Return the complete scene and its labeled banana Object Point Cloud.""" + points, labels, metadata = load_scene_record(dataset_dir) + timestamp = float(metadata["timestamp"]) + scene = PointCloud2.from_numpy(points, frame_id="world", timestamp=timestamp) + object_cloud = PointCloud2.from_numpy( + points[labels == BANANA_LABEL], + frame_id="world", + timestamp=timestamp, + ) + return scene, object_cloud diff --git a/dimos/manipulation/demo_graspgenx/render.py b/dimos/manipulation/demo_graspgenx/render.py new file mode 100644 index 0000000000..dde6d58029 --- /dev/null +++ b/dimos/manipulation/demo_graspgenx/render.py @@ -0,0 +1,250 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Render one static point-cloud image with grasp proposal annotations.""" + +from __future__ import annotations + +from collections.abc import Sequence +import os +from pathlib import Path +import tempfile +from typing import Any, Protocol + +from matplotlib.backends.backend_agg import FigureCanvasAgg +from matplotlib.figure import Figure +from matplotlib.lines import Line2D +import numpy as np + +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + +DISPLAYED_CANDIDATES = 5 + + +class SweepVolumeLike(Protocol): + extents_open: tuple[float, float, float] + offset_open: tuple[float, float, float] + extents_half_open: tuple[float, float, float] + offset_half_open: tuple[float, float, float] + + +def _rotation(q: Any) -> np.ndarray: + x, y, z, w = (float(q.x), float(q.y), float(q.z), float(q.w)) + return np.asarray( + [ + [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)], + [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)], + [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)], + ], + dtype=float, + ) + + +def _score_colors(scores: np.ndarray) -> np.ndarray: + if not len(scores): + return np.empty((0, 3), dtype=float) + low, high = float(scores.min()), float(scores.max()) + normalized = ( + np.linspace(1.0, 0.0, len(scores)) if high == low else (scores - low) / (high - low) + ) + from matplotlib import colormaps + + return np.asarray(colormaps["viridis"](normalized))[:, :3] + + +def _fork_strips_local(gripper: SweepVolumeLike) -> tuple[np.ndarray, ...]: + open_extents = np.asarray(gripper.extents_open, dtype=float) + half_extents = np.asarray(gripper.extents_half_open, dtype=float) + open_offset = np.asarray(gripper.offset_open, dtype=float) + half_offset = np.asarray(gripper.offset_half_open, dtype=float) + + rear_center = np.asarray([half_offset[0], 0.0, half_offset[2] - half_extents[2] / 2.0]) + mouth_center = np.asarray([open_offset[0], 0.0, open_offset[2] + open_extents[2] / 2.0]) + if mouth_center[2] <= rear_center[2]: + raise ValueError("configured sweep profiles must open toward increasing local +Z") + + rear_half_width = max(float(half_extents[0]) / 2.0, 1e-6) + mouth_half_width = max(float(open_extents[0]) / 2.0, rear_half_width) + rear_left = rear_center + np.asarray([-rear_half_width, 0.0, 0.0]) + rear_right = rear_center + np.asarray([rear_half_width, 0.0, 0.0]) + mouth_left = mouth_center + np.asarray([-mouth_half_width, 0.0, 0.0]) + mouth_right = mouth_center + np.asarray([mouth_half_width, 0.0, 0.0]) + return ( + np.asarray([rear_center, [0.0, 0.0, 0.0]]), + np.asarray([rear_left, rear_right]), + np.asarray([rear_left, mouth_left]), + np.asarray([rear_right, mouth_right]), + ) + + +def _fork_strips(candidate: Any, gripper: SweepVolumeLike) -> tuple[np.ndarray, ...]: + return gripper_wireframe_strips(candidate, gripper, np.eye(4, dtype=float)) + + +def gripper_wireframe_strips( + candidate: Any, + gripper: SweepVolumeLike, + grasp_frame_to_tcp: Sequence[Sequence[float]] | np.ndarray, +) -> tuple[np.ndarray, ...]: + """Convert one TCP proposal into world-frame sweep-volume wireframe strips.""" + p, q = candidate.pose.position, candidate.pose.orientation + world_to_tcp = np.eye(4, dtype=float) + world_to_tcp[:3, :3] = _rotation(q) + world_to_tcp[:3, 3] = np.asarray([p.x, p.y, p.z], dtype=float) + grasp_to_tcp = np.asarray(grasp_frame_to_tcp, dtype=float) + if grasp_to_tcp.shape != (4, 4): + raise ValueError("grasp_frame_to_tcp must have shape (4, 4)") + world_to_grasp = world_to_tcp @ np.linalg.inv(grasp_to_tcp) + rotation = world_to_grasp[:3, :3] + translation = world_to_grasp[:3, 3] + return tuple((rotation @ strip.T).T + translation for strip in _fork_strips_local(gripper)) + + +def gripper_wireframe_geometry( + candidate: Any, + gripper: SweepVolumeLike, + grasp_frame_to_tcp: Sequence[Sequence[float]] | np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """Return indexed world-frame vertices and edges for one grasp proposal.""" + strips = gripper_wireframe_strips(candidate, gripper, grasp_frame_to_tcp) + vertices = np.vstack(strips).astype(np.float32) + edges = np.arange(len(vertices), dtype=np.int32).reshape((-1, 2)) + return vertices, edges + + +def _set_equal_axes(axis: Any, points: np.ndarray) -> None: + minimum = points.min(axis=0) + maximum = points.max(axis=0) + center = (minimum + maximum) / 2.0 + radius = max(float(np.max(maximum - minimum)) / 2.0, 0.05) + axis.set_xlim(center[0] - radius, center[0] + radius) + axis.set_ylim(center[1] - radius, center[1] + radius) + axis.set_zlim(center[2] - radius, center[2] + radius) + + +def render_grasp_image( + output_path: Path, + scene: PointCloud2, + object_cloud: PointCloud2, + candidates: GraspCandidateArray, + gripper: SweepVolumeLike, + *, + ranks: Sequence[int] | None = None, + title: str | None = None, +) -> Path: + """Write a single PNG containing the scene, target cloud, and top proposals.""" + scene_points, _ = scene.as_numpy() + object_points, _ = object_cloud.as_numpy() + top = list(candidates.candidates[:DISPLAYED_CANDIDATES]) + displayed_ranks = list(ranks) if ranks is not None else list(range(1, len(top) + 1)) + if len(displayed_ranks) != len(top): + raise ValueError("one displayed rank is required for each rendered candidate") + scores = np.asarray([candidate.score for candidate in top], dtype=float) + colors = _score_colors(scores) + + figure = Figure(figsize=(10, 8), dpi=150, facecolor="#10171c") + FigureCanvasAgg(figure) + axis = figure.add_subplot(111, projection="3d", facecolor="#10171c") + for pane in (axis.xaxis.pane, axis.yaxis.pane, axis.zaxis.pane): + pane.set_facecolor("#10171c") + pane.set_edgecolor("#58707b") + axis.scatter( + scene_points[:, 0], + scene_points[:, 1], + scene_points[:, 2], + s=1.0, + c="#6f8793", + alpha=0.18, + depthshade=False, + ) + axis.scatter( + object_points[:, 0], + object_points[:, 1], + object_points[:, 2], + s=2.5, + c="#f5bf1f", + alpha=0.9, + depthshade=False, + ) + + annotation_points = [scene_points] + legend_handles: list[Line2D] = [] + for rank, candidate, color in zip(displayed_ranks, top, colors, strict=True): + strips = _fork_strips(candidate, gripper) + annotation_points.extend(strips) + for strip in strips: + axis.plot( + strip[:, 0], + strip[:, 1], + strip[:, 2], + color=color, + linewidth=2.5, + ) + legend_handles.append( + Line2D( + [0], + [0], + color=color, + linewidth=3, + label=f"#{rank} score {candidate.score:.3f}", + ) + ) + + if legend_handles: + legend = axis.legend( + handles=legend_handles, + loc="upper right", + frameon=True, + facecolor="#10171c", + edgecolor="#58707b", + labelcolor="white", + ) + legend.get_frame().set_alpha(0.9) + + all_points = np.vstack(annotation_points) + _set_equal_axes(axis, all_points) + axis.view_init(elev=24, azim=-58) + axis.set_xlabel("X (m)", color="#b8c7ce") + axis.set_ylabel("Y (m)", color="#b8c7ce") + axis.set_zlabel("Z (m)", color="#b8c7ce") + axis.tick_params(colors="#90a4ae") + figure.suptitle( + title or f"GraspGenX proposals — {len(candidates)} candidates, top {len(top)} shown", + color="white", + fontsize=16, + y=0.98, + ) + figure.tight_layout(rect=(0.0, 0.0, 1.0, 0.95)) + + final_path = output_path.expanduser().resolve() + final_path.parent.mkdir(parents=True, exist_ok=True) + temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile( + dir=final_path.parent, + prefix=f".{final_path.name}.", + suffix=".png", + delete=False, + ) as handle: + temporary = Path(handle.name) + figure.savefig(temporary, format="png", facecolor=figure.get_facecolor()) + temporary.chmod(0o644) + os.replace(temporary, final_path) + temporary = None + finally: + figure.clear() + if temporary is not None: + temporary.unlink(missing_ok=True) + return final_path diff --git a/dimos/manipulation/demo_graspgenx/test_demo.py b/dimos/manipulation/demo_graspgenx/test_demo.py new file mode 100644 index 0000000000..31ba8f5b0f --- /dev/null +++ b/dimos/manipulation/demo_graspgenx/test_demo.py @@ -0,0 +1,160 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hermetic coverage for the one-shot GraspGenX image demo.""" + +from pathlib import Path +import stat + +import numpy as np +import pytest +from pytest_mock import MockerFixture + +import dimos.manipulation.demo_graspgenx.demo as demo +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.std_msgs.Header import Header + +from . import __main__ +from .demo import DemoResult, deployment_config, run_contributor_demo, run_demo +from .fixture import load_demo_clouds, load_scene_record + + +class FakeGraspProposer: + def __init__(self) -> None: + self.calls = 0 + + def propose_grasps(self, object_pointcloud: PointCloud2) -> GraspCandidateArray: + self.calls += 1 + center = object_pointcloud.pointcloud.get_center() + candidates = [ + GraspCandidate( + Pose( + { + "position": [float(center[0]), float(center[1]), float(center[2] + dz)], + "orientation": [0.0, 0.0, 0.0, 1.0], + } + ), + score, + ) + for dz, score in ((0.16, 0.91), (0.15, 0.73), (0.14, 0.52)) + ] + return GraspCandidateArray( + Header(float(object_pointcloud.ts), "world"), + candidates, + ) + + +def test_standard_data_fixture_is_deterministic() -> None: + scene, object_cloud = load_demo_clouds() + scene_again, object_again = load_demo_clouds() + points, labels, metadata = load_scene_record() + + assert len(scene) == 3804 + assert len(object_cloud) == 3500 + assert metadata["counts"] == { + "banana": 3500, + "table": 256, + "distractor": 48, + "total": 3804, + } + assert metadata["timestamp"] == 1700000000.25 + assert scene.frame_id == object_cloud.frame_id == "world" + assert scene.ts == object_cloud.ts == metadata["timestamp"] + assert np.count_nonzero(labels == 0) == 3500 + np.testing.assert_array_equal(points, scene_again.points_f32()) + np.testing.assert_array_equal(object_cloud.points_f32(), object_again.points_f32()) + + +def test_demo_runs_inference_once_and_writes_png(tmp_path: Path) -> None: + proposer = FakeGraspProposer() + output = tmp_path / "graspgenx.png" + result = run_demo(proposer, output, gripper=deployment_config().gripper) + + assert proposer.calls == 1 + assert result == DemoResult(output.resolve(), 3804, 3500, 3, 0.91, "world") + assert output.read_bytes().startswith(b"\x89PNG\r\n\x1a\n") + assert output.stat().st_size > 10_000 + assert stat.S_IMODE(output.stat().st_mode) == 0o644 + + +def test_contributor_stops_adapter_after_render_failure( + mocker: MockerFixture, tmp_path: Path +) -> None: + adapter = mocker.patch.object(demo, "GraspGenXModule").return_value + mocker.patch.object(demo, "run_demo", side_effect=RuntimeError("render")) + + with pytest.raises(RuntimeError, match="render"): + run_contributor_demo(output_path=tmp_path / "failure.png") + + adapter.start.assert_called_once_with() + adapter.stop.assert_called_once_with() + + +def test_one_config_drives_adapter_and_wireframe(mocker: MockerFixture, tmp_path: Path) -> None: + base = deployment_config() + config = base.model_copy( + update={ + "gripper": base.gripper.model_copy(update={"offset_open": (0.0, 0.0, 0.2)}), + } + ) + adapter_class = mocker.patch.object(demo, "GraspGenXModule") + run = mocker.patch.object( + demo, + "run_demo", + return_value=DemoResult(tmp_path / "config.png", 1, 1, 1, 1.0, "world"), + ) + + run_contributor_demo( + output_path=tmp_path / "config.png", + config=config, + ) + + assert adapter_class.call_args.kwargs["gripper"] == config.gripper.model_dump() + run.assert_called_once_with( + adapter_class.return_value, + tmp_path / "config.png", + gripper=config.gripper, + ) + + +def test_python_module_entrypoint_is_direct_and_user_visible( + mocker: MockerFixture, tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + output = tmp_path / "entrypoint.png" + run = mocker.patch.object( + __main__, + "run_contributor_demo", + return_value=DemoResult(output, 1, 1, 3, 0.9, "world"), + ) + + assert __main__.main(["--output", str(output)]) == 0 + run.assert_called_once_with(output_path=output) + assert f"candidates=3 image={output}" in capsys.readouterr().out + + +def test_empty_result_fails_explicitly(tmp_path: Path) -> None: + class Empty: + def propose_grasps(self, cloud: object) -> GraspCandidateArray: + _, object_cloud = load_demo_clouds() + return GraspCandidateArray(Header(object_cloud.ts, "world"), []) + + with pytest.raises(ValueError, match="no grasp candidates"): + run_demo( + Empty(), # type: ignore[arg-type] + tmp_path / "empty.png", + gripper=deployment_config().gripper, + ) diff --git a/dimos/manipulation/demo_graspgenx/test_render.py b/dimos/manipulation/demo_graspgenx/test_render.py new file mode 100644 index 0000000000..99e82dcefd --- /dev/null +++ b/dimos/manipulation/demo_graspgenx/test_render.py @@ -0,0 +1,97 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Focused checks for static grasp proposal annotations.""" + +from types import SimpleNamespace + +import numpy as np + +from .render import ( + _fork_strips, + _fork_strips_local, + _score_colors, + gripper_wireframe_geometry, +) + + +def _candidate( + position: tuple[float, float, float], + orientation: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0), +) -> SimpleNamespace: + return SimpleNamespace( + pose=SimpleNamespace( + position=SimpleNamespace(x=position[0], y=position[1], z=position[2]), + orientation=SimpleNamespace( + x=orientation[0], + y=orientation[1], + z=orientation[2], + w=orientation[3], + ), + ) + ) + + +def _gripper() -> SimpleNamespace: + return SimpleNamespace( + extents_open=(0.2, 0.3, 0.4), + offset_open=(0.1, 0.0, 0.3), + extents_half_open=(0.1, 0.15, 0.2), + offset_half_open=(0.0, 0.2, 0.1), + ) + + +def test_wireframe_opens_along_local_positive_z() -> None: + strips = _fork_strips_local(_gripper()) + rear_bridge, left_rail, right_rail = strips[1:] + rear_width = np.ptp(rear_bridge[:, 0]) + mouth_width = abs(right_rail[-1, 0] - left_rail[-1, 0]) + + assert all(np.allclose(strip[:, 1], 0.0) for strip in strips) + assert mouth_width > rear_width + assert left_rail[-1, 2] > left_rail[0, 2] + assert right_rail[-1, 2] > right_rail[0, 2] + + +def test_wireframe_uses_candidate_full_rigid_pose() -> None: + first = _fork_strips(_candidate((1.0, 2.0, 3.0)), _gripper()) + second = _fork_strips( + _candidate((4.0, 5.0, 6.0), (2**-0.5, 0.0, 0.0, 2**-0.5)), + _gripper(), + ) + rotation_x_90 = np.asarray([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]]) + for first_strip, second_strip in zip(first, second, strict=True): + local = first_strip - [1.0, 2.0, 3.0] + expected = (rotation_x_90 @ local.T).T + np.asarray([4.0, 5.0, 6.0]) + np.testing.assert_allclose(second_strip, expected, rtol=0.0, atol=1e-6) + + +def test_score_colors_are_relative_and_monotonic() -> None: + colors = _score_colors(np.asarray([0.1, 0.5, 0.9])) + assert len({tuple(color) for color in colors}) == 3 + assert np.all(np.diff(colors.mean(axis=1)) > 0) + + +def test_wireframe_converts_tcp_pose_back_to_grasp_frame() -> None: + grasp_frame_to_tcp = np.eye(4) + grasp_frame_to_tcp[0, 3] = 0.2 + + vertices, edges = gripper_wireframe_geometry( + _candidate((1.0, 2.0, 3.0)), + _gripper(), + grasp_frame_to_tcp, + ) + + np.testing.assert_allclose(vertices[1], [0.8, 2.0, 3.0], atol=1e-6) + np.testing.assert_array_equal(edges, [[0, 1], [2, 3], [4, 5], [6, 7]]) diff --git a/dimos/manipulation/grasping/grasp_gen_spec.py b/dimos/manipulation/grasping/grasp_gen_spec.py index 37c81c85bc..07e3882a9a 100644 --- a/dimos/manipulation/grasping/grasp_gen_spec.py +++ b/dimos/manipulation/grasping/grasp_gen_spec.py @@ -14,14 +14,10 @@ from typing import Protocol -from dimos.msgs.geometry_msgs.PoseArray import PoseArray +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.spec.utils import Spec class GraspGenSpec(Spec, Protocol): - def generate_grasps( - self, - pointcloud: PointCloud2, - scene_pointcloud: PointCloud2 | None = None, - ) -> PoseArray | None: ... + def propose_grasps(self, object_pointcloud: PointCloud2) -> GraspCandidateArray: ... diff --git a/dimos/manipulation/grasping/grasp_gen_x.py b/dimos/manipulation/grasping/grasp_gen_x.py new file mode 100644 index 0000000000..a383e14d12 --- /dev/null +++ b/dimos/manipulation/grasping/grasp_gen_x.py @@ -0,0 +1,192 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Import-safe DimOS adapter for GraspGenX grasp proposals.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Annotated, Any, Literal, TypeAlias + +import numpy as np +from pydantic import Field, FiniteFloat, field_validator + +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.manipulation.grasping.grasp_gen_spec import GraspGenSpec +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.std_msgs.Header import Header +from dimos.protocol.service.spec import BaseConfig + +if TYPE_CHECKING: + from dimos.manipulation.grasping.grasp_gen_x_runtime import GraspGenXRuntime + +GRASPGENX_MODEL_REPO = "adithyamurali/GraspGenXModel" +GRASPGENX_MODEL_REVISION = "7c834043c11a11417e31d6d5ea9355801e40a2c1" +GRASPGENX_MODEL_VERSION = "release" + +BoundedExtent = Annotated[FiniteFloat, Field(gt=0.0, le=0.5)] +BoundedOffset = Annotated[FiniteFloat, Field(ge=-0.5, le=0.5)] +PositiveCount = Annotated[int, Field(gt=0, strict=True)] +SweepExtents: TypeAlias = tuple[BoundedExtent, BoundedExtent, BoundedExtent] +SweepOffset: TypeAlias = tuple[BoundedOffset, BoundedOffset, BoundedOffset] +Vector4: TypeAlias = tuple[FiniteFloat, FiniteFloat, FiniteFloat, FiniteFloat] +RigidTransform: TypeAlias = tuple[Vector4, Vector4, Vector4, Vector4] +GripperFamily: TypeAlias = Literal["parallel_2f", "revolute_2f", "revolute_3f"] + +IDENTITY_TRANSFORM: RigidTransform = ( + (1.0, 0.0, 0.0, 0.0), + (0.0, 1.0, 0.0, 0.0), + (0.0, 0.0, 1.0, 0.0), + (0.0, 0.0, 0.0, 1.0), +) + + +class SweepVolumeGripperConfig(BaseConfig): + """Axis-aligned open and half-open sweep-volume description.""" + + extents_open: SweepExtents + offset_open: SweepOffset + extents_half_open: SweepExtents + offset_half_open: SweepOffset + fingertip_depth: BoundedExtent + family: GripperFamily = "parallel_2f" + + +class GraspGenXConfig(ModuleConfig): + """GraspGenX deployment settings, serializable by DimOS blueprints.""" + + gripper: SweepVolumeGripperConfig + grasp_frame_to_tcp: RigidTransform = IDENTITY_TRANSFORM + max_candidates: PositiveCount = 100 + load_on_start: bool = True + + # Relational matrix properties cannot be expressed through scalar Field constraints. + @field_validator("grasp_frame_to_tcp") + @classmethod + def _validate_rigid_transform(cls, value: RigidTransform) -> RigidTransform: + matrix = np.asarray(value, dtype=float) + if not np.allclose(matrix[3], [0.0, 0.0, 0.0, 1.0], atol=1e-7): + raise ValueError("grasp_frame_to_tcp must be homogeneous") + rotation = matrix[:3, :3] + if not np.allclose(rotation.T @ rotation, np.eye(3), atol=1e-6) or not np.isclose( + np.linalg.det(rotation), 1.0, atol=1e-6 + ): + raise ValueError("grasp_frame_to_tcp rotation must be orthonormal with determinant +1") + return value + + +class GraspGenXError(RuntimeError): + """Base error for model loading and inference failures.""" + + +def _create_runtime(config: GraspGenXConfig) -> GraspGenXRuntime: + # This import is the intentional first-use boundary for the optional GPU runtime. + from dimos.manipulation.grasping.grasp_gen_x_runtime import GraspGenXRuntime + + return GraspGenXRuntime(config) + + +class GraspGenXModule(Module, GraspGenSpec): + """Direct adapter whose optional runtime is loaded synchronously by ``start``.""" + + dedicated_worker = True + config: GraspGenXConfig # type: ignore[assignment] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._runtime: GraspGenXRuntime | None = None + + @rpc + def start(self) -> None: + super().start() + if self.config.load_on_start: + self._ensure_runtime() + + def _ensure_runtime(self) -> None: + if self._runtime is not None: + return + try: + self._runtime = _create_runtime(self.config) + except Exception as exc: + raise GraspGenXError("failed to initialize GraspGenX") from exc + + @rpc + def stop(self) -> None: + if self._runtime is not None: + self._runtime.stop() + self._runtime = None + super().stop() + + @rpc + def propose_grasps(self, object_pointcloud: PointCloud2) -> GraspCandidateArray: + if object_pointcloud.ts is None: + raise ValueError("object pointcloud must have a timestamp") + if not object_pointcloud.frame_id: + raise ValueError("object pointcloud frame_id must not be empty") + + points = object_pointcloud.points_f32() + if points.ndim != 2 or points.shape[1] != 3 or len(points) == 0: + raise ValueError("object pointcloud must contain at least one XYZ point") + if not np.all(np.isfinite(points)): + raise ValueError("object pointcloud XYZ values must be finite floats in metres") + + self._ensure_runtime() + assert self._runtime is not None + try: + poses, scores = self._runtime.infer(points) + except Exception as exc: + raise GraspGenXError("GraspGenX inference failed") from exc + scores = scores.reshape(-1) + + if poses.size == 0 and scores.size == 0: + return GraspCandidateArray( + Header(float(object_pointcloud.ts), object_pointcloud.frame_id), + [], + ) + if poses.shape != (len(scores), 4, 4): + raise ValueError("backend poses must have shape (N, 4, 4)") + if not np.all(np.isfinite(poses)) or not np.all(np.isfinite(scores)): + raise ValueError("backend returned non-finite poses or scores") + if not np.allclose(poses[:, 3, :], np.array([0.0, 0.0, 0.0, 1.0]), atol=1e-7): + raise ValueError("backend poses must be homogeneous") + rotations = poses[:, :3, :3] + if not np.allclose(np.einsum("nij,nkj->nik", rotations, rotations), np.eye(3), atol=1e-5): + raise ValueError("backend poses must have orthonormal rotations") + if not np.allclose(np.linalg.det(rotations), 1.0, atol=1e-5): + raise ValueError("backend poses must have proper rotations") + + tcp_poses = poses @ np.asarray(self.config.grasp_frame_to_tcp) + order = np.argsort(-scores, kind="stable")[: self.config.max_candidates] + candidates = [ + GraspCandidate(self._pose_from_matrix(tcp_poses[index]), float(scores[index])) + for index in order + ] + return GraspCandidateArray( + Header(float(object_pointcloud.ts), object_pointcloud.frame_id), + candidates, + ) + + @staticmethod + def _pose_from_matrix(matrix: np.ndarray) -> Pose: + return Pose( + { + "position": Vector3(matrix[:3, 3]), + "orientation": Quaternion.from_rotation_matrix(matrix[:3, :3]), + } + ) diff --git a/dimos/manipulation/grasping/grasp_gen_x_runtime.py b/dimos/manipulation/grasping/grasp_gen_x_runtime.py new file mode 100644 index 0000000000..b7d89d6d52 --- /dev/null +++ b/dimos/manipulation/grasping/grasp_gen_x_runtime.py @@ -0,0 +1,96 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""First-use in-process GraspGenX runtime.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import numpy as np + +from dimos.manipulation.grasping.grasp_gen_x import ( + GRASPGENX_MODEL_REPO, + GRASPGENX_MODEL_REVISION, + GRASPGENX_MODEL_VERSION, + GraspGenXConfig, +) + +_GRIPPER_TYPES = { + "parallel_2f": 0, + "revolute_2f": 1, + "revolute_3f": 2, +} + + +class GraspGenXRuntime: + """Loaded GraspGenX sampler and exact tensor conversion boundary.""" + + def __init__(self, config: GraspGenXConfig) -> None: + from huggingface_hub import snapshot_download + + snapshot_root = Path( + snapshot_download( + repo_id=GRASPGENX_MODEL_REPO, + revision=GRASPGENX_MODEL_REVISION, + allow_patterns=[ + f"{GRASPGENX_MODEL_VERSION}/gen/*", + f"{GRASPGENX_MODEL_VERSION}/dis/*", + ], + ) + ).resolve() + checkpoint_root = snapshot_root / GRASPGENX_MODEL_VERSION + gen_dir = checkpoint_root / "gen" + dis_dir = checkpoint_root / "dis" + if not gen_dir.is_dir() or not dis_dir.is_dir(): + raise FileNotFoundError( + f"GraspGenX checkpoint must contain release/gen and release/dis: {snapshot_root}" + ) + + # Sweep-volume grippers do not use named gripper assets, so avoid the + # upstream package's on-import Git clone. + os.environ["GRASPGENX_CHECKPOINT_DIR"] = str(snapshot_root) + os.environ["GRASPGENX_GRIPPER_CFG_DIR"] = str(snapshot_root) + from graspgenx.grasp_server import SWEEP_VOLUME_ONLY_BACKBONES, GraspGenXSampler + from graspgenx.utils.checkpoint_io import load_model_cfg + from graspgenx.x_grippers import make_sweep_volume_gripper_info + + model_config = load_model_cfg(gen_dir, dis_dir, gen_pth=None, dis_pth=None) + for component in ("diffusion", "discriminator"): + backbone = getattr(model_config, component).gripper_backbone + if backbone not in SWEEP_VOLUME_ONLY_BACKBONES: + raise ValueError( + f"GraspGenX {component}.gripper_backbone={backbone!r} " + "requires an asset-backed gripper" + ) + gripper_info = make_sweep_volume_gripper_info( + extents_open=config.gripper.extents_open, + offset_open=config.gripper.offset_open, + extents_mid=config.gripper.extents_half_open, + offset_mid=config.gripper.offset_half_open, + gripper_type=_GRIPPER_TYPES[config.gripper.family], + fingertip_depth=config.gripper.fingertip_depth, + ) + self._sampler_type = GraspGenXSampler + self._sampler = GraspGenXSampler(model_config, gripper_info=gripper_info) + + def infer(self, points: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Run inference and copy the known torch tensors to CPU NumPy arrays.""" + poses, scores = self._sampler_type.run_inference(points, self._sampler) + return poses.detach().cpu().numpy(), scores.detach().cpu().numpy() + + def stop(self) -> None: + """Release the sampler when the module stops.""" + del self._sampler diff --git a/dimos/manipulation/grasping/grasping.py b/dimos/manipulation/grasping/grasping.py deleted file mode 100644 index 5cefa526d2..0000000000 --- a/dimos/manipulation/grasping/grasping.py +++ /dev/null @@ -1,137 +0,0 @@ -# Copyright 2025-2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Grasping skill module - -Provides @skill interface for agents and orchestrates the grasp generation pipeline: -perception (get pointcloud) to graspgen (generate grasps in Docker) to output grasps -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -from dimos.agents.annotation import skill -from dimos.core.core import rpc -from dimos.core.module import Module -from dimos.core.stream import Out -from dimos.manipulation.grasping.grasp_gen_spec import GraspGenSpec -from dimos.msgs.geometry_msgs.PoseArray import PoseArray -from dimos.perception.experimental.object_scene_registration_spec import ObjectSceneRegistrationSpec -from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import quaternion_to_euler - -if TYPE_CHECKING: - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - -logger = setup_logger() - - -class GraspingModule(Module): - """Grasping skill and orchestrator module""" - - grasps: Out[PoseArray] - - _scene_registration: ObjectSceneRegistrationSpec - _grasp_gen: GraspGenSpec - - @rpc - def start(self) -> None: - super().start() - logger.info("GraspingModule started") - - @rpc - def stop(self) -> None: - super().stop() - logger.info("GraspingModule stopped") - - @skill - def generate_grasps( - self, - object_name: str = "object", - object_id: str | None = None, - filter_collisions: bool = True, - ) -> str: - """Generate grasp poses for the specified object. - - Args: - object_name: Name of the object to grasp (e.g. "coke can", "cup", "bottle"). - object_id: Optional unique object ID from perception. If provided, uses this - instead of object_name for lookup. - filter_collisions: Whether to filter grasps that collide with scene geometry. - - """ - # Get object pointcloud from perception - pc = self._get_object_pointcloud(object_name, object_id) - if pc is None: - msg = f"No pointcloud found for '{object_id or object_name}'" - logger.warning(msg) - return msg - - # Get scene pointcloud for collision filtering - scene_pc = None - if filter_collisions: - scene_pc = self._get_scene_pointcloud(exclude_object_id=object_id) - - # Call GraspGenModule (running in Docker) - try: - result = self._grasp_gen.generate_grasps(pc, scene_pc) - except Exception as e: - msg = f"Grasp generation failed: {e}" - logger.error(msg) - return msg - - if result is None or len(result.poses) == 0: - msg = f"No grasps generated for '{object_name}'" - logger.info(msg) - return msg - - self.grasps.publish(result) - logger.info(f"Generated {len(result.poses)} grasps for '{object_name}'") - - # Format result for agent/human - return self._format_grasp_result(result, object_name) - - def _get_object_pointcloud( - self, object_name: str, object_id: str | None = None - ) -> PointCloud2 | None: - """Fetch object pointcloud from perception.""" - try: - if object_id is not None: - return self._scene_registration.get_object_pointcloud_by_object_id(object_id) - - return self._scene_registration.get_object_pointcloud_by_name(object_name) - except Exception as e: - logger.error(f"Failed to get object pointcloud: {e}") - return None - - def _get_scene_pointcloud(self, exclude_object_id: str | None = None) -> PointCloud2 | None: - """Fetch scene pointcloud from perception for collision filtering.""" - try: - return self._scene_registration.get_full_scene_pointcloud( - exclude_object_id=exclude_object_id - ) - except Exception as e: - logger.debug(f"Could not get scene pointcloud: {e}") - return None - - def _format_grasp_result(self, grasps: PoseArray, object_name: str) -> str: - """Format grasp result for agent/human consumption.""" - best = grasps.poses[0] - pos = best.position - rpy = quaternion_to_euler(best.orientation, degrees=True) - return ( - f"Generated {len(grasps.poses)}" - f"Best grasp: pos=({pos.x:.4f}, {pos.y:.4f}, {pos.z:.4f}), " - f"rpy=({rpy.x:.1f}, {rpy.y:.1f}, {rpy.z:.1f}) degrees" - ) diff --git a/dimos/manipulation/grasping/test_grasp_gen_x.py b/dimos/manipulation/grasping/test_grasp_gen_x.py new file mode 100644 index 0000000000..4becb2b784 --- /dev/null +++ b/dimos/manipulation/grasping/test_grasp_gen_x.py @@ -0,0 +1,301 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hermetic tests for the import-safe GraspGenX adapter.""" + +from __future__ import annotations + +import inspect +import subprocess +import sys +from typing import Any + +import numpy as np +import pytest +from pytest_mock import MockerFixture + +from dimos.manipulation.grasping.grasp_gen_spec import GraspGenSpec +import dimos.manipulation.grasping.grasp_gen_x as grasp_gen_x +from dimos.manipulation.grasping.grasp_gen_x import ( + GraspGenXConfig, + GraspGenXError, + GraspGenXModule, +) +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.std_msgs.Header import Header + + +def config(**overrides: object) -> GraspGenXConfig: + values: dict[str, object] = { + "gripper": { + "extents_open": (0.1, 0.1, 0.1), + "offset_open": (0.0, 0.0, 0.0), + "extents_half_open": (0.1, 0.1, 0.1), + "offset_half_open": (0.0, 0.0, 0.0), + "fingertip_depth": 0.1, + }, + } + values.update(overrides) + return GraspGenXConfig(**values) # type: ignore[arg-type] + + +def module_args(value: GraspGenXConfig | None = None) -> dict[str, Any]: + return (value or config()).model_dump(exclude={"rpc_transport", "tf_transport", "g"}) + + +def cloud(points: np.ndarray | None = None) -> PointCloud2: + xyz = np.zeros((1, 3), dtype=np.float32) if points is None else points + return PointCloud2.from_numpy(xyz, frame_id="camera", timestamp=12.5) + + +def test_public_adapter_import_does_not_load_optional_runtime() -> None: + result = subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys; " + "import dimos.manipulation.grasping.grasp_gen_x; " + "assert 'dimos.manipulation.grasping.grasp_gen_x_runtime' not in sys.modules" + ), + ], + check=False, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stderr + + +@pytest.fixture +def runtime(mocker: MockerFixture) -> Any: + create_runtime = mocker.patch.object(grasp_gen_x, "_create_runtime") + instance = create_runtime.return_value + instance.infer.return_value = ( + np.repeat(np.eye(4, dtype=np.float32)[None], 1, axis=0), + np.asarray([0.5], dtype=np.float32), + ) + return create_runtime + + +def test_messages_round_trip_empty_and_score() -> None: + value = GraspCandidateArray(Header(3.0, "camera"), [GraspCandidate(Pose(1, 2, 3), 0.25)]) + decoded = GraspCandidateArray.decode(value.encode()) + + assert decoded.header.frame_id == "camera" + assert decoded.header.timestamp == pytest.approx(3.0) + assert decoded.candidates[0].score == pytest.approx(0.25) + assert ( + GraspCandidateArray.decode( + GraspCandidateArray(Header(3.0, "camera"), []).encode() + ).candidates + == [] + ) + + +def test_candidate_array_lcm_round_trip() -> None: + value = GraspCandidateArray( + Header(3.0, "camera"), [GraspCandidate(Pose(1, 2, 3), 0.25)], selected_index=1 + ) + + decoded = GraspCandidateArray.lcm_decode(value.lcm_encode()) + + assert decoded.header.frame_id == "camera" + assert decoded.candidates[0].score == pytest.approx(0.25) + assert decoded.selected_index == 1 + + +def test_spec_signature() -> None: + signature = inspect.signature(GraspGenSpec.propose_grasps) + + assert list(signature.parameters) == ["self", "object_pointcloud"] + assert signature.parameters["object_pointcloud"].annotation.__name__ == "PointCloud2" + assert signature.return_annotation is GraspCandidateArray + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("family", "unsupported"), + ("extents_open", (0.1, 0.1)), + ("extents_open", (0.0, 0.1, 0.1)), + ("extents_open", (0.6, 0.1, 0.1)), + ("offset_open", (0.6, 0.0, 0.0)), + ("offset_open", (np.nan, 0.0, 0.0)), + ("fingertip_depth", 0.0), + ], +) +def test_gripper_constraints_are_declared_by_fields(field: str, value: object) -> None: + gripper = config().gripper.model_dump() + + with pytest.raises(ValueError): + config(gripper={**gripper, field: value}) + + +@pytest.mark.parametrize("value", [0, -1, True]) +def test_candidate_limit_is_a_strict_positive_integer(value: object) -> None: + with pytest.raises(ValueError): + config(max_candidates=value) + + +def test_rigid_transform_relational_validation() -> None: + with pytest.raises(ValueError, match="orthonormal"): + config( + grasp_frame_to_tcp=( + (2.0, 0.0, 0.0, 0.0), + (0.0, 1.0, 0.0, 0.0), + (0.0, 0.0, 1.0, 0.0), + (0.0, 0.0, 0.0, 1.0), + ) + ) + + +def test_start_is_synchronous_and_idempotent(runtime: Any) -> None: + module = GraspGenXModule(**module_args()) + try: + module.start() + module.start() + + runtime.assert_called_once_with(module.config) + assert len(module.propose_grasps(cloud())) == 1 + finally: + module.stop() + runtime.return_value.stop.assert_called_once_with() + + +def test_start_failure_is_explicit(runtime: Any) -> None: + runtime.side_effect = RuntimeError("CUDA unavailable") + module = GraspGenXModule(**module_args()) + try: + with pytest.raises(GraspGenXError, match="initialize"): + module.start() + finally: + module.stop() + + +def test_adapter_sorts_stably_truncates_and_applies_tcp_transform(runtime: Any) -> None: + poses = np.repeat(np.eye(4, dtype=np.float32)[None], 3, axis=0) + poses[:, 0, 3] = [1.0, 2.0, 3.0] + runtime.return_value.infer.return_value = ( + poses, + np.asarray([0.5, 0.5, 0.9], dtype=np.float32), + ) + cfg = config( + max_candidates=2, + grasp_frame_to_tcp=( + (0.0, -1.0, 0.0, 10.0), + (1.0, 0.0, 0.0, 0.0), + (0.0, 0.0, 1.0, 0.0), + (0.0, 0.0, 0.0, 1.0), + ), + ) + module = GraspGenXModule(**module_args(cfg)) + try: + module.start() + result = module.propose_grasps(cloud()) + + assert [candidate.score for candidate in result] == pytest.approx([0.9, 0.5]) + assert [candidate.pose.position.x for candidate in result] == pytest.approx([13.0, 11.0]) + assert result.header.frame_id == "camera" + assert result.header.timestamp == pytest.approx(12.5) + finally: + module.stop() + + +def test_empty_backend_result_preserves_input_header(runtime: Any) -> None: + runtime.return_value.infer.return_value = ( + np.empty((0, 4, 4), dtype=np.float32), + np.empty(0, dtype=np.float32), + ) + module = GraspGenXModule(**module_args()) + try: + module.start() + result = module.propose_grasps(cloud()) + + assert result.header.frame_id == "camera" + assert result.header.timestamp == pytest.approx(12.5) + assert result.candidates == [] + finally: + module.stop() + + +@pytest.mark.parametrize( + "points", + [ + np.array([[np.nan, 0.0, 0.0]], dtype=np.float32), + np.empty((0, 3), dtype=np.float32), + np.zeros((2, 2), dtype=np.float32), + ], +) +def test_invalid_cloud_points_are_rejected(runtime: Any, points: np.ndarray) -> None: + module = GraspGenXModule(**module_args()) + try: + module.start() + with pytest.raises(ValueError, match="pointcloud|XYZ"): + module.propose_grasps(cloud(points)) + runtime.return_value.infer.assert_not_called() + finally: + module.stop() + + +@pytest.mark.parametrize( + "backend", + [ + (np.ones((2, 4, 4)), np.ones(1)), + (np.full((1, 4, 4), np.nan), np.ones(1)), + (np.ones((1, 4, 4)), np.array([np.inf])), + ], +) +def test_invalid_backend_outputs_are_rejected( + runtime: Any, backend: tuple[np.ndarray, np.ndarray] +) -> None: + runtime.return_value.infer.return_value = backend + module = GraspGenXModule(**module_args()) + try: + module.start() + with pytest.raises(ValueError): + module.propose_grasps(cloud()) + finally: + module.stop() + + +def test_inference_failure_is_wrapped(runtime: Any) -> None: + runtime.return_value.infer.side_effect = RuntimeError("backend") + module = GraspGenXModule(**module_args()) + try: + module.start() + with pytest.raises(GraspGenXError, match="inference"): + module.propose_grasps(cloud()) + finally: + module.stop() + + +def test_lazy_runtime_loading_and_missing_metadata_are_rejected(runtime: Any) -> None: + module = GraspGenXModule(**module_args()) + missing_frame = cloud() + missing_frame.frame_id = "" + missing_timestamp = cloud() + missing_timestamp.ts = None + try: + assert len(module.propose_grasps(cloud())) == 1 + runtime.assert_called_once_with(module.config) + with pytest.raises(ValueError, match="frame_id"): + module.propose_grasps(missing_frame) + with pytest.raises(ValueError, match="timestamp"): + module.propose_grasps(missing_timestamp) + finally: + module.stop() diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 137d8af72c..ed4a155d6a 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -25,6 +25,7 @@ from __future__ import annotations from collections.abc import Mapping, Sequence +from dataclasses import dataclass from enum import Enum import math import threading @@ -90,6 +91,7 @@ NoManipulationVisualizationConfig, ) from dimos.manipulation.visualization.factory import create_manipulation_visualization +from dimos.manipulation.visualization.layers import VisualizationLayer from dimos.manipulation.visualization.operator import ManipulationOperator from dimos.manipulation.visualization.types import TargetEvaluation from dimos.msgs.geometry_msgs.Pose import Pose @@ -136,6 +138,15 @@ class ManipulationState(Enum): FAULT = 4 +@dataclass(frozen=True) +class ConnectedPoseSequenceResult: + """Motion-free result for an ordered sequence of pose plans.""" + + failed_index: int | None + endpoint: JointState | None + paths: tuple[tuple[JointState, ...], ...] + + class ManipulationModuleConfig(ModuleConfig): """Configuration for ManipulationModule.""" @@ -998,6 +1009,14 @@ def plan_to_pose(self, pose: Pose, robot_name: RobotName | None = None) -> bool: return False return self.plan_to_pose_targets({group_id: pose}) + @rpc + def set_visualization_layer(self, layer: VisualizationLayer) -> bool: + """Replace one display-only layer in the active manipulation visualizer.""" + if self._world_monitor is None or self._world_monitor.visualization is None: + return False + self._world_monitor.visualization.set_layer(layer) + return True + @rpc def plan_to_pose_targets( self, @@ -1122,6 +1141,152 @@ def generate_plan_to_pose_targets( logger.info(f"IK solved, error: {ik.position_error:.4f}m") return self._plan_selected_path(group_ids, start, ik.joint_state, planning_epoch) + def _check_connected_pose_sequence( + self, + poses: Sequence[Pose], + robot_name: RobotName, + start: JointState | None = None, + ) -> tuple[int | None, JointState | None]: + """Dry-run a connected pose sequence and return failure index and endpoint. + + The check does not store a plan or change module state. Each IK solve and + path plan starts at the preceding path's endpoint. When ``start`` is + omitted, the sequence begins at the robot's current authoritative state. + """ + result = self._plan_connected_pose_sequence(poses, robot_name, start) + return result.failed_index, result.endpoint + + def _plan_connected_pose_sequence( + self, + poses: Sequence[Pose], + robot_name: RobotName, + start: JointState | None = None, + ) -> ConnectedPoseSequenceResult: + """Dry-run connected pose plans and retain each successful segment path.""" + if not poses: + return ConnectedPoseSequenceResult(None, start, ()) + if self._world_monitor is None or self._kinematics is None or self._planner is None: + logger.warning("Connected pose planning is unavailable") + return ConnectedPoseSequenceResult(0, None, ()) + try: + group_id = self._require_unique_pose_group_id_for_robot(robot_name) + selection = self._world_monitor.planning_groups.select((group_id,)) + if start is None: + current = self._world_monitor.current_global_joint_state() + start = filter_joint_state_to_selected_joints(current, selection.joint_names) + else: + start = filter_joint_state_to_selected_joints(start, selection.joint_names) + except (KeyError, ValueError) as exc: + logger.warning("Failed to initialize connected pose planning: %s", exc) + return ConnectedPoseSequenceResult(0, None, ()) + + paths: list[tuple[JointState, ...]] = [] + for index, pose in enumerate(poses): + target = PoseStamped( + frame_id="world", + position=pose.position, + orientation=pose.orientation, + ) + ik = self.inverse_kinematics( + pose_targets={group_id: target}, + seed=start, + check_collision=True, + ) + if not ik.is_success() or ik.joint_state is None: + logger.info( + "Connected pose planning failed IK at index %d: %s%s", + index, + ik.status.name, + f": {ik.message}" if ik.message else "", + ) + return ConnectedPoseSequenceResult(index, None, tuple(paths)) + result = self._planner.plan_selected_joint_path( + world=self._world_monitor.world, + selection=selection, + start=start, + goal=ik.joint_state, + timeout=self.config.planning_timeout, + ) + if not result.is_success() or not result.path: + logger.info( + "Connected pose planning failed path at index %d: %s%s", + index, + result.status.name, + f": {result.message}" if result.message else "", + ) + return ConnectedPoseSequenceResult(index, None, tuple(paths)) + try: + start = filter_joint_state_to_selected_joints( + result.path[-1], selection.joint_names + ) + except ValueError as exc: + logger.info( + "Connected pose planning returned an invalid endpoint at index %d: %s", + index, + exc, + ) + return ConnectedPoseSequenceResult(index, None, tuple(paths)) + paths.append(tuple(result.path)) + return ConnectedPoseSequenceResult(None, start, tuple(paths)) + + @rpc + def plan_cartesian_targets( + self, + targets: Mapping[PlanningGroupID | PlanningGroup, CartesianTarget], + config: CartesianPathConfig, + auxiliary_groups: Sequence[PlanningGroupID | PlanningGroup] = (), + ) -> bool: + """Plan TCP motion through absolute or relative Cartesian waypoints.""" + return self.generate_cartesian_plan(targets, config, auxiliary_groups) is not None + + def generate_cartesian_plan( + self, + targets: Mapping[PlanningGroupID | PlanningGroup, CartesianTarget], + config: CartesianPathConfig, + auxiliary_groups: Sequence[PlanningGroupID | PlanningGroup] = (), + ) -> GeneratedPlan | None: + """Generate and store a timed Cartesian plan through PlannerSpec.""" + if self._world_monitor is None or self._planner is None: + return None + if not targets: + self._fail("At least one Cartesian target is required") + return None + normalized_targets = { + planning_group_id_from_selector(group): target for group, target in targets.items() + } + if len(normalized_targets) != len(targets): + self._fail("Cartesian target groups must be unique") + return None + auxiliary_ids = tuple(planning_group_id_from_selector(group) for group in auxiliary_groups) + group_ids = tuple((*normalized_targets.keys(), *auxiliary_ids)) + planning_epoch = self._begin_group_planning() + if planning_epoch is None: + return None + resolved = self._resolve_group_plan_start(group_ids, planning_epoch) + if resolved is None: + return None + selection, start = resolved + result = self._planner.plan_cartesian_path( + world=self._world_monitor.world, + selection=selection, + start=start, + targets=normalized_targets, + config=config, + auxiliary_groups=auxiliary_ids, + ) + if not result.is_success(): + detail = f": {result.message}" if result.message else "" + self._fail_planning_epoch( + planning_epoch, f"Cartesian planning failed: {result.status.name}{detail}" + ) + return None + return self._store_generated_plan( + group_ids, + result, + planning_epoch, + preserve_timing=True, + ) + @rpc def plan_cartesian_targets( self, @@ -1541,6 +1706,11 @@ def execute_plan(self, plan: GeneratedPlan | None = None) -> bool: self._error_message = result.message return bool(result and result.accepted) + @rpc + def execute_and_wait(self, timeout: float = 60.0) -> bool: + """Execute the stored plan and wait for its expected trajectory duration.""" + return self.execute_plan() and self._wait_for_trajectory_completion(timeout) + @property def world_monitor(self) -> WorldMonitor | None: """Access the world monitor for advanced obstacle/world operations.""" @@ -1578,6 +1748,41 @@ def add_obstacle( ) return self._world_monitor.add_obstacle(obstacle) + @skill(uses=["movement"]) + def set_table_collision( + self, + center_x: float, + center_y: float, + tabletop_z: float, + width: float, + depth: float, + safety_margin: float = 0.0, + thickness: float = 0.20, + ) -> bool: + """Install or update a horizontal table collision slab. + + All dimensions are meters. ``tabletop_z`` is the measured physical tabletop height. The default + uses no added clearance so a grasp can descend to an object's measured contact pose. + """ + if self._world_monitor is None: + return False + if width <= 0.0 or depth <= 0.0 or thickness <= 0.0 or safety_margin < 0.0: + raise ValueError("table dimensions must be positive and safety_margin non-negative") + protected_top = tabletop_z + safety_margin + table = Obstacle( + name="calibrated-table", + obstacle_type=ObstacleType.BOX, + pose=PoseStamped( + position=Vector3(center_x, center_y, protected_top - thickness / 2), + orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + ), + dimensions=(width, depth, thickness), + color=(0.2, 0.5, 0.9, 0.35), + ) + if self._world_monitor.update_obstacle(table): + return True + return bool(self._world_monitor.add_obstacle(table)) + @rpc def update_obstacle( self, @@ -1710,24 +1915,43 @@ def _wait_for_trajectory_completion(self, timeout: float = 60.0) -> bool: time.sleep(wait_time) return True - def _lift_if_low( + def _safety_lift_pose( self, robot_name: RobotName | None = None, min_z: float = 0.05 - ) -> SkillResult[ManipulationSkillError]: - """If the end-effector is below *min_z*, plan and execute a short lift.""" + ) -> Pose | None: + """Return the required safety-lift target, if the end effector is low.""" ee = self.get_ee_pose(robot_name) if ee is None or ee.position.z >= min_z: - return SkillResult.ok() + return None lift_z = min_z + 0.05 logger.info(f"EE z={ee.position.z:.3f} < {min_z}, lifting to z={lift_z:.3f}") - lift_pose = Pose(Vector3(ee.position.x, ee.position.y, lift_z), ee.orientation) + return Pose(Vector3(ee.position.x, ee.position.y, lift_z), ee.orientation) + + def _lift_if_low( + self, robot_name: RobotName | None = None, min_z: float = 0.05 + ) -> SkillResult[ManipulationSkillError]: + """If the end-effector is below *min_z*, plan and execute a short lift.""" + lift_pose = self._safety_lift_pose(robot_name, min_z) + if lift_pose is None: + return SkillResult.ok() if not self.plan_to_pose(lift_pose, robot_name): return SkillResult.fail( "PLANNING_FAILED", - f"Failed to plan lift from z={ee.position.z:.3f}", + f"Failed to plan safety lift to z={lift_pose.position.z:.3f}", ) return self._preview_execute_wait(robot_name) + def _planner_fault_result(self) -> SkillResult[ManipulationSkillError] | None: + """Return an actionable failure while the planner requires recovery.""" + with self._lock: + if self._state != ManipulationState.FAULT: + return None + detail = self._error_message or "unknown planner error" + return SkillResult.fail( + "INVALID_STATE", + f"Planner is FAULT ({detail}). Call reset before issuing another motion command.", + ) + def _preview_execute_wait( self, robot_name: RobotName | None = None, preview_duration: float = 0.5 ) -> SkillResult[ManipulationSkillError]: @@ -1806,6 +2030,9 @@ def move_to_pose( yaw: Target yaw in radians (omit to keep current orientation). robot_name: Robot to move (only needed for multi-arm setups). """ + if fault := self._planner_fault_result(): + return fault + logger.info(f"Planning motion to ({x:.3f}, {y:.3f}, {z:.3f})...") # If no orientation specified, preserve the current EE orientation. @@ -1863,6 +2090,9 @@ def move_to_joints( joints: Comma-separated joint positions in radians, e.g. "0.1, -0.5, 1.2, 0.0, 0.3, -0.1". robot_name: Robot to move (only needed for multi-arm setups). """ + if fault := self._planner_fault_result(): + return fault + try: joint_values = [float(j.strip()) for j in joints.split(",")] except ValueError: @@ -1899,6 +2129,9 @@ def go_home(self, robot_name: str | None = None) -> SkillResult[ManipulationSkil Args: robot_name: Robot to move (only needed for multi-arm setups). """ + if fault := self._planner_fault_result(): + return fault + robot = self._get_robot(robot_name) if robot is None: return SkillResult.fail("ROBOT_NOT_FOUND", "Robot not found") @@ -1935,6 +2168,9 @@ def go_init(self, robot_name: str | None = None) -> SkillResult[ManipulationSkil Args: robot_name: Robot to move (only needed for multi-arm setups). """ + if fault := self._planner_fault_result(): + return fault + robot = self._get_robot(robot_name) if robot is None: return SkillResult.fail("ROBOT_NOT_FOUND", "Robot not found") diff --git a/dimos/manipulation/obstacle_world_spec.py b/dimos/manipulation/obstacle_world_spec.py new file mode 100644 index 0000000000..f0870ee7fb --- /dev/null +++ b/dimos/manipulation/obstacle_world_spec.py @@ -0,0 +1,43 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Planner obstacle mutation protocol for scene-derived geometry.""" + +from typing import Literal, Protocol + +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.spec.utils import Spec + + +class ObstacleWorldSpec(Spec, Protocol): + def add_obstacle( + self, + name: str, + pose: Pose, + shape: Literal["box", "sphere", "cylinder", "mesh"], + dimensions: list[float] | None = None, + mesh_path: str | None = None, + ) -> str: ... + + def update_obstacle( + self, + name: str, + pose: Pose, + shape: Literal["box", "sphere", "cylinder", "mesh"], + dimensions: list[float] | None = None, + mesh_path: str | None = None, + color: list[float] | None = None, + ) -> bool: ... + + def remove_obstacle(self, obstacle_id: str) -> bool: ... diff --git a/dimos/manipulation/pick_and_place_module.py b/dimos/manipulation/pick_and_place_module.py index c788277b94..6394a43217 100644 --- a/dimos/manipulation/pick_and_place_module.py +++ b/dimos/manipulation/pick_and_place_module.py @@ -22,14 +22,22 @@ from __future__ import annotations +from collections import Counter +from dataclasses import dataclass, field +from enum import Enum import math +import threading import time -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np +from pydantic import Field, FiniteFloat, model_validator from dimos.agents.annotation import skill from dimos.agents.skill_result import SkillResult from dimos.core.core import rpc from dimos.core.stream import In +from dimos.manipulation.grasping.grasp_gen_spec import GraspGenSpec from dimos.manipulation.manipulation_module import ( ManipulationModule, ManipulationModuleConfig, @@ -38,13 +46,18 @@ from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate from dimos.perception.experimental.object import ( Object as DetObject, ) +from dimos.perception.experimental.object_scene_registration_spec import ObjectSceneRegistrationSpec +from dimos.protocol.service.spec import BaseConfig from dimos.utils.logging_config import setup_logger +from dimos.utils.transform_utils import offset_distance if TYPE_CHECKING: from dimos.msgs.geometry_msgs.PoseArray import PoseArray + from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 logger = setup_logger() @@ -62,9 +75,102 @@ _TALL_OBJECT_MIN_HEIGHT = 0.06 +class GraspVerificationConfig(BaseConfig): + """Robot-specific gripper closure verification settings.""" + + enabled: bool = False + open_position: FiniteFloat = 0.85 + closed_position: FiniteFloat = 0.0 + held_threshold: FiniteFloat = 0.02 + timeout: FiniteFloat = Field(default=2.0, gt=0.0) + poll_interval: FiniteFloat = Field(default=0.05, gt=0.0) + + @model_validator(mode="after") + def _validate_threshold(self) -> GraspVerificationConfig: + low = min(self.open_position, self.closed_position) + high = max(self.open_position, self.closed_position) + if self.open_position == self.closed_position: + raise ValueError("gripper open_position and closed_position must differ") + if not low < self.held_threshold < high: + raise ValueError("held_threshold must lie between open_position and closed_position") + if self.poll_interval > self.timeout: + raise ValueError("poll_interval must not exceed timeout") + return self + + class PickAndPlaceModuleConfig(ManipulationModuleConfig): """Configuration for PickAndPlaceModule.""" + heuristic_grasp_fallback: bool = False + planning_frame: str = "world" + max_object_pointcloud_age: FiniteFloat = Field(default=10.0, gt=0.0) + max_grasp_candidates_to_check: int = Field(default=5, gt=0) + grasp_pre_grasp_offset: FiniteFloat | None = Field(default=None, gt=0.0) + grasp_retreat_offset: FiniteFloat | None = Field(default=None, gt=0.0) + grasp_approach_vector: tuple[FiniteFloat, FiniteFloat, FiniteFloat] = (0.0, 0.0, -1.0) + grasp_verification: GraspVerificationConfig = Field(default_factory=GraspVerificationConfig) + + @model_validator(mode="after") + def _validate_grasp_pipeline(self) -> PickAndPlaceModuleConfig: + if not self.planning_frame.strip(): + raise ValueError("planning_frame must not be empty") + vector = np.asarray(self.grasp_approach_vector, dtype=float) + if not np.isclose(np.linalg.norm(vector), 1.0, atol=1e-6): + raise ValueError("grasp_approach_vector must be a unit vector") + return self + + +class _PickPhase(str, Enum): + RESOLVE = "RESOLVE" + PROPOSE = "PROPOSE" + SELECT = "SELECT" + PREPARE = "PREPARE" + APPROACH = "APPROACH" + GRASP = "GRASP" + CLOSE = "CLOSE" + VERIFY = "VERIFY" + RETREAT = "RETREAT" + DONE = "DONE" + + +class _CandidateRejection(str, Enum): + INVALID = "invalid" + PRE_GRASP_INFEASIBLE = "pre_grasp_infeasible" + GRASP_INFEASIBLE = "grasp_infeasible" + RETREAT_INFEASIBLE = "retreat_infeasible" + + +@dataclass(frozen=True) +class _FeasibleGrasp: + candidate: GraspCandidate + rank: int + pre_grasp_pose: Pose + retreat_pose: Pose + + +@dataclass(frozen=True) +class _GraspVerification: + held: bool + position: float | None + detail: str + + +@dataclass +class _PickTransaction: + object_id: str = "" + object_name: str = "" + proposal_source: Literal["grasp_provider", "heuristic"] = "grasp_provider" + phase: _PickPhase = _PickPhase.RESOLVE + selected: _FeasibleGrasp | None = None + rejections: Counter[str] = field(default_factory=Counter) + gripper_closed: bool = False + + +class _PickPipelineError(RuntimeError): + def __init__(self, code: ManipulationSkillError, message: str) -> None: + super().__init__(message) + self.code = code + class PickAndPlaceModule(ManipulationModule): """Manipulation module with perception integration and pick-and-place skills. @@ -76,6 +182,8 @@ class PickAndPlaceModule(ManipulationModule): """ config: PickAndPlaceModuleConfig + _object_scene: ObjectSceneRegistrationSpec | None = None + _grasp_generator: GraspGenSpec | None = None # Input: Objects from perception (for obstacle integration) objects: In[list[DetObject]] @@ -90,6 +198,7 @@ def __init__(self, **kwargs: Any) -> None: # The live detection cache is volatile (labels change every frame), # so pick/place use this stable snapshot instead. self._detection_snapshot: list[DetObject] = [] + self._pick_guard = threading.Lock() @rpc def start(self) -> None: @@ -177,7 +286,12 @@ def generate_grasps( "GraspGen Docker support removed; see issue #1266 for re-implementation as NativeModule subclass" ) - def _compute_pre_grasp_pose(self, grasp_pose: Pose, offset: float = 0.10) -> Pose: + def _compute_pre_grasp_pose( + self, + grasp_pose: Pose, + offset: float = 0.10, + approach_vector: Vector3 | None = None, + ) -> Pose: """Compute a pre-grasp pose offset along the approach direction (local -Z). Args: @@ -187,9 +301,11 @@ def _compute_pre_grasp_pose(self, grasp_pose: Pose, offset: float = 0.10) -> Pos Returns: Pre-grasp pose offset from the grasp pose """ - from dimos.utils.transform_utils import offset_distance - - return offset_distance(grasp_pose, offset) + return offset_distance( + grasp_pose, + offset, + approach_vector if approach_vector is not None else Vector3(0.0, 0.0, -1.0), + ) def _find_object_in_detections( self, object_name: str, object_id: str | None = None @@ -224,10 +340,19 @@ def _find_object_in_detections( logger.warning(f"Ambiguous object_id prefix '{object_id}' matches {ids}") return None - # Second pass: match by name - for det in self._detection_snapshot: - if object_name.lower() in det.name.lower() or det.name.lower() in object_name.lower(): - return det + # Second pass: require a unique name match. + normalized = object_name.casefold() + name_matches = [ + det + for det in self._detection_snapshot + if normalized in det.name.casefold() or det.name.casefold() in normalized + ] + if len(name_matches) == 1: + return name_matches[0] + if len(name_matches) > 1: + ids = [det.object_id for det in name_matches] + logger.warning("Ambiguous object name", object_name=object_name, object_ids=ids) + return None available = [det.name for det in self._detection_snapshot] logger.warning(f"Object '{object_name}' not found in snapshot. Available: {available}") @@ -454,7 +579,8 @@ def scan_objects( for det in detections: c = det.center lines.append( - f" - {det.name}: ({c.x:.3f}, {c.y:.3f}, {c.z:.3f}) [{det.detections_count} views]" + f" - {det.name} [id={det.object_id[:8]}]: " + f"({c.x:.3f}, {c.y:.3f}, {c.z:.3f}) [{det.detections_count} views]" ) if obstacles: @@ -462,6 +588,282 @@ def scan_objects( return SkillResult.ok("\n".join(lines)) + def _require_pick_object(self, object_name: str, object_id: str | None) -> DetObject: + detection = self._find_object_in_detections(object_name, object_id) + if detection is not None: + return detection + selector = f"id '{object_id}'" if object_id else f"name '{object_name}'" + raise _PickPipelineError( + "OBJECT_NOT_DETECTED", + f"No unique current detection matches {selector}; scan again and use an object ID", + ) + + def _provider_candidates( + self, detection: DetObject, transaction: _PickTransaction + ) -> list[GraspCandidate]: + if self._grasp_generator is None: + if not self.config.heuristic_grasp_fallback: + raise _PickPipelineError( + "GRASP_PROVIDER_UNAVAILABLE", + "No grasp proposal provider is connected and heuristic fallback is disabled", + ) + transaction.proposal_source = "heuristic" + poses = self._generate_grasps_for_pick(detection.name, detection.object_id) + if not poses: + raise _PickPipelineError( + "GRASP_GENERATION_FAILED", + f"Heuristic grasp generation failed for '{detection.name}'", + ) + return [GraspCandidate(pose=pose, score=0.0) for pose in poses] + + if self._object_scene is None: + raise _PickPipelineError( + "GRASP_PROVIDER_UNAVAILABLE", + "No object-scene provider is connected for learned grasp input", + ) + + pointcloud = self._object_scene.get_object_pointcloud_by_object_id(detection.object_id) + if pointcloud is None: + raise _PickPipelineError( + "GRASP_INPUT_INVALID", + f"No point cloud is available for object '{detection.object_id}'", + ) + points = pointcloud.points_f32() + if points.ndim != 2 or points.shape[1] != 3 or len(points) == 0: + raise _PickPipelineError( + "GRASP_INPUT_INVALID", + f"Object '{detection.object_id}' has an empty or invalid point cloud", + ) + if ( + pointcloud.ts is None + or time.time() - pointcloud.ts > self.config.max_object_pointcloud_age + ): + raise _PickPipelineError( + "GRASP_INPUT_INVALID", + f"Object '{detection.object_id}' point cloud is stale", + ) + if pointcloud.frame_id != self.config.planning_frame: + raise _PickPipelineError( + "GRASP_FRAME_MISMATCH", + f"Object cloud frame '{pointcloud.frame_id}' does not match " + f"planning frame '{self.config.planning_frame}'", + ) + + try: + proposals = self._grasp_generator.propose_grasps(pointcloud) + except Exception as exc: + raise _PickPipelineError( + "GRASP_GENERATION_FAILED", f"Grasp proposal failed: {exc}" + ) from exc + if proposals.header.frame_id != self.config.planning_frame: + raise _PickPipelineError( + "GRASP_FRAME_MISMATCH", + f"Proposal frame '{proposals.header.frame_id}' does not match " + f"planning frame '{self.config.planning_frame}'", + ) + if not proposals.candidates: + raise _PickPipelineError( + "GRASP_GENERATION_FAILED", + f"No grasp proposals were generated for '{detection.name}'", + ) + return sorted(proposals.candidates, key=lambda candidate: candidate.score, reverse=True) + + @staticmethod + def _valid_candidate(candidate: GraspCandidate) -> bool: + pose = candidate.pose + values = np.asarray( + [ + pose.position.x, + pose.position.y, + pose.position.z, + pose.orientation.x, + pose.orientation.y, + pose.orientation.z, + pose.orientation.w, + candidate.score, + ], + dtype=float, + ) + quaternion = values[3:7] + return bool( + np.all(np.isfinite(values)) and np.isclose(np.linalg.norm(quaternion), 1.0, atol=1e-5) + ) + + def _select_feasible_grasp( + self, + candidates: list[GraspCandidate], + robot_name: str, + robot_pre_grasp_offset: float, + transaction: _PickTransaction, + sequence_start: JointState | None = None, + ) -> _FeasibleGrasp: + vector = Vector3(self.config.grasp_approach_vector) + pre_offset = self.config.grasp_pre_grasp_offset or robot_pre_grasp_offset + retreat_offset = self.config.grasp_retreat_offset or pre_offset + limit = min(len(candidates), self.config.max_grasp_candidates_to_check) + + for rank, candidate in enumerate(candidates[:limit], start=1): + if not self._valid_candidate(candidate): + transaction.rejections[_CandidateRejection.INVALID.value] += 1 + continue + pre_grasp = self._compute_pre_grasp_pose(candidate.pose, pre_offset, vector) + retreat = self._compute_pre_grasp_pose(candidate.pose, retreat_offset, vector) + rejections = ( + _CandidateRejection.PRE_GRASP_INFEASIBLE, + _CandidateRejection.GRASP_INFEASIBLE, + _CandidateRejection.RETREAT_INFEASIBLE, + ) + failed_index, _ = self._check_connected_pose_sequence( + (pre_grasp, candidate.pose, retreat), + robot_name, + start=sequence_start, + ) + if failed_index is not None: + transaction.rejections[rejections[failed_index].value] += 1 + continue + return _FeasibleGrasp(candidate, rank, pre_grasp, retreat) + + summary = ", ".join( + f"{reason}={count}" for reason, count in sorted(transaction.rejections.items()) + ) + raise _PickPipelineError( + "GRASP_ATTEMPTS_EXHAUSTED", + f"No feasible grasp among {limit} candidate(s)" + (f" ({summary})" if summary else ""), + ) + + def _verify_grasp(self, robot_name: str) -> _GraspVerification: + verification = self.config.grasp_verification + if not verification.enabled: + return _GraspVerification(True, None, "gripper feedback verification disabled") + + deadline = time.monotonic() + verification.timeout + last_position: float | None = None + while time.monotonic() < deadline: + last_position = self.get_gripper(robot_name) + if last_position is not None: + closes_upward = verification.closed_position > verification.open_position + empty = ( + last_position >= verification.held_threshold + if closes_upward + else last_position <= verification.held_threshold + ) + if empty: + return _GraspVerification( + False, last_position, "gripper reached the empty-closed region" + ) + time.sleep(verification.poll_interval) + + if last_position is None: + return _GraspVerification(False, None, "gripper feedback was unavailable") + movement = abs(last_position - verification.open_position) + if movement < 1e-3: + return _GraspVerification( + False, last_position, "gripper did not leave the open position" + ) + closes_upward = verification.closed_position > verification.open_position + held = ( + last_position < verification.held_threshold + if closes_upward + else last_position > verification.held_threshold + ) + detail = ( + "grasp verified by gripper closure feedback" + if held + else "gripper reached the empty-closed region" + ) + return _GraspVerification(held, last_position, detail) + + @staticmethod + def _phase_failure( + transaction: _PickTransaction, + code: ManipulationSkillError, + message: str, + ) -> SkillResult[ManipulationSkillError]: + may_hold = transaction.gripper_closed + suffix = "; object may be held" if may_hold else "" + result = SkillResult[ManipulationSkillError].fail( + code, f"{transaction.phase.value}: {message}{suffix}" + ) + result.metadata = { + "phase": transaction.phase.value, + "object_id": transaction.object_id, + "proposal_source": transaction.proposal_source, + "object_may_be_held": may_hold, + "rejections": dict(transaction.rejections), + } + if transaction.selected is not None: + result.metadata.update( + candidate_rank=transaction.selected.rank, + candidate_score=transaction.selected.candidate.score, + ) + return result + + def _execute_selected_pick( + self, transaction: _PickTransaction, robot_name: str + ) -> SkillResult[ManipulationSkillError]: + assert transaction.selected is not None + selected = transaction.selected + verification = self.config.grasp_verification + + transaction.phase = _PickPhase.PREPARE + lift = self._lift_if_low(robot_name) + if not lift.is_success(): + return self._phase_failure( + transaction, lift.error_code or "EXECUTION_FAILED", lift.message + ) + if not self._set_gripper_position(float(verification.open_position), robot_name): + return self._phase_failure(transaction, "GRIPPER_FAILED", "open command failed") + + transaction.phase = _PickPhase.APPROACH + if not self.plan_to_pose(selected.pre_grasp_pose, robot_name): + return self._phase_failure(transaction, "PLANNING_FAILED", "pre-grasp planning failed") + execution = self._preview_execute_wait(robot_name) + if not execution.is_success(): + return self._phase_failure( + transaction, execution.error_code or "EXECUTION_FAILED", execution.message + ) + + transaction.phase = _PickPhase.GRASP + if not self.plan_to_pose(selected.candidate.pose, robot_name): + return self._phase_failure(transaction, "PLANNING_FAILED", "grasp planning failed") + execution = self._preview_execute_wait(robot_name) + if not execution.is_success(): + return self._phase_failure( + transaction, execution.error_code or "EXECUTION_FAILED", execution.message + ) + + transaction.phase = _PickPhase.CLOSE + if not self._set_gripper_position(float(verification.closed_position), robot_name): + return self._phase_failure(transaction, "GRIPPER_FAILED", "close command failed") + transaction.gripper_closed = True + + transaction.phase = _PickPhase.VERIFY + verified = self._verify_grasp(robot_name) + if not verified.held: + return self._phase_failure(transaction, "GRASP_VERIFICATION_FAILED", verified.detail) + + transaction.phase = _PickPhase.RETREAT + if not self.plan_to_pose(selected.retreat_pose, robot_name): + return self._phase_failure(transaction, "PLANNING_FAILED", "retreat planning failed") + execution = self._preview_execute_wait(robot_name) + if not execution.is_success(): + return self._phase_failure( + transaction, execution.error_code or "EXECUTION_FAILED", execution.message + ) + + transaction.phase = _PickPhase.DONE + self._last_pick_pose = selected.candidate.pose + return SkillResult.ok( + f"Pick complete — grasped '{transaction.object_name}' using candidate " + f"{selected.rank} (score={selected.candidate.score:.4f}); {verified.detail}", + object_id=transaction.object_id, + proposal_source=transaction.proposal_source, + candidate_rank=selected.rank, + candidate_score=selected.candidate.score, + verification=verified.detail, + rejections=dict(transaction.rejections), + ) + @skill def pick( self, @@ -479,80 +881,64 @@ def pick( object_id: Optional unique object ID from perception for precise identification. robot_name: Robot to use (only needed for multi-arm setups). """ - robot = self._get_robot(robot_name) - if robot is None: - return SkillResult.fail("ROBOT_NOT_FOUND", "Robot not found") - rname, _, config, _ = robot - pre_grasp_offset = config.pre_grasp_offset - - # 1. Generate grasps (uses already-cached detections — call scan_objects first) - logger.info(f"Generating grasp poses for '{object_name}'...") - grasp_poses = self._generate_grasps_for_pick(object_name, object_id) - if not grasp_poses: - return SkillResult.fail( - "GRASP_GENERATION_FAILED", - f"No grasp poses found for '{object_name}'. Object may not be detected.", - ) + if not self._pick_guard.acquire(blocking=False): + return SkillResult.fail("PICK_BUSY", "Another pick transaction is active") - # Lift if EE is low before approaching - lift = self._lift_if_low(rname) - if not lift.is_success(): - return lift - - # 2. Try each grasp candidate - max_attempts = min(len(grasp_poses), 5) - for i, grasp_pose in enumerate(grasp_poses[:max_attempts]): - # Reduce pre-grasp height for far objects (arm can't reach high + far) - gp = grasp_pose.position - xy_dist = (gp.x**2 + gp.y**2) ** 0.5 - offset = pre_grasp_offset if xy_dist < _FAR_REACH_XY_THRESHOLD else 0.05 - pre_grasp_pose = self._compute_pre_grasp_pose(grasp_pose, offset) - - logger.info(f"Planning approach to pre-grasp (attempt {i + 1}/{max_attempts})...") - if not self.plan_to_pose(pre_grasp_pose, rname): - logger.info(f"Grasp candidate {i + 1} approach planning failed, trying next") - continue # Try next candidate - - # 3. Open gripper before approach - logger.info("Opening gripper...") - self._set_gripper_position(0.85, rname) - time.sleep(0.5) - - # 4. Execute approach to pre-grasp - exec_result = self._preview_execute_wait(rname) - if not exec_result.is_success(): - return exec_result - - # 5. Move to grasp pose - logger.info("Moving to grasp position...") - if not self.plan_to_pose(grasp_pose, rname): - return SkillResult.fail("PLANNING_FAILED", "Grasp pose planning failed") - exec_result = self._preview_execute_wait(rname) - if not exec_result.is_success(): - return exec_result - - # 6. Close gripper - logger.info("Closing gripper...") - self._set_gripper_position(0.0, rname) - time.sleep(1.5) # Wait for gripper to close - - # 7. Retract to pre-grasp - logger.info("Retracting with object...") - if not self.plan_to_pose(pre_grasp_pose, rname): - return SkillResult.fail("PLANNING_FAILED", "Retract planning failed") - exec_result = self._preview_execute_wait(rname) - if not exec_result.is_success(): - return exec_result - - # Store pick pose so place_back() can return with same orientation - self._last_pick_pose = grasp_pose - - return SkillResult.ok(f"Pick complete — grasped '{object_name}' successfully") - - return SkillResult.fail( - "GRASP_ATTEMPTS_EXHAUSTED", - f"All {max_attempts} grasp attempts failed for '{object_name}'", - ) + transaction = _PickTransaction() + suppression = None + result: SkillResult[ManipulationSkillError] + try: + robot = self._get_robot(robot_name) + if robot is None: + return SkillResult.fail("ROBOT_NOT_FOUND", "Robot not found") + rname, _, robot_config, _ = robot + + detection = self._require_pick_object(object_name, object_id) + transaction.object_id = detection.object_id + transaction.object_name = detection.name + transaction.phase = _PickPhase.PROPOSE + candidates = self._provider_candidates(detection, transaction) + + if self._world_monitor is None: + raise _PickPipelineError( + "WORLD_MONITOR_UNAVAILABLE", "Planning world monitor is unavailable" + ) + + with self._world_monitor.suppress_object_obstacle(detection.object_id) as suppression: + sequence_start = None + lift_pose = self._safety_lift_pose(rname) + if lift_pose is not None: + transaction.phase = _PickPhase.PREPARE + failed_index, sequence_start = self._check_connected_pose_sequence( + (lift_pose,), rname + ) + if failed_index is not None: + raise _PickPipelineError( + "PLANNING_FAILED", + "Required safety-lift planning failed", + ) + transaction.phase = _PickPhase.SELECT + transaction.selected = self._select_feasible_grasp( + candidates, + rname, + robot_config.pre_grasp_offset, + transaction, + sequence_start, + ) + result = self._execute_selected_pick(transaction, rname) + if suppression.cleanup_error is not None: + if result.is_success(): + return self._phase_failure( + transaction, "WORLD_MONITOR_UNAVAILABLE", suppression.cleanup_error + ) + result.message = f"{result.message}; cleanup: {suppression.cleanup_error}" + return result + except _PickPipelineError as exc: + return self._phase_failure(transaction, exc.code, str(exc)) + except RuntimeError as exc: + return self._phase_failure(transaction, "WORLD_MONITOR_UNAVAILABLE", str(exc)) + finally: + self._pick_guard.release() @skill def place( diff --git a/dimos/manipulation/pick_execution_spec.py b/dimos/manipulation/pick_execution_spec.py new file mode 100644 index 0000000000..e09236f39d --- /dev/null +++ b/dimos/manipulation/pick_execution_spec.py @@ -0,0 +1,42 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Manipulation operations used by the selected-object pick transaction.""" + +from typing import Protocol + +from dimos.agents.skill_result import SkillResult +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.spec.utils import Spec + + +class PickExecutionSpec(Spec, Protocol): + def open_gripper(self, robot_name: str | None = None) -> SkillResult: ... + + def close_gripper(self, robot_name: str | None = None) -> SkillResult: ... + + def get_gripper(self, robot_name: str | None = None) -> float | None: ... + + def get_ee_pose(self, robot_name: str | None = None) -> Pose | None: ... + + def move_to_pose( + self, + x: float, + y: float, + z: float, + roll: float | None = None, + pitch: float | None = None, + yaw: float | None = None, + robot_name: str | None = None, + ) -> SkillResult: ... diff --git a/dimos/manipulation/picknplace.py b/dimos/manipulation/picknplace.py new file mode 100644 index 0000000000..8fc5bb5054 --- /dev/null +++ b/dimos/manipulation/picknplace.py @@ -0,0 +1,936 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Request-driven perception interface for the pick-and-place workflow.""" + +import math +import threading +import time +from typing import Literal + +import numpy as np +from pydantic import AliasChoices, Field + +from dimos.agents.annotation import skill +from dimos.agents.capabilities import CAP_MOVEMENT, CAP_PERCEPTION +from dimos.agents.skill_result import SkillResult +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.manipulation.candidate_filter_spec import GraspCandidateFilterSpec +from dimos.manipulation.grasping.grasp_gen_spec import GraspGenSpec +from dimos.manipulation.obstacle_world_spec import ObstacleWorldSpec +from dimos.manipulation.pick_execution_spec import PickExecutionSpec +from dimos.manipulation.visualization.layers import ( + LineSetElement, + MeshElement, + PointCloudElement, + VisualizationLayer, +) +from dimos.manipulation.visualization.pose_overlay import draw_pose_axes +from dimos.manipulation.visualization_spec import ManipulationVisualizationSpec +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray +from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.perception.experimental.object import ( + Object as DetObject, + to_detection3d_array, +) +from dimos.perception.experimental.object_scene_registration_spec import ObjectSceneRegistrationSpec + + +def _estimate_table_surface(points: np.ndarray) -> dict[str, float] | None: + """Fit the dominant horizontal support plane and return a conservative footprint.""" + if points.ndim != 2 or points.shape[1] != 3 or len(points) < 30: + return None + import open3d as o3d # type: ignore[import-untyped] + + cloud = o3d.geometry.PointCloud() + cloud.points = o3d.utility.Vector3dVector(points) + plane, inliers = cloud.segment_plane(distance_threshold=0.01, ransac_n=3, num_iterations=1000) + normal = np.asarray(plane[:3], dtype=np.float64) + normal /= np.linalg.norm(normal) + if abs(normal[2]) < 0.98 or len(inliers) < 30: + return None + surface = points[np.asarray(inliers)] + x_low, y_low = np.quantile(surface[:, :2], 0.02, axis=0) + x_high, y_high = np.quantile(surface[:, :2], 0.98, axis=0) + # Extend the observed tabletop patch so collision protection includes its edges. + margin = 0.10 + return { + "center_x": float((x_low + x_high) / 2), + "center_y": float((y_low + y_high) / 2), + "tabletop_z": float(np.median(surface[:, 2])), + "width": float(max(x_high - x_low + 2 * margin, 0.20)), + "depth": float(max(y_high - y_low + 2 * margin, 0.20)), + "inlier_count": float(len(inliers)), + } + + +def _table_midpoint_grasp_z( + points: np.ndarray, tabletop_z: float | None, fallback_z: float +) -> float: + """Return the midpoint from the physical table plane to an object's observed top surface.""" + if tabletop_z is None or points.ndim != 2 or points.shape[1] != 3 or len(points) < 10: + return fallback_z + top_z = float(np.quantile(points[:, 2], 0.95)) + if top_z <= tabletop_z: + return fallback_z + return tabletop_z + (top_z - tabletop_z) / 2.0 + + +def _primitive_mesh( + shape: Literal["box", "sphere", "cylinder"], + center: Vector3, + dimensions: tuple[float, ...], + orientation: Quaternion, +) -> tuple[np.ndarray, np.ndarray]: + """Create a display mesh for one planner primitive.""" + if shape == "box": + x, y, z = (dimension / 2.0 for dimension in dimensions) + vertices = np.asarray( + [ + [-x, -y, -z], + [x, -y, -z], + [x, y, -z], + [-x, y, -z], + [-x, -y, z], + [x, -y, z], + [x, y, z], + [-x, y, z], + ] + ) + triangles = np.asarray( + [ + [0, 1, 2], + [0, 2, 3], + [4, 6, 5], + [4, 7, 6], + [0, 4, 5], + [0, 5, 1], + [1, 5, 6], + [1, 6, 2], + [2, 6, 7], + [2, 7, 3], + [3, 7, 4], + [3, 4, 0], + ] + ) + else: + segments = 16 + angles = np.linspace(0.0, 2.0 * math.pi, segments, endpoint=False) + radius = dimensions[0] + if shape == "cylinder": + half_height = dimensions[1] / 2.0 + vertices = np.vstack( + ( + np.column_stack( + (radius * np.cos(angles), radius * np.sin(angles), -half_height) + ), + np.column_stack( + (radius * np.cos(angles), radius * np.sin(angles), half_height) + ), + [[0.0, 0.0, -half_height], [0.0, 0.0, half_height]], + ) + ) + bottom_center, top_center = 2 * segments, 2 * segments + 1 + triangles = np.asarray( + [ + triangle + for index in range(segments) + for triangle in ( + [index, (index + 1) % segments, segments + index], + [ + (index + 1) % segments, + segments + (index + 1) % segments, + segments + index, + ], + [bottom_center, (index + 1) % segments, index], + [top_center, segments + index, segments + (index + 1) % segments], + ) + ] + ) + else: + rings = 8 + phi = np.linspace(0.0, math.pi, rings + 1) + vertices = np.asarray( + [ + [ + radius * math.sin(p) * math.cos(a), + radius * math.sin(p) * math.sin(a), + radius * math.cos(p), + ] + for p in phi + for a in angles + ] + ) + triangles = np.asarray( + [ + triangle + for ring in range(rings) + for index in range(segments) + for triangle in ( + [ + ring * segments + index, + ring * segments + (index + 1) % segments, + (ring + 1) * segments + index, + ], + [ + ring * segments + (index + 1) % segments, + (ring + 1) * segments + (index + 1) % segments, + (ring + 1) * segments + index, + ], + ) + ] + ) + transformed = vertices @ orientation.to_rotation_matrix().T + return transformed + np.asarray(center.as_tuple), triangles + + +class PickNPlaceConfig(ModuleConfig): + """Configuration for PickNPlaceModule.""" + + align_grasp_yaw: bool = False + grasp: Literal["obb_center", "graspgenx"] = Field( + default="obb_center", validation_alias=AliasChoices("grasp", "grasp_strategy") + ) + graspgenx_pregrasp_offset: float = 0.10 + graspgenx_ik_filter_limit: int = 10 + grasp_empty_closed_threshold: float = 0.01 + grasp_feedback_delay: float = 0.5 + + +class PickNPlaceModule(Module): + """Provide request-driven perception and target selection for pick and place.""" + + config: PickNPlaceConfig + _scene: ObjectSceneRegistrationSpec + _grasp_generator: GraspGenSpec | None + _grasp_filter: GraspCandidateFilterSpec + _pick_execution: PickExecutionSpec + _obstacle_world: ObstacleWorldSpec + _visualization: ManipulationVisualizationSpec + objects: In[list[DetObject]] + camera_info: In[CameraInfo] + basic_grasp_overlay: Out[Image] + graspgenx_candidates: Out[GraspCandidateArray] + + def __init__(self, **kwargs: object) -> None: + super().__init__(**kwargs) + self._objects_condition = threading.Condition() + self._latest_objects: tuple[DetObject, ...] = () + self._objects_version = 0 + self._camera_info: CameraInfo | None = None + self._goal_pose: PoseStamped | None = None + self._pre_grasp_pose: PoseStamped | None = None + self._grasp_candidates: GraspCandidateArray | None = None + self._selected_object: DetObject | None = None + self._held_object_size: Vector3 | None = None + self._tabletop_z: float | None = None + self._open_box: dict[str, float] | None = None + self._scene_geometry_ids: set[str] = set() + + @rpc + def start(self) -> None: + super().start() + self.objects.subscribe(self._on_objects) + self.camera_info.subscribe(self._on_camera_info) + + def _on_objects(self, objects: list[DetObject]) -> None: + with self._objects_condition: + self._latest_objects = tuple(objects) + self._objects_version += 1 + self._objects_condition.notify_all() + + def _on_camera_info(self, camera_info: CameraInfo) -> None: + with self._objects_condition: + self._camera_info = camera_info + + @rpc + def scan_scene( + self, prompt: str | None = None, prompts: list[str] | None = None + ) -> Detection3DArray: + """Run one RGB-D detection pass, optionally targeting one or more text prompts.""" + if prompt is not None and prompts is not None: + raise ValueError("Specify either prompt or prompts, not both") + with self._objects_condition: + objects_version = self._objects_version + if prompts is not None: + self._scene.set_prompts(prompts) + elif prompt: + self._scene.set_prompts([prompt]) + detections = self._scene.scan_scene() + with self._objects_condition: + received_result = self._objects_condition.wait_for( + lambda: self._objects_version > objects_version, + timeout=5.0, + ) + objects = self._latest_objects + if received_result: + # Stream delivery crosses process boundaries and can lag the OSR RPC response. + # Return the same snapshot used by the object/grasp APIs, not the prior response. + return to_detection3d_array( + list(objects), + frame_id=objects[0].frame_id if objects else detections.frame_id, + ts=objects[0].ts if objects else detections.ts, + ) + return detections + + @skill(uses=[CAP_PERCEPTION]) + def scan(self, prompt: str) -> SkillResult: + """Detect a prompted object from one RGB-D frame without moving the robot. + + Returns numbered objects. Use a returned number with ``select_object`` to create a grasp target + or ``get_object_geometry`` to inspect a container target. + """ + if not prompt.strip(): + return SkillResult.fail("INVALID_INPUT", "A nonempty object prompt is required") + try: + detections = self.scan_scene(prompt) + except RuntimeError as exc: + return SkillResult.fail("PERCEPTION_FAILED", str(exc)) + self._publish_scene_objects() + return SkillResult.ok( + f"Detected {detections.detections_length} object(s)", objects=self.get_scene_info() + ) + + @skill(uses=[CAP_PERCEPTION]) + def scan_objects(self, object_names: list[str]) -> SkillResult: + """Detect instances of simple object names from one RGB-D frame. + + Pass one short noun phrase per item, for example ``["wooden block", "white box"]``. Each name is + an independent Moondream query and can return multiple instances. Do not pass instructions, + exclusions, counting requests, or full sentences as object names. + """ + names = [name.strip() for name in object_names if name.strip()] + if not names: + return SkillResult.fail("INVALID_INPUT", "At least one simple object name is required") + try: + detections = self.scan_scene(prompts=names) + except RuntimeError as exc: + return SkillResult.fail("PERCEPTION_FAILED", str(exc)) + self._publish_scene_objects() + return SkillResult.ok( + f"Detected {detections.detections_length} object(s)", + queried_names=names, + objects=self.get_scene_info(), + ) + + @rpc + def get_scene_info(self) -> list[dict[str, object]]: + """Return the number, name, and confidence for current detections.""" + with self._objects_condition: + objects = self._latest_objects + return [ + { + "number": number, + "name": obj.name, + "confidence": obj.confidence, + } + for number, obj in enumerate(objects, 1) + ] + + @skill + def describe_scene(self, question: str = "What objects are visible on the table?") -> str: + """Answer an open-ended question about the latest camera image without moving the robot. + + Requires ``osr.det=moondream`` and is descriptive only; use ``scan`` for numbered 3D objects. + """ + return self._scene.describe_scene(question) + + @skill + def get_object_geometry(self, number: int) -> dict[str, object] | None: + """Return a scanned object's center and OBB size without moving the robot. + + ``number`` must come from the latest ``scan`` result. ``center`` and ``size`` are ``[x, y, z]`` + lists in meters in the returned planning frame; use this to derive a container placement target. + """ + with self._objects_condition: + if number < 1 or number > len(self._latest_objects): + return None + obj = self._latest_objects[number - 1] + return { + "number": number, + "name": obj.name, + "frame_id": obj.frame_id, + "center": [obj.center.x, obj.center.y, obj.center.z], + "size": [obj.size.x, obj.size.y, obj.size.z], + } + + @skill + def install_object_obstacle( + self, number: int, shape: Literal["box", "sphere", "cylinder"] = "box" + ) -> SkillResult: + """Install one measured object as a planner obstacle and render the same primitive in Viser. + + ``number`` must come from the latest ``scan`` result. Choose ``box`` for rectangular objects, + ``cylinder`` for upright round objects, and ``sphere`` only for near-spherical objects. + """ + obj = self._object_for_number(number) + if obj is None: + return SkillResult.fail("INVALID_INPUT", f"No detected object numbered {number}") + orientation = self._upright_orientation(obj) + if shape == "box": + dimensions = (obj.size.x, obj.size.y, obj.size.z) + elif shape == "sphere": + dimensions = (max(obj.size.x, obj.size.y, obj.size.z) / 2.0,) + else: + dimensions = (max(obj.size.x, obj.size.y) / 2.0, obj.size.z) + name = f"scene-object-{number}" + center = Vector3(obj.center) + if not self._install_geometry(name, center, orientation, shape, dimensions): + return SkillResult.fail("EXECUTION_FAILED", f"Failed to install obstacle '{name}'") + return SkillResult.ok( + "Obstacle installed", + name=name, + shape=shape, + center=[center.x, center.y, center.z], + dimensions=list(dimensions), + ) + + @skill + def install_open_box(self, number: int, wall_thickness: float = 0.01) -> SkillResult: + """Measure an open rectangular box and render it as a display-only solid box in Viser. + + Call ``estimate_table`` first. The result describes the free opening for top-down placement, but + does not add box walls to the planning world. + """ + if wall_thickness <= 0.0: + return SkillResult.fail("INVALID_INPUT", "wall_thickness must be positive") + if self._tabletop_z is None: + return SkillResult.fail( + "INVALID_STATE", "Estimate the table before modeling an open box" + ) + obj = self._object_for_number(number) + if obj is None: + return SkillResult.fail("INVALID_INPUT", f"No detected object numbered {number}") + width, depth = obj.size.x, obj.size.y + if width <= 2.0 * wall_thickness or depth <= 2.0 * wall_thickness: + return SkillResult.fail( + "INVALID_INPUT", "Box opening is smaller than twice wall_thickness" + ) + points = obj.pointcloud.points_f32() + rim_z = ( + float(np.quantile(points[:, 2], 0.95)) if len(points) else obj.center.z + obj.size.z / 2 + ) + height = rim_z - self._tabletop_z + if height <= 0.0: + return SkillResult.fail( + "PERCEPTION_FAILED", "Box rim is not above the estimated tabletop" + ) + center = Vector3(obj.center.x, obj.center.y, self._tabletop_z + height / 2.0) + orientation = self._upright_orientation(obj) + + vertices, triangles = _primitive_mesh("box", center, (width, depth, height), orientation) + self._visualization.set_visualization_layer( + VisualizationLayer( + "picknplace/open-box", + "world", + ( + MeshElement( + "box-envelope", + vertices, + triangles, + color=np.asarray([230, 230, 230]), + opacity=0.25, + ), + ), + ) + ) + self._open_box = { + "center_x": center.x, + "center_y": center.y, + "tabletop_z": self._tabletop_z, + "rim_z": rim_z, + "opening_width": width - 2.0 * wall_thickness, + "opening_depth": depth - 2.0 * wall_thickness, + } + return SkillResult.ok( + "Open box measured and displayed", + center=[center.x, center.y], + rim_z=rim_z, + opening_width=width - 2.0 * wall_thickness, + opening_depth=depth - 2.0 * wall_thickness, + ) + + @skill + def clear_scene_geometry(self) -> SkillResult: + """Remove temporary scene obstacles and the display-only open-box marker.""" + removed = [ + geometry_id + for geometry_id in tuple(self._scene_geometry_ids) + if self._obstacle_world.remove_obstacle(geometry_id) + ] + self._scene_geometry_ids.difference_update(removed) + self._visualization.set_visualization_layer( + VisualizationLayer("picknplace/open-box", "world", ()) + ) + self._open_box = None + return SkillResult.ok("Temporary scene geometry cleared", removed=removed) + + @rpc + def get_goal_pose(self, number: int) -> PoseStamped | None: + """Select an object and return its downward-facing, floor-clamped grasp goal.""" + selection = self._basic_grasp(number) + if selection is None: + return None + grasp, obj = selection + if self.config.grasp == "graspgenx": + if self._grasp_generator is None: + raise RuntimeError("GraspGenX is not configured for this pick-and-place blueprint") + candidates = self._grasp_generator.propose_grasps(obj.pointcloud) + self._selected_object = obj + self._grasp_candidates = self._filter_graspgenx_candidates(candidates) + if not self._grasp_candidates.candidates: + self.graspgenx_candidates.publish(self._grasp_candidates) + return None + return self._select_graspgenx_candidate(0) + yaw = self._grasp_yaw(obj) if self.config.align_grasp_yaw else 0.0 + pick_execution = getattr(self, "_pick_execution", None) + current_pose = pick_execution.get_ee_pose() if pick_execution is not None else None + if current_pose is not None: + yaw = self._closest_parallel_jaw_yaw(yaw, current_pose.orientation.to_euler().z) + self._grasp_candidates = None + self._selected_object = obj + self.graspgenx_candidates.publish(GraspCandidateArray()) + grasp_z = _table_midpoint_grasp_z( + obj.pointcloud.points_f32(), self._tabletop_z, grasp.position.z + ) + self._goal_pose = PoseStamped( + ts=grasp.ts, + frame_id=grasp.frame_id, + position=Vector3(grasp.position.x, grasp.position.y, max(grasp_z, 0.100)), + orientation=Quaternion.from_euler(Vector3(-math.pi, 0.0, yaw)), + ) + self._pre_grasp_pose = None + return self._goal_pose + + @skill + def select_object(self, number: int) -> SkillResult: + """Select a scanned object and return grasp and pre-grasp targets without moving the robot. + + ``number`` must come from the latest ``scan`` result. Returned target values are XYZ in meters and + roll/pitch/yaw in radians. Move to ``pre_grasp`` first, then move to ``goal`` for gripper contact; + ``pre_grasp`` is 100 mm above the object and is not a grasp pose. + """ + goal = self.get_goal_pose(number) + if goal is None: + return SkillResult.fail("INVALID_INPUT", f"No selectable object numbered {number}") + pre_grasp = self.get_pre_grasp_pose() + if pre_grasp is None: + return SkillResult.fail("INVALID_STATE", "Selected object has no pre-grasp target") + + def pose_target(pose: PoseStamped) -> dict[str, float]: + euler = pose.orientation.to_euler() + return { + "x": pose.position.x, + "y": pose.position.y, + "z": pose.position.z, + "roll": euler.x, + "pitch": euler.y, + "yaw": euler.z, + } + + return SkillResult.ok( + "Object selected. Move to pre_grasp, then goal for gripper contact before closing the gripper.", + goal=pose_target(goal), + pre_grasp=pose_target(pre_grasp), + ) + + @skill(uses=[CAP_MOVEMENT]) + def pick_selected(self, robot_name: str | None = None) -> SkillResult: + """Pick the object most recently selected with ``select_object``. + + Executes the full pre-grasp, contact-grasp, close, feedback verification, and retreat sequence. + A gripper position at or below the empty-closed threshold means no object was picked up and returns + ``GRASP_VERIFICATION_FAILED``. Call ``select_object`` before this tool; do not manually recreate + the grasp sequence with individual motion and gripper tools. + """ + goal = self._goal_pose + pre_grasp = self._pre_grasp_pose + if goal is None or pre_grasp is None: + return SkillResult.fail("INVALID_STATE", "Select an object before starting a pick") + + def move(pose: PoseStamped) -> SkillResult: + euler = pose.orientation.to_euler() + return self._pick_execution.move_to_pose( + pose.position.x, + pose.position.y, + pose.position.z, + euler.x, + euler.y, + euler.z, + robot_name, + ) + + opened = self._pick_execution.open_gripper(robot_name) + if not opened.is_success(): + return opened + approach = move(pre_grasp) + if not approach.is_success(): + return approach + contact = move(goal) + if not contact.is_success(): + return contact + closed = self._pick_execution.close_gripper(robot_name) + if not closed.is_success(): + return closed + + time.sleep(self.config.grasp_feedback_delay) + gripper_position = self._pick_execution.get_gripper(robot_name) + if gripper_position is None: + return SkillResult.fail( + "GRIPPER_FAILED", "Cannot verify pickup: gripper feedback unavailable" + ) + if gripper_position <= self.config.grasp_empty_closed_threshold: + self._pick_execution.open_gripper(robot_name) + recovery = move(pre_grasp) + result = SkillResult.fail( + "GRASP_VERIFICATION_FAILED", + "Pickup failed: gripper reached the empty-closed position; rescan and select before retrying", + ) + result.metadata = { + "gripper_position": gripper_position, + "rescan_required": True, + "recovered_to_pre_grasp": recovery.is_success(), + } + return result + + retreat = move(pre_grasp) + if not retreat.is_success(): + return retreat + selected_object = self._selected_object + if selected_object is None: + return SkillResult.fail( + "INVALID_STATE", "Selected object details are unavailable after grasp" + ) + self._held_object_size = Vector3(selected_object.size) + return SkillResult.ok( + "Pick complete: grasp verified and object retreated from the table", + gripper_position=gripper_position, + ) + + @skill(uses=[CAP_MOVEMENT]) + def place_selected(self, robot_name: str | None = None) -> SkillResult: + """Drop the verified held object into the most recently measured open box. + + Call ``install_open_box`` for the destination and complete ``pick_selected`` first. This tool moves + above the remembered opening and releases above the rim. It first lifts the held object for transit, + then lowers only at the box center. It never lowers the end effector into the box. The box remains + display-only and does not add planner collisions. + """ + box = self._open_box + held_size = self._held_object_size + if box is None: + return SkillResult.fail( + "INVALID_STATE", "Measure the destination with install_open_box first" + ) + if held_size is None: + return SkillResult.fail( + "INVALID_STATE", "No verified held object is available to place" + ) + if held_size.x > box["opening_width"] or held_size.y > box["opening_depth"]: + return SkillResult.fail( + "INVALID_INPUT", "Held object does not fit inside the measured box opening" + ) + + def move(x: float, y: float, z: float) -> SkillResult: + return self._pick_execution.move_to_pose(x, y, z, robot_name=robot_name) + + # The held object's bottom remains above the rim throughout lateral travel. + drop_z = box["rim_z"] + held_size.z / 2.0 + 0.02 + transit_z = drop_z + 0.10 + current_pose = self._pick_execution.get_ee_pose(robot_name) + if current_pose is not None and current_pose.position.z < transit_z: + lift = move(current_pose.position.x, current_pose.position.y, transit_z) + if not lift.is_success(): + return lift + approach = move(box["center_x"], box["center_y"], transit_z) + if not approach.is_success(): + return approach + lower = move(box["center_x"], box["center_y"], drop_z) + if not lower.is_success(): + return lower + opened = self._pick_execution.open_gripper(robot_name) + if not opened.is_success(): + return opened + self._held_object_size = None + return SkillResult.ok( + "Drop complete: object released above the measured box opening", + drop_z=drop_z, + object_bottom_clearance=0.02, + ) + + @rpc + def select_grasp_candidate(self, rank: int) -> PoseStamped | None: + """Select one ranked GraspGenX proposal as the goal and Rerun highlight.""" + return self._select_graspgenx_candidate(rank) + + def _select_graspgenx_candidate(self, rank: int) -> PoseStamped | None: + candidates = self._grasp_candidates + if candidates is None or rank < 0 or rank >= len(candidates.candidates): + return None + candidates.selected_index = rank + self.graspgenx_candidates.publish(candidates) + candidate = candidates.candidates[rank] + self._goal_pose = PoseStamped( + ts=candidates.header.timestamp, + frame_id=candidates.header.frame_id, + position=candidate.pose.position, + orientation=candidate.pose.orientation, + ) + self._pre_grasp_pose = None + return self._goal_pose + + def _filter_graspgenx_candidates(self, candidates: GraspCandidateArray) -> GraspCandidateArray: + """Keep only top-ranked proposals whose TCP IK is collision-free in the live world.""" + accepted = [] + for candidate in candidates.candidates[: self.config.graspgenx_ik_filter_limit]: + result = self._grasp_filter.inverse_kinematics_single( + candidate.pose, "arm", check_collision=True + ) + if result.is_success(): + accepted.append(candidate) + return GraspCandidateArray(candidates.header, accepted) + + @rpc + def get_pre_grasp_pose(self) -> PoseStamped | None: + """Return the selected goal offset 100 mm opposite its final approach direction.""" + if self._goal_pose is None: + return None + if self.config.grasp == "graspgenx": + offset = self._goal_pose.orientation.rotate_vector( + # GraspGenX local +Z points in the direction of the final + # approach. A pre-grasp retreats along the opposite axis. + Vector3(0.0, 0.0, -self.config.graspgenx_pregrasp_offset) + ) + else: + offset = Vector3(0.0, 0.0, 0.100) + self._pre_grasp_pose = PoseStamped( + ts=self._goal_pose.ts, + frame_id=self._goal_pose.frame_id, + position=Vector3( + self._goal_pose.position.x + offset.x, + self._goal_pose.position.y + offset.y, + self._goal_pose.position.z + offset.z, + ), + orientation=self._goal_pose.orientation, + ) + self._publish_viser_selection() + return self._pre_grasp_pose + + def _publish_viser_selection(self) -> None: + """Show the selected object and TCP targets without mutating the planning scene.""" + obj = self._selected_object + goal = self._goal_pose + pre_grasp = self._pre_grasp_pose + if obj is None or goal is None or pre_grasp is None: + return + points = obj.pointcloud.points_f32() + if len(points) == 0: + return + cloud_colors = np.repeat(np.array([[255, 190, 70]], dtype=np.uint8), len(points), axis=0) + vertices: list[np.ndarray] = [] + edges: list[list[int]] = [] + colors: list[list[int]] = [] + for pose, color in ((goal, [255, 70, 70]), (pre_grasp, [70, 255, 120])): + start = len(vertices) + origin = np.asarray(pose.position.as_tuple, dtype=np.float32) + axes = pose.orientation.to_rotation_matrix().astype(np.float32) * 0.06 + vertices.extend((origin, origin + axes[:, 0], origin + axes[:, 1], origin + axes[:, 2])) + edges.extend(((start, start + 1), (start, start + 2), (start, start + 3))) + colors.extend((color, color, color)) + self._visualization.set_visualization_layer( + VisualizationLayer( + "picknplace/selection", + "world", + ( + PointCloudElement("object", points, cloud_colors, point_size=0.003), + LineSetElement( + "tcp-targets", + np.asarray(vertices), + np.asarray(edges), + np.asarray(colors), + ), + ), + ) + ) + + def _publish_scene_objects(self) -> None: + """Display the latest measured object envelopes without affecting planning.""" + visualization = getattr(self, "_visualization", None) + if visualization is None: + return + with self._objects_condition: + objects = tuple(self._latest_objects) + colors = ([255, 180, 70], [80, 180, 255], [130, 230, 130], [230, 150, 230]) + elements: list[MeshElement] = [] + for number, obj in enumerate(objects, 1): + dimensions = (obj.size.x, obj.size.y, obj.size.z) + if any(dimension <= 0.0 for dimension in dimensions): + continue + vertices, triangles = _primitive_mesh( + "box", Vector3(obj.center), dimensions, self._upright_orientation(obj) + ) + elements.append( + MeshElement( + f"object-{number}", + vertices, + triangles, + color=np.asarray(colors[(number - 1) % len(colors)]), + opacity=0.30, + ) + ) + visualization.set_visualization_layer( + VisualizationLayer("picknplace/scene-objects", "world", tuple(elements)) + ) + + @rpc + def get_grasp_candidates(self) -> GraspCandidateArray: + """Return the GraspGenX proposals generated for the selected object.""" + return self._grasp_candidates or GraspCandidateArray() + + @rpc + def estimate_table_surface(self) -> dict[str, float] | None: + """Estimate a horizontal tabletop from the latest full RGB-D scene cloud.""" + scene = self._scene.get_full_scene_pointcloud(voxel_size=0.01) + if scene is None: + return None + estimate = _estimate_table_surface(scene.points_f32()) + if estimate is None: + return None + self._tabletop_z = estimate["tabletop_z"] + z = estimate["tabletop_z"] + half_width = estimate["width"] / 2 + half_depth = estimate["depth"] / 2 + x = estimate["center_x"] + y = estimate["center_y"] + vertices = np.asarray( + [ + [x - half_width, y - half_depth, z], + [x + half_width, y - half_depth, z], + [x + half_width, y + half_depth, z], + [x - half_width, y + half_depth, z], + ] + ) + self._visualization.set_visualization_layer( + VisualizationLayer( + "picknplace/table-estimate", + "world", + ( + MeshElement( + "tabletop-fill", + vertices, + np.asarray([[0, 1, 2], [0, 2, 3]]), + color=np.asarray([80, 180, 255]), + opacity=1.0, + ), + LineSetElement( + "tabletop", + vertices, + np.asarray([[0, 1], [1, 2], [2, 3], [3, 0]]), + colors=np.asarray([[80, 180, 255]] * 4), + line_width=2.0, + ), + ), + ) + ) + return estimate + + @skill(uses=[CAP_PERCEPTION]) + def estimate_table(self) -> SkillResult: + """Run a fresh RGB-D scan, then estimate the tabletop without moving the robot. + + Pass the returned ``center_x``, ``center_y``, ``tabletop_z``, ``width``, and ``depth`` directly to + ``set_table_collision`` before requesting motion near the table. Do not call this concurrently with + ``scan``; both tools exclusively use the perception pipeline. + """ + try: + self.scan_scene() + except RuntimeError as exc: + return SkillResult.fail("PERCEPTION_FAILED", str(exc)) + estimate = self.estimate_table_surface() + if estimate is None: + return SkillResult.fail( + "PERCEPTION_FAILED", "No horizontal tabletop estimate is available" + ) + return SkillResult.ok("Table estimated", **estimate) + + def _basic_grasp(self, number: int) -> tuple[PoseStamped, DetObject] | None: + """Return the selected cloud's OBB-center grasp frame and object geometry.""" + with self._objects_condition: + if number < 1 or number > len(self._latest_objects): + return None + obj = self._latest_objects[number - 1] + camera_info = self._camera_info + grasp = PoseStamped( + ts=obj.ts, + frame_id=obj.frame_id, + position=obj.center, + orientation=obj.pose.orientation, + ) + if camera_info is not None and obj.camera_transform is not None and obj.image is not None: + if overlay := draw_pose_axes( + obj.image, grasp, obj.camera_transform.inverse(), camera_info + ): + self.basic_grasp_overlay.publish(overlay) + return grasp, obj + + def _object_for_number(self, number: int) -> DetObject | None: + with self._objects_condition: + if number < 1 or number > len(self._latest_objects): + return None + return self._latest_objects[number - 1] + + @staticmethod + def _upright_orientation(obj: DetObject) -> Quaternion: + """Keep measured horizontal yaw while constraining scene primitives upright.""" + return Quaternion.from_euler(Vector3(0.0, 0.0, obj.pose.orientation.to_euler().z)) + + def _install_geometry( + self, + name: str, + center: Vector3, + orientation: Quaternion, + shape: Literal["box", "sphere", "cylinder"], + dimensions: tuple[float, ...], + ) -> bool: + pose = Pose(center, orientation) + if self._obstacle_world.update_obstacle(name, pose, shape, list(dimensions)): + self._scene_geometry_ids.add(name) + return True + obstacle_id = self._obstacle_world.add_obstacle(name, pose, shape, list(dimensions)) + if obstacle_id: + self._scene_geometry_ids.add(obstacle_id) + return True + return False + + @staticmethod + def _grasp_yaw(obj: DetObject) -> float: + """Align the gripper's local Y closing axis with the narrowest horizontal OBB axis.""" + rotation = obj.pose.orientation.to_rotation_matrix() + extents = (obj.size.x, obj.size.y, obj.size.z) + horizontal_axes = sorted(range(3), key=lambda axis: abs(rotation[2, axis]))[:2] + narrow_axis = min(horizontal_axes, key=lambda axis: extents[axis]) + return math.atan2(rotation[1, narrow_axis], rotation[0, narrow_axis]) - math.pi / 2 + + @staticmethod + def _closest_parallel_jaw_yaw(target_yaw: float, current_yaw: float) -> float: + """Choose the equivalent parallel-jaw yaw requiring the smallest wrist rotation.""" + return target_yaw + math.pi * round((current_yaw - target_yaw) / math.pi) diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index b14000496d..34a5a58523 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -35,6 +35,7 @@ VisualizationStateFrame, ) from dimos.manipulation.planning.spec.protocols import VisualizationSpec +from dimos.manipulation.visualization.layers import VisualizationLayer from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Vector3 import Vector3 @@ -220,6 +221,12 @@ def remove_vis_obstacle(self, obstacle_id: str) -> None: def clear_vis_obstacles(self) -> None: self.calls.append(("clear_vis_obstacles",)) + def set_layer(self, layer: VisualizationLayer) -> None: + self.calls.append(("set_layer", layer)) + + def clear_layer(self, layer_id: str) -> None: + self.calls.append(("clear_layer", layer_id)) + def _robot_config() -> RobotModelConfig: return RobotModelConfig( diff --git a/dimos/manipulation/planning/monitor/test_world_obstacle_suppression.py b/dimos/manipulation/planning/monitor/test_world_obstacle_suppression.py new file mode 100644 index 0000000000..995a137931 --- /dev/null +++ b/dimos/manipulation/planning/monitor/test_world_obstacle_suppression.py @@ -0,0 +1,178 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import threading +import time +from types import SimpleNamespace + +import open3d as o3d +import pytest +from pytest_mock import MockerFixture + +from dimos.manipulation.planning.monitor.world_obstacle_monitor import WorldObstacleMonitor +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.perception.experimental.object import Object + + +def _object(object_id: str) -> Object: + return Object( + name=f"object-{object_id}", + object_id=object_id, + center=Vector3(0.4, 0.0, 0.2), + size=Vector3(0.05, 0.05, 0.1), + pose=PoseStamped(), + pointcloud=PointCloud2(o3d.geometry.PointCloud()), + bbox=(0.0, 0.0, 1.0, 1.0), + track_id=0, + class_id=0, + confidence=1.0, + ts=time.time(), + image=Image(), + ) + + +def _monitor(mocker: MockerFixture) -> tuple[WorldObstacleMonitor, SimpleNamespace]: + parent = SimpleNamespace( + _lock=threading.RLock(), + add_obstacle=mocker.Mock( + side_effect=lambda obstacle: f"world-{obstacle.name}-{time.monotonic_ns()}" + ), + remove_obstacle=mocker.Mock(return_value=True), + ) + monitor = WorldObstacleMonitor(parent) # type: ignore[arg-type] + monitor.start() + return monitor, parent + + +def test_suppression_skips_target_but_refreshes_other_objects( + mocker: MockerFixture, +) -> None: + monitor, parent = _monitor(mocker) + target = _object("target") + other = _object("other") + monitor.on_objects([target, other]) + monitor.refresh_obstacles() + parent.add_obstacle.reset_mock() + + with monitor.suppress_object_obstacle("target") as suppression: + refreshed = monitor.refresh_obstacles() + + assert suppression.removed is True + assert [item["object_id"] for item in refreshed] == ["other"] + assert set(monitor._object_obstacles) == {"other"} + + assert set(monitor._object_obstacles) == {"target", "other"} + assert parent.remove_obstacle.call_count >= 1 + + +def test_suppression_wins_race_with_in_progress_refresh(mocker: MockerFixture) -> None: + monitor, _ = _monitor(mocker) + monitor.on_objects([_object("target"), _object("other")]) + conversion_started = threading.Event() + continue_conversion = threading.Event() + original_conversion = monitor._object_to_obstacle + + def delayed_conversion(obj: Object): + if obj.object_id == "target": + conversion_started.set() + assert continue_conversion.wait(timeout=1.0) + return original_conversion(obj) + + mocker.patch.object(monitor, "_object_to_obstacle", side_effect=delayed_conversion) + refreshed: list[list[dict[str, object]]] = [] + thread = threading.Thread(target=lambda: refreshed.append(monitor.refresh_obstacles())) + thread.start() + assert conversion_started.wait(timeout=1.0) + + with monitor.suppress_object_obstacle("target"): + continue_conversion.set() + thread.join(timeout=1.0) + + assert not thread.is_alive() + assert [item["object_id"] for item in refreshed[0]] == ["other"] + assert set(monitor._object_obstacles) == {"other"} + + assert set(monitor._object_obstacles) == {"target", "other"} + + +def test_nested_suppression_removes_and_restores_once(mocker: MockerFixture) -> None: + monitor, parent = _monitor(mocker) + monitor.on_objects([_object("target")]) + monitor.refresh_obstacles() + parent.add_obstacle.reset_mock() + parent.remove_obstacle.reset_mock() + + with monitor.suppress_object_obstacle("target"): + with monitor.suppress_object_obstacle("target"): + assert "target" not in monitor._object_obstacles + assert "target" not in monitor._object_obstacles + + assert parent.remove_obstacle.call_count == 1 + assert parent.add_obstacle.call_count == 1 + assert "target" in monitor._object_obstacles + + +def test_suppression_restores_after_cancellation(mocker: MockerFixture) -> None: + class Cancelled(BaseException): + pass + + monitor, _ = _monitor(mocker) + monitor.on_objects([_object("target")]) + monitor.refresh_obstacles() + + with pytest.raises(Cancelled): + with monitor.suppress_object_obstacle("target"): + raise Cancelled + + assert monitor._object_suppressions == {} + assert "target" in monitor._object_obstacles + + +def test_suppression_reports_restore_failure_without_masking_body( + mocker: MockerFixture, +) -> None: + monitor, parent = _monitor(mocker) + monitor.on_objects([_object("target")]) + monitor.refresh_obstacles() + parent.add_obstacle.side_effect = None + parent.add_obstacle.return_value = "" + + with monitor.suppress_object_obstacle("target") as suppression: + body_completed = True + + assert body_completed is True + assert suppression.cleanup_error == "failed to restore obstacle for object 'target'" + assert "target" not in monitor._object_obstacles + + +def test_failed_suppression_removal_restores_internal_tracking( + mocker: MockerFixture, +) -> None: + monitor, parent = _monitor(mocker) + monitor.on_objects([_object("target")]) + monitor.refresh_obstacles() + parent.remove_obstacle.return_value = False + + with pytest.raises(RuntimeError, match="failed to suppress") as exc_info: + with monitor.suppress_object_obstacle("target"): + raise AssertionError("suppression body must not run") + + assert str(exc_info.value) == "failed to suppress obstacle for object 'target'" + assert monitor._object_suppressions == {} + assert "target" in monitor._object_obstacles diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index bcbb2b7ed8..dc1c616230 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -16,7 +16,7 @@ from __future__ import annotations -from collections.abc import Sequence +from collections.abc import Iterator, Sequence from contextlib import contextmanager import threading from typing import TYPE_CHECKING, Any @@ -29,7 +29,10 @@ from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry from dimos.manipulation.planning.groups.utils import filter_joint_state_to_selected_joints from dimos.manipulation.planning.monitor.robot_state_monitor import RobotStateMonitor -from dimos.manipulation.planning.monitor.world_obstacle_monitor import WorldObstacleMonitor +from dimos.manipulation.planning.monitor.world_obstacle_monitor import ( + ObjectObstacleSuppression, + WorldObstacleMonitor, +) from dimos.manipulation.planning.spec.models import ( PlanningSceneInfo, VisualizationSession, @@ -313,6 +316,15 @@ def remove_object_obstacle(self, object_id: str) -> bool: return self._obstacle_monitor.remove_object_obstacle(object_id) return False + @contextmanager + def suppress_object_obstacle(self, object_id: str) -> Iterator[ObjectObstacleSuppression]: + """Temporarily exclude one perception object from collision checking.""" + if self._obstacle_monitor is None: + yield ObjectObstacleSuppression(object_id=object_id) + return + with self._obstacle_monitor.suppress_object_obstacle(object_id) as suppression: + yield suppression + def clear_perception_obstacles(self) -> int: """Remove all perception obstacles. Returns count removed.""" if self._obstacle_monitor is not None: diff --git a/dimos/manipulation/planning/monitor/world_obstacle_monitor.py b/dimos/manipulation/planning/monitor/world_obstacle_monitor.py index 481e6976ca..6f870b171c 100644 --- a/dimos/manipulation/planning/monitor/world_obstacle_monitor.py +++ b/dimos/manipulation/planning/monitor/world_obstacle_monitor.py @@ -26,7 +26,10 @@ from __future__ import annotations -from dataclasses import replace +from collections import Counter +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, replace import time from typing import TYPE_CHECKING, Any @@ -49,6 +52,15 @@ logger = setup_logger() +@dataclass +class ObjectObstacleSuppression: + """Result of a scoped object-obstacle suppression.""" + + object_id: str + removed: bool = False + cleanup_error: str | None = None + + class WorldObstacleMonitor: """Monitors world obstacles and updates its parent WorldMonitor. @@ -95,6 +107,7 @@ def __init__( self._object_cache: dict[str, tuple[Object, float, float]] = {} # object_id -> obstacle_id (objects currently added to Drake world) self._object_obstacles: dict[str, str] = {} + self._object_suppressions: Counter[str] = Counter() # Running state self._running = False @@ -123,6 +136,7 @@ def clear_tracking(self) -> None: self._perception_objects.clear() self._perception_timestamps.clear() self._object_obstacles.clear() + self._object_suppressions.clear() def on_collision_object(self, msg: CollisionObjectMessage) -> None: """Handle explicit collision object message. @@ -497,6 +511,8 @@ def refresh_obstacles(self, min_duration: float = 0.0) -> list[dict[str, Any]]: for oid, (obj, first_seen, last_seen) in self._object_cache.items(): if not isinstance(obj, Object): continue + if self._object_suppressions[oid] > 0: + continue if last_seen - first_seen < min_duration: continue eligible.append((oid, obj)) @@ -517,6 +533,10 @@ def refresh_obstacles(self, min_duration: float = 0.0) -> list[dict[str, Any]]: result: list[dict[str, Any]] = [] for oid, obj, obstacle in prepared: + # Suppression may have started while obstacle geometry was + # computed outside the lock. + if self._object_suppressions[oid] > 0: + continue assert isinstance(obj, Object) obs_id = self._parent.add_obstacle(obstacle) if not obs_id: @@ -552,6 +572,58 @@ def remove_object_obstacle(self, object_id: str) -> bool: logger.info(f"Removed obstacle for object '{object_id}'") return True + @contextmanager + def suppress_object_obstacle(self, object_id: str) -> Iterator[ObjectObstacleSuppression]: + """Exclude one cached object obstacle for the lifetime of the context. + + Nested callers share one removal. Live refreshes skip suppressed object + IDs, and the outermost exit restores the latest cached geometry. + """ + handle = ObjectObstacleSuppression(object_id=object_id) + with self._lock: + depth = self._object_suppressions[object_id] + self._object_suppressions[object_id] = depth + 1 + if depth == 0: + obstacle_id = self._object_obstacles.get(object_id) + if obstacle_id is not None: + if not self._parent.remove_obstacle(obstacle_id): + del self._object_suppressions[object_id] + raise RuntimeError(f"failed to suppress obstacle for object '{object_id}'") + del self._object_obstacles[object_id] + handle.removed = True + try: + yield handle + finally: + self._release_object_suppression(handle) + + def _release_object_suppression(self, handle: ObjectObstacleSuppression) -> None: + object_id = handle.object_id + cached: Object | None = None + with self._lock: + depth = self._object_suppressions.get(object_id, 0) + if depth > 1: + self._object_suppressions[object_id] = depth - 1 + return + self._object_suppressions.pop(object_id, None) + entry = self._object_cache.get(object_id) + if entry is not None: + cached = entry[0] + + if cached is None: + return + obstacle = self._object_to_obstacle(cached) + with self._lock: + if self._object_suppressions.get(object_id, 0) > 0: + return + if object_id in self._object_obstacles: + return + obstacle_id = self._parent.add_obstacle(obstacle) + if obstacle_id: + self._object_obstacles[object_id] = obstacle_id + return + handle.cleanup_error = f"failed to restore obstacle for object '{object_id}'" + logger.error(handle.cleanup_error) + def clear_perception_obstacles(self) -> int: """Remove all object obstacles from the planning world. diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py index ff33953025..36e663fc7d 100644 --- a/dimos/manipulation/planning/spec/protocols.py +++ b/dimos/manipulation/planning/spec/protocols.py @@ -42,6 +42,7 @@ VisualizationStateFrame, WorldRobotID, ) + from dimos.manipulation.visualization.layers import VisualizationLayer from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory @@ -224,6 +225,14 @@ def clear_vis_obstacles(self) -> None: """Clear obstacle representations from the visualization.""" ... + def set_layer(self, layer: VisualizationLayer) -> None: + """Replace one complete display-only visualization layer.""" + ... + + def clear_layer(self, layer_id: str) -> None: + """Clear one display-only layer while retaining viewer-owned state.""" + ... + def get_visualization_url(self) -> str | None: """Get visualization URL if enabled.""" ... diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index b856379cc7..6c4c64b1b2 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -57,6 +57,7 @@ VisualizationSession, VisualizationStateFrame, ) + from dimos.manipulation.visualization.layers import VisualizationLayer try: from pydrake.geometry import ( @@ -1208,6 +1209,14 @@ def clear_vis_obstacles(self) -> None: """Embedded Meshcat observes native WorldSpec obstacle mutations.""" return None + def set_layer(self, layer: VisualizationLayer) -> None: + """Embedded Meshcat ignores generic display-only layers.""" + return None + + def clear_layer(self, layer_id: str) -> None: + """Embedded Meshcat ignores generic display-only layers.""" + return None + def get_visualization_url(self) -> str | None: """Get visualization URL if enabled.""" if self._meshcat is not None: diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index 3f1580c7af..2361626bd6 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -450,7 +450,7 @@ def plan_joint_path( goal: JointState, timeout: float = 10.0, ) -> PlanningResult: - """Plan using the legacy robot-scoped local-name contract.""" + """Plan between explicit states using the legacy robot-scoped contract.""" if world is not self: return PlanningResult( status=PlanningStatus.NO_SOLUTION, @@ -470,12 +470,6 @@ def plan_joint_path( message="RoboPlan planning scene is not ready: authoritative state is incomplete", ) robot = self._get_robot(robot_id) - current = self._live_context.q_by_robot[robot_id] - if not np.allclose(q_start, current, atol=1e-6, rtol=0.0): - return PlanningResult( - status=PlanningStatus.INVALID_START, - message="Requested start state does not match current scene state", - ) group = self._legacy_group(robot.config.name) return self._plan_group( group, @@ -494,7 +488,7 @@ def plan_selected_joint_path( timeout: float = 10.0, max_iterations: int = 5000, ) -> PlanningResult: - """Plan one or more non-overlapping groups through RoboPlan RRT.""" + """Plan selected groups between explicit states through RoboPlan RRT.""" if world is not self: return PlanningResult( status=PlanningStatus.UNSUPPORTED, diff --git a/dimos/manipulation/pnpconsole.py b/dimos/manipulation/pnpconsole.py new file mode 100644 index 0000000000..1e0d2e6559 --- /dev/null +++ b/dimos/manipulation/pnpconsole.py @@ -0,0 +1,314 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Interactive RPC client for the stepwise ``picknplace`` pipeline. + +Start the blueprint first, then run: + + uv run --no-sync python -m dimos.manipulation.pnpconsole +""" + +from __future__ import annotations + +from pprint import pprint +import time +from typing import Any + +from dimos import Dimos +from dimos.manipulation.planning.planners.config import RoboPlanCartesianPathConfig +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped + + +def _object_number() -> int | None: + value = input("Object number: ").strip() + try: + number = int(value) + except ValueError: + print("Enter a positive whole number.") + return None + if number < 1: + print("Enter a positive whole number.") + return None + return number + + +def _grasp_rank(candidate_count: int) -> int | None: + value = input(f"Grasp rank [0-{candidate_count - 1}, Enter=0]: ").strip() + if not value: + return 0 + try: + rank = int(value) + except ValueError: + print("Enter a whole-number grasp rank.") + return None + if rank < 0 or rank >= candidate_count: + print(f"Choose a rank from 0 to {candidate_count - 1}.") + return None + return rank + + +def _print_pose(pose: Any) -> None: + if pose is None: + print("No pose is available.") + return + if frame_id := getattr(pose, "frame_id", None): + print(f"frame: {frame_id}") + print(f"position: {pose.position.as_tuple}") + print( + "orientation: " + f"({pose.orientation.x}, {pose.orientation.y}, {pose.orientation.z}, {pose.orientation.w})" + ) + + +def _cartesian_waypoints(manipulation: Any, target: Any) -> list[PoseStamped] | None: + current = manipulation.get_ee_pose("arm") + if current is None: + return None + return [ + PoseStamped(frame_id="world", position=current.position, orientation=current.orientation), + PoseStamped(frame_id="world", position=target.position, orientation=target.orientation), + ] + + +def _preview(manipulation: Any) -> None: + print(f"Viser preview: {manipulation.get_visualization_url()}") + print(manipulation.preview_plan(duration=2.0)) + + +def _print_grasp_candidates(candidates: Any) -> None: + if not candidates.candidates: + return + print(f"GraspGenX proposals: {len(candidates.candidates)}") + for rank, candidate in enumerate(candidates.candidates[:10]): + pose = candidate.pose + print( + f"{rank}: score={candidate.score:.3f} position={pose.position.as_tuple} " + f"orientation={pose.orientation.to_tuple()}" + ) + + +def main() -> None: + """Connect to PickNPlaceModule and run one explicit pick-pipeline stage.""" + print("Connecting to PickNPlaceModule...") + app = Dimos.connect() + pnp = app.pnp + manipulation = app.ManipulationModule + goal = None + pre_grasp = None + approach_planned = False + approach_executed = False + descent_planned = False + descent_executed = False + gripper_closed = False + ascent_planned = False + ascent_executed = False + print("Connected. Every planned motion is previewed in Viser before execution.") + + while True: + print("\n1) Scan 2) Info 3) Select target 4) Plan/preview approach") + print("5) Execute approach 6) Plan/preview descent 7) Execute descent 8) Close") + print("9) Plan/preview ascent 10) Execute ascent 11) Open 12) Current EE 13) Go home") + print("14) Scan/estimate/install table collision 15) Grasp + lift now (no preview)") + print("16) Describe current camera scene q) Quit") + choice = input("Select: ").strip().lower() + try: + if choice == "q": + return + if choice == "1": + prompt = input("Object prompt (blank = current detector prompt): ").strip() + detections = pnp.scan_scene(prompt or None) + print(f"Detected {detections.detections_length} object(s).") + elif choice == "2": + pprint(pnp.get_scene_info()) + elif choice == "3": + if (number := _object_number()) is not None: + goal = pnp.get_goal_pose(number) + candidates = pnp.get_grasp_candidates() + _print_grasp_candidates(candidates) + if candidates.candidates: + rank = _grasp_rank(min(10, len(candidates.candidates))) + if rank is None: + continue + goal = pnp.select_grasp_candidate(rank) + pre_grasp = pnp.get_pre_grasp_pose() + approach_planned = False + approach_executed = False + descent_planned = False + descent_executed = False + gripper_closed = False + ascent_planned = False + ascent_executed = False + print("Goal:") + _print_pose(goal) + print("Pre-grasp:") + _print_pose(pre_grasp) + elif choice == "4": + if pre_grasp is None: + print("Select a target first.") + else: + approach_planned = manipulation.plan_to_pose( + Pose(pre_grasp.position, pre_grasp.orientation), "arm" + ) + print(approach_planned) + if approach_planned: + _preview(manipulation) + elif choice == "5": + if not approach_planned: + print("Plan the approach first.") + else: + approach_executed = manipulation.execute_and_wait() + print(approach_executed) + elif choice == "6": + if goal is None or not approach_executed: + print("Execute the approach first.") + else: + waypoints = _cartesian_waypoints(manipulation, goal) + descent_planned = False + descent_planned = manipulation.plan_cartesian_targets( + {"arm/manipulator": waypoints}, + RoboPlanCartesianPathConfig(max_linear_speed=0.03), + ) + print(descent_planned) + if descent_planned: + _preview(manipulation) + elif choice == "7": + if not descent_planned: + print("Plan the descent first.") + else: + descent_executed = manipulation.execute_and_wait() + print(descent_executed) + elif choice == "8": + if not descent_executed: + print("Execute the descent first.") + else: + gripper_closed = manipulation.close_gripper("arm").is_success() + print(gripper_closed) + elif choice == "9": + if pre_grasp is None or not gripper_closed: + print("Close the gripper before planning ascent.") + else: + waypoints = _cartesian_waypoints(manipulation, pre_grasp) + ascent_planned = False + ascent_planned = manipulation.plan_cartesian_targets( + {"arm/manipulator": waypoints}, + RoboPlanCartesianPathConfig(max_linear_speed=0.03), + ) + print(ascent_planned) + if ascent_planned: + _preview(manipulation) + elif choice == "10": + if not ascent_planned: + print("Plan the ascent first.") + else: + # Gripper commands are asynchronous on xArm. Reassert close before + # lift and let that command settle before dispatching the trajectory. + gripper_closed = manipulation.close_gripper("arm").is_success() + if not gripper_closed: + print("Failed to keep the gripper closed; ascent was not executed.") + else: + time.sleep(1.5) + ascent_executed = manipulation.execute_and_wait() + print(ascent_executed) + elif choice == "11": + if not ascent_executed: + print("Execute the ascent before opening the gripper.") + else: + print(manipulation.open_gripper("arm")) + elif choice == "12": + _print_pose(manipulation.get_ee_pose("arm")) + elif choice == "13": + print(manipulation.go_home("arm")) + elif choice == "14": + # A fresh RGB-D snapshot is required before fitting the table plane. + pnp.scan_scene() + estimate = pnp.estimate_table_surface() + if estimate is None: + print("No horizontal tabletop estimate. Scan the scene and try again.") + else: + print( + "Table estimate: " + f"z={estimate['tabletop_z']:.3f} m, center=({estimate['center_x']:.3f}, " + f"{estimate['center_y']:.3f}) m, size=({estimate['width']:.3f}, " + f"{estimate['depth']:.3f}) m" + ) + clearance_text = input( + "Table clearance in mm [10 recommended, 0 = no clearance]: " + ).strip() + try: + clearance_mm = 10.0 if not clearance_text else float(clearance_text) + except ValueError: + print("Enter a non-negative clearance in millimeters.") + continue + if clearance_mm < 0.0: + print("Enter a non-negative clearance in millimeters.") + continue + print( + manipulation.set_table_collision( + estimate["center_x"], + estimate["center_y"], + estimate["tabletop_z"], + estimate["width"], + estimate["depth"], + safety_margin=clearance_mm / 1000.0, + ) + ) + elif choice == "15": + if goal is None or pre_grasp is None or not approach_executed: + print("Execute the approach first.") + continue + print("Executing descent, gripper close, and ascent without previews.") + descent_planned = False + if (waypoints := _cartesian_waypoints(manipulation, goal)) is not None: + descent_planned = manipulation.plan_cartesian_targets( + {"arm/manipulator": waypoints}, + RoboPlanCartesianPathConfig(max_linear_speed=0.03), + ) + if not descent_planned: + print("Could not plan the descent; grasp sequence stopped.") + continue + descent_executed = manipulation.execute_and_wait() + if not descent_executed: + print("Descent failed; grasp sequence stopped.") + continue + gripper_closed = manipulation.close_gripper("arm").is_success() + if not gripper_closed: + print("Failed to close the gripper; grasp sequence stopped.") + continue + # Gripper commands are asynchronous; wait before lifting the object. + time.sleep(1.5) + ascent_planned = False + if (waypoints := _cartesian_waypoints(manipulation, pre_grasp)) is not None: + ascent_planned = manipulation.plan_cartesian_targets( + {"arm/manipulator": waypoints}, + RoboPlanCartesianPathConfig(max_linear_speed=0.03), + ) + if not ascent_planned: + print("Could not plan the ascent; grasp sequence stopped with gripper closed.") + continue + ascent_executed = manipulation.execute_and_wait() + print(ascent_executed) + elif choice == "16": + question = input( + "Scene question [What objects are visible on the table?]: " + ).strip() + print(pnp.describe_scene(question or "What objects are visible on the table?")) + else: + print("Choose 1-16 or q.") + except Exception as exc: + print(f"RPC failed: {exc}") + + +if __name__ == "__main__": + main() diff --git a/dimos/manipulation/skill_errors.py b/dimos/manipulation/skill_errors.py index 9a17085ec5..c980149c78 100644 --- a/dimos/manipulation/skill_errors.py +++ b/dimos/manipulation/skill_errors.py @@ -35,6 +35,11 @@ "COLLISION_AT_START", "GRASP_GENERATION_FAILED", "GRASP_ATTEMPTS_EXHAUSTED", + "GRASP_PROVIDER_UNAVAILABLE", + "GRASP_INPUT_INVALID", + "GRASP_FRAME_MISMATCH", + "GRASP_VERIFICATION_FAILED", + "PICK_BUSY", "GRIPPER_FAILED", "WORLD_MONITOR_UNAVAILABLE", ] diff --git a/dimos/manipulation/test_connected_grasp_sequence_integration.py b/dimos/manipulation/test_connected_grasp_sequence_integration.py new file mode 100644 index 0000000000..5877842f77 --- /dev/null +++ b/dimos/manipulation/test_connected_grasp_sequence_integration.py @@ -0,0 +1,150 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Guarded GraspGenX-to-RoboPlan connected-sequence integration coverage.""" + +from __future__ import annotations + +import importlib.util +import os + +import numpy as np +import pytest +import torch + +from dimos.manipulation.demo_graspgenx.fixture import load_demo_clouds +from dimos.manipulation.grasping.grasp_gen_x import GraspGenXModule +from dimos.manipulation.pick_and_place_module import PickAndPlaceModule +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.robot.manipulators.xarm.config import make_xarm7_sim_robot_config +from dimos.robot.manipulators.xarm.grasp_config import make_xarm_graspgenx_config + +pytestmark = pytest.mark.self_hosted_large + + +def _require_real_graspgenx_environment() -> None: + if os.environ.get("DIMOS_RUN_GRASPGENX_INTEGRATION") != "1": + pytest.skip("set DIMOS_RUN_GRASPGENX_INTEGRATION=1 to run model integration") + if importlib.util.find_spec("graspgenx") is None: + pytest.skip("graspgenx optional dependency is not installed") + if not torch.cuda.is_available(): + pytest.skip("GraspGenX integration requires CUDA") + + +@pytest.fixture(scope="module") +def recorded_grasp_proposals() -> tuple[np.ndarray, GraspCandidateArray]: + """Run the real model once against the repository's recorded object cloud.""" + _require_real_graspgenx_environment() + try: + _, object_cloud = load_demo_clouds() + except FileNotFoundError: + pytest.skip("recorded graspgenx_ycb_banana_scene data is unavailable") + + config = make_xarm_graspgenx_config() + module = GraspGenXModule(**config.model_dump(exclude={"rpc_transport", "tf_transport", "g"})) + try: + module.start() + proposals = module.propose_grasps(object_cloud) + finally: + module.stop() + return object_cloud.points_f32(), proposals + + +def test_real_graspgenx_proposals_preserve_recorded_cloud_contract( + recorded_grasp_proposals: tuple[np.ndarray, GraspCandidateArray], +) -> None: + """Smoke-test structural invariants without pinning stochastic poses.""" + points, proposals = recorded_grasp_proposals + + assert points.shape == (3500, 3) + assert proposals.header.frame_id == "world" + assert proposals.candidates + scores = np.asarray([candidate.score for candidate in proposals.candidates]) + assert np.all(np.isfinite(scores)) + assert np.all(scores[:-1] >= scores[1:]) + for candidate in proposals.candidates: + pose = candidate.pose + values = np.asarray( + [ + pose.position.x, + pose.position.y, + pose.position.z, + pose.orientation.x, + pose.orientation.y, + pose.orientation.z, + pose.orientation.w, + ] + ) + assert np.all(np.isfinite(values)) + assert np.linalg.norm(values[3:]) == pytest.approx(1.0, abs=1e-5) + + +def test_real_graspgenx_has_a_connected_sequence_in_open_roboplan_scene( + recorded_grasp_proposals: tuple[np.ndarray, GraspCandidateArray], +) -> None: + """Relocate recorded proposals into xArm workspace and require one full path.""" + points, proposals = recorded_grasp_proposals + robot_config = make_xarm7_sim_robot_config() + module = PickAndPlaceModule( + robots=[robot_config], + planning_timeout=10.0, + visualization={"backend": "none"}, + floor_z=None, + ) + module.coordinator_joint_state = None + module.objects = None + try: + module.start() + home = robot_config.home_joints + assert home is not None + module._on_joint_state( + JointState( + name=list(robot_config.get_coordinator_joint_names()), + position=list(home), + ) + ) + source_center = np.mean(points, axis=0) + workspace_center = np.asarray([0.45, 0.0, 0.25]) + translation = workspace_center - source_center + approach = Vector3(0.0, 0.0, -1.0) + + feasible: GraspCandidate | None = None + for candidate in proposals.candidates[:20]: + pose = candidate.pose + relocated = Pose( + Vector3( + pose.position.x + translation[0], + pose.position.y + translation[1], + pose.position.z + translation[2], + ), + Quaternion(pose.orientation), + ) + pre_grasp = module._compute_pre_grasp_pose(relocated, 0.05, approach) + retreat = module._compute_pre_grasp_pose(relocated, 0.05, approach) + failed_index, _ = module._check_connected_pose_sequence( + (pre_grasp, relocated, retreat), + "arm", + ) + if failed_index is None: + feasible = candidate + break + + assert feasible is not None + finally: + module.stop() diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py index 8c886aeafa..efae275b0f 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -17,6 +17,7 @@ from __future__ import annotations from pathlib import Path +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -185,6 +186,22 @@ def _generated_plan_trajectory(joint_names: list[str], *points: list[float]) -> ) +def _connected_sequence_module( + robot_config: RobotModelConfig, +) -> tuple[ManipulationModule, list[str], JointState]: + module = _make_module() + module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} + module._world_monitor = MagicMock() + module._world_monitor.world = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([robot_config]) + names = ["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"] + live = JointState(name=names, position=[0.0, 0.0, 0.0]) + module._world_monitor.current_global_joint_state.return_value = live + module._kinematics = MagicMock() + module._planner = MagicMock() + return module, names, live + + def _make_trajectory(*points: tuple[float, list[float]]) -> JointTrajectory: joint_names = [f"j{i}" for i in range(len(points[0][1]))] if points else [] return JointTrajectory( @@ -196,6 +213,63 @@ def _make_trajectory(*points: tuple[float, list[float]]) -> JointTrajectory: ) +def _make_module() -> ManipulationModule: + """Create a module for pure planning helpers without opening transports.""" + module = ManipulationModule() + module.stop() + return module + + +class TestSafetyLift: + @pytest.mark.parametrize( + ("current_z", "expected_z"), + [(0.01, 0.10), (0.05, None), (0.20, None)], + ) + def test_safety_lift_target_is_derived_without_motion( + self, + mocker: MockerFixture, + current_z: float, + expected_z: float | None, + ) -> None: + module = _make_module() + orientation = Quaternion(0.0, 0.0, 0.0, 1.0) + mocker.patch.object( + module, + "get_ee_pose", + return_value=Pose(Vector3(0.2, -0.1, current_z), orientation), + ) + plan = mocker.patch.object(module, "plan_to_pose") + + target = module._safety_lift_pose("test_arm") + + if expected_z is None: + assert target is None + else: + assert target is not None + assert target.position.as_tuple == pytest.approx((0.2, -0.1, expected_z)) + assert ( + target.orientation.x, + target.orientation.y, + target.orientation.z, + target.orientation.w, + ) == pytest.approx((0.0, 0.0, 0.0, 1.0)) + plan.assert_not_called() + + +class TestAgentMotionRecovery: + def test_move_to_pose_explains_how_to_recover_from_fault(self, module_factory) -> None: + module = module_factory() + module._state = ManipulationState.FAULT + module._error_message = "Trajectory execution timed out" + + result = module.move_to_pose(0.2, 0.0, 0.1) + + assert not result.is_success() + assert result.error_code == "INVALID_STATE" + assert "FAULT" in result.message + assert "reset" in result.message + + class TestObstacleUpdates: def test_complete_update_forwards_new_obstacle_value(self, module_factory) -> None: module = module_factory() @@ -761,6 +835,306 @@ def test_plan_to_pose_targets_uses_group_ik_and_selected_path( _, planner_kwargs = module._planner.plan_selected_joint_path.call_args assert planner_kwargs["goal"] is ik_goal + def test_connected_pose_check_chains_each_plan_endpoint_into_the_next_start( + self, robot_config, mocker: MockerFixture + ): + module, names, live = _connected_sequence_module(robot_config) + goals = [ + JointState(name=names, position=[0.1, 0.0, 0.0]), + JointState(name=names, position=[0.2, 0.1, 0.0]), + JointState(name=names, position=[0.1, 0.2, 0.1]), + ] + solve = mocker.patch.object( + module, + "inverse_kinematics", + side_effect=[IKResult(status=IKStatus.SUCCESS, joint_state=goal) for goal in goals], + ) + module._planner.plan_selected_joint_path.side_effect = lambda **kwargs: PlanningResult( + status=PlanningStatus.SUCCESS, + path=[kwargs["start"], kwargs["goal"]], + ) + poses = [ + Pose(position=Vector3(x=0.4 + index * 0.1), orientation=Quaternion()) + for index in range(3) + ] + + failed_index, endpoint = module._check_connected_pose_sequence(poses, "test_arm") + + assert failed_index is None + assert endpoint is not None + assert endpoint.position == goals[-1].position + ik_seeds = [call.kwargs["seed"] for call in solve.call_args_list] + planner_starts = [ + call.kwargs["start"] for call in module._planner.plan_selected_joint_path.call_args_list + ] + assert [seed.position for seed in ik_seeds] == [ + live.position, + goals[0].position, + goals[1].position, + ] + assert [start.position for start in planner_starts] == [ + live.position, + goals[0].position, + goals[1].position, + ] + assert module._state == ManipulationState.IDLE + assert module._last_plan is None + + def test_connected_pose_plan_exposes_paths_without_storing_them( + self, robot_config, mocker: MockerFixture + ) -> None: + module, names, live = _connected_sequence_module(robot_config) + goals = [ + JointState(name=names, position=[0.1, 0.0, 0.0]), + JointState(name=names, position=[0.2, 0.1, 0.0]), + ] + mocker.patch.object( + module, + "inverse_kinematics", + side_effect=[IKResult(status=IKStatus.SUCCESS, joint_state=goal) for goal in goals], + ) + module._planner.plan_selected_joint_path.side_effect = lambda **kwargs: PlanningResult( + status=PlanningStatus.SUCCESS, + path=[kwargs["start"], kwargs["goal"]], + ) + + result = module._plan_connected_pose_sequence( + [ + Pose(position=Vector3(x=0.4), orientation=Quaternion()), + Pose(position=Vector3(x=0.5), orientation=Quaternion()), + ], + "test_arm", + ) + + assert result.failed_index is None + assert result.endpoint is not None + assert result.endpoint.position == goals[-1].position + assert [[state.position for state in path] for path in result.paths] == [ + [live.position, goals[0].position], + [goals[0].position, goals[1].position], + ] + assert module._state == ManipulationState.IDLE + assert module._last_plan is None + + def test_connected_pose_check_accepts_explicit_start(self, robot_config, mocker: MockerFixture): + module, names, _ = _connected_sequence_module(robot_config) + explicit_start = JointState(name=names, position=[0.3, 0.2, 0.1]) + goal = JointState(name=names, position=[0.4, 0.2, 0.1]) + solve = mocker.patch.object( + module, + "inverse_kinematics", + return_value=IKResult(status=IKStatus.SUCCESS, joint_state=goal), + ) + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, + path=[explicit_start, goal], + ) + + failed_index, endpoint = module._check_connected_pose_sequence( + [Pose(position=Vector3(x=0.5), orientation=Quaternion())], + "test_arm", + explicit_start, + ) + + assert failed_index is None + assert endpoint is not None + assert endpoint.position == goal.position + assert solve.call_args.kwargs["seed"].position == explicit_start.position + assert ( + module._planner.plan_selected_joint_path.call_args.kwargs["start"].position + == explicit_start.position + ) + module._world_monitor.current_global_joint_state.assert_not_called() + + def test_connected_pose_check_returns_first_ik_failure_with_diagnostics( + self, + robot_config, + mocker: MockerFixture, + ): + module, names, live = _connected_sequence_module(robot_config) + first_goal = JointState(name=names, position=[0.1, 0.0, 0.0]) + mocker.patch.object( + module, + "inverse_kinematics", + side_effect=[ + IKResult(status=IKStatus.SUCCESS, joint_state=first_goal), + IKResult(status=IKStatus.NO_SOLUTION, message="blocked IK"), + ], + ) + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, + path=[live, first_goal], + ) + poses = [ + Pose(position=Vector3(x=0.4), orientation=Quaternion()), + Pose(position=Vector3(x=0.5), orientation=Quaternion()), + ] + log_info = mocker.patch("dimos.manipulation.manipulation_module.logger.info") + + failed_index, endpoint = module._check_connected_pose_sequence(poses, "test_arm") + + assert (failed_index, endpoint) == (1, None) + log_info.assert_called_with( + "Connected pose planning failed IK at index %d: %s%s", + 1, + "NO_SOLUTION", + ": blocked IK", + ) + assert module._planner.plan_selected_joint_path.call_count == 1 + + def test_connected_pose_check_returns_path_failure_with_diagnostics( + self, + robot_config, + mocker: MockerFixture, + ): + module, names, _ = _connected_sequence_module(robot_config) + goal = JointState(name=names, position=[0.1, 0.0, 0.0]) + mocker.patch.object( + module, + "inverse_kinematics", + return_value=IKResult(status=IKStatus.SUCCESS, joint_state=goal), + ) + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.NO_SOLUTION, + message="blocked path", + ) + log_info = mocker.patch("dimos.manipulation.manipulation_module.logger.info") + + failed_index, endpoint = module._check_connected_pose_sequence( + [Pose(position=Vector3(x=0.4), orientation=Quaternion())], + "test_arm", + ) + + assert (failed_index, endpoint) == (0, None) + log_info.assert_called_with( + "Connected pose planning failed path at index %d: %s%s", + 0, + "NO_SOLUTION", + ": blocked path", + ) + + @pytest.mark.parametrize( + "path", + [ + [], + [ + JointState( + name=["wrong/joint1", "wrong/joint2", "wrong/joint3"], + position=[0.1, 0.0, 0.0], + ) + ], + ], + ids=["empty", "malformed-endpoint"], + ) + def test_connected_pose_check_rejects_invalid_success_path( + self, + robot_config, + mocker: MockerFixture, + path: list[JointState], + ): + module, names, _ = _connected_sequence_module(robot_config) + goal = JointState(name=names, position=[0.1, 0.0, 0.0]) + mocker.patch.object( + module, + "inverse_kinematics", + return_value=IKResult(status=IKStatus.SUCCESS, joint_state=goal), + ) + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, + path=path, + ) + + result = module._check_connected_pose_sequence( + [Pose(position=Vector3(x=0.4), orientation=Quaternion())], + "test_arm", + ) + + assert result == (0, None) + + @pytest.mark.parametrize( + "blocked_index", + [None, 0, 1, 2], + ids=["open-scene", "pre-grasp-blocked", "grasp-blocked", "retreat-blocked"], + ) + def test_connected_pose_check_uses_one_reusable_scene_with_optional_blocker( + self, + robot_config, + mocker: MockerFixture, + blocked_index: int | None, + ) -> None: + module, names, live = _connected_sequence_module(robot_config) + goals = [ + JointState(name=names, position=[0.1 + index * 0.1, 0.0, 0.0]) for index in range(3) + ] + mocker.patch.object( + module, + "inverse_kinematics", + side_effect=[IKResult(status=IKStatus.SUCCESS, joint_state=goal) for goal in goals], + ) + starts = [live, goals[0], goals[1]] + module._planner.plan_selected_joint_path.side_effect = [ + ( + PlanningResult(status=PlanningStatus.NO_SOLUTION, message="scene blocker") + if index == blocked_index + else PlanningResult( + status=PlanningStatus.SUCCESS, + path=[starts[index], goals[index]], + ) + ) + for index in range(3) + ] + poses = [ + Pose(position=Vector3(x=0.4 + index * 0.1), orientation=Quaternion()) + for index in range(3) + ] + + failed_index, endpoint = module._check_connected_pose_sequence(poses, "test_arm") + + assert failed_index == blocked_index + if blocked_index is None: + assert endpoint is not None + assert endpoint.position == goals[-1].position + else: + assert endpoint is None + assert module._planner.plan_selected_joint_path.call_count == blocked_index + 1 + + def test_connected_pose_check_observes_live_scene_on_each_segment( + self, robot_config, mocker: MockerFixture + ) -> None: + module, names, live = _connected_sequence_module(robot_config) + world = SimpleNamespace(revision=0) + module._world_monitor.world = world + goals = [ + JointState(name=names, position=[0.1 + index * 0.1, 0.0, 0.0]) for index in range(3) + ] + mocker.patch.object( + module, + "inverse_kinematics", + side_effect=[IKResult(status=IKStatus.SUCCESS, joint_state=goal) for goal in goals], + ) + starts = [live, goals[0], goals[1]] + observed_revisions: list[int] = [] + + def plan_with_live_scene(**kwargs) -> PlanningResult: + observed_revisions.append(kwargs["world"].revision) + index = len(observed_revisions) - 1 + world.revision += 1 + return PlanningResult( + status=PlanningStatus.SUCCESS, + path=[starts[index], goals[index]], + ) + + module._planner.plan_selected_joint_path.side_effect = plan_with_live_scene + poses = [ + Pose(position=Vector3(x=0.4 + index * 0.1), orientation=Quaternion()) + for index in range(3) + ] + + failed_index, _ = module._check_connected_pose_sequence(poses, "test_arm") + + assert failed_index is None + assert observed_revisions == [0, 1, 2] + def test_failed_plan_materialization_clears_generated_plan(self, robot_config, module_factory): module = module_factory() registry = PlanningGroupRegistry([robot_config]) diff --git a/dimos/manipulation/test_pick_and_place_unit.py b/dimos/manipulation/test_pick_and_place_unit.py index 6ee984f0e2..1dd84ae38c 100644 --- a/dimos/manipulation/test_pick_and_place_unit.py +++ b/dimos/manipulation/test_pick_and_place_unit.py @@ -16,18 +16,42 @@ from __future__ import annotations +from collections import Counter +from contextlib import nullcontext +import json +from types import SimpleNamespace from unittest.mock import patch +import numpy as np import open3d as o3d import pytest +from pytest_mock import MockerFixture +from dimos.agents.skill_result import SkillResult +from dimos.core.coordination.blueprints import BlueprintAtom, autoconnect +from dimos.core.coordination.module_coordinator import _resolve_single_ref from dimos.core.module import ModuleBase -from dimos.manipulation.pick_and_place_module import PickAndPlaceModule +from dimos.manipulation.grasping.grasp_gen_x import GraspGenXModule +from dimos.manipulation.pick_and_place_module import ( + GraspVerificationConfig, + PickAndPlaceModule, + PickAndPlaceModuleConfig, + _FeasibleGrasp, + _GraspVerification, +) +from dimos.manipulation.skill_errors import ManipulationSkillError +from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.std_msgs.Header import Header from dimos.perception.experimental.object import Object as DetObject +from dimos.perception.experimental.object_scene_registration import ObjectSceneRegistrationModule def _make_det_object( @@ -57,7 +81,9 @@ def _make_det_object( def module() -> PickAndPlaceModule: """Create a PickAndPlaceModule with heavy base init (RPC, config) patched out.""" with patch.object(ModuleBase, "__init__", lambda self, config_args: None): - return PickAndPlaceModule() + result = PickAndPlaceModule() + result.config = PickAndPlaceModuleConfig() + return result class TestFindObjectInDetections: @@ -99,6 +125,14 @@ def test_find_missing_returns_none(self, module): result = module._find_object_in_detections("keyboard") assert result is None + def test_find_by_name_requires_unique_match(self, module): + module._detection_snapshot = [ + _make_det_object(name="cup", object_id="first"), + _make_det_object(name="red cup", object_id="second"), + ] + + assert module._find_object_in_detections("cup") is None + def test_empty_snapshot_returns_none(self, module): module._detection_snapshot = [] @@ -158,3 +192,639 @@ def test_place_back_no_pick_pose_errors(self, module): assert not result.is_success() assert result.error_code == "NO_PRIOR_POSE" assert "pick" in result.message.lower() + + +def test_grasp_pipeline_error_agent_encoding_is_structured() -> None: + result = SkillResult[ManipulationSkillError].fail("PICK_BUSY", "pick in progress") + + payload = json.loads(result.agent_encode()[0]["text"]) + + assert payload == { + "success": False, + "message": "pick in progress", + "error_code": "PICK_BUSY", + "duration_ms": 0.0, + } + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"planning_frame": " "}, "planning_frame"), + ({"grasp_approach_vector": (0.0, 0.0, 2.0)}, "unit vector"), + ( + { + "grasp_verification": { + "open_position": 0.85, + "closed_position": 0.0, + "held_threshold": 0.9, + } + }, + "held_threshold", + ), + ], +) +def test_grasp_pipeline_config_rejects_invalid_values( + kwargs: dict[str, object], message: str +) -> None: + with pytest.raises(ValueError, match=message): + PickAndPlaceModuleConfig(**kwargs) + + +def test_pick_module_declares_optional_perception_and_grasp_specs() -> None: + atom = BlueprintAtom.create(PickAndPlaceModule, kwargs={}) + + refs = {ref.name: ref for ref in atom.module_refs} + + assert refs["_object_scene"].optional is True + assert refs["_grasp_generator"].optional is True + + +@pytest.mark.parametrize( + ("ref_name", "provider"), + [ + ("_grasp_generator", GraspGenXModule), + ("_object_scene", ObjectSceneRegistrationModule), + ], +) +def test_optional_provider_resolves_when_absent_present_or_ambiguous( + ref_name: str, provider: type[ModuleBase] +) -> None: + consumer = BlueprintAtom.create(PickAndPlaceModule, kwargs={}) + module_ref = next(ref for ref in consumer.module_refs if ref.name == ref_name) + + absent = autoconnect(PickAndPlaceModule.blueprint()) + assert _resolve_single_ref(consumer, module_ref, module_ref.spec, absent, set()) is None + + present = autoconnect(PickAndPlaceModule.blueprint(), provider.blueprint()) + assert ( + _resolve_single_ref(consumer, module_ref, module_ref.spec, present, set()) == provider.name + ) + + ambiguous = autoconnect( + PickAndPlaceModule.blueprint(), + provider.blueprint(instance_name="provider-a"), + provider.blueprint(instance_name="provider-b"), + ) + with pytest.raises(Exception, match="Multiple modules met that spec"): + _resolve_single_ref(consumer, module_ref, module_ref.spec, ambiguous, set()) + + +def _pointcloud(frame_id: str = "world", timestamp: float | None = None) -> PointCloud2: + return PointCloud2.from_numpy( + np.asarray([[0.4, 0.0, 0.2], [0.41, 0.01, 0.2]], dtype=np.float32), + frame_id=frame_id, + timestamp=timestamp, + ) + + +def _candidate(x: float, score: float) -> GraspCandidate: + return GraspCandidate( + Pose(Vector3(x, 0.0, 0.2), Quaternion(0.0, 0.0, 0.0, 1.0)), + score, + ) + + +class TestProposalSelection: + def test_provider_receives_real_world_frame_cloud( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + now = 100.0 + cloud = _pointcloud(timestamp=now) + detection = _make_det_object() + scene = mocker.Mock() + scene.get_object_pointcloud_by_object_id.return_value = cloud + generator = mocker.Mock() + generator.propose_grasps.return_value = GraspCandidateArray( + Header(now, "world"), [_candidate(0.4, 0.8)] + ) + module._object_scene = scene + module._grasp_generator = generator + mocker.patch("dimos.manipulation.pick_and_place_module.time.time", return_value=now + 0.1) + + candidates = module._provider_candidates( + detection, SimpleNamespace(proposal_source="grasp_provider") + ) + + generator.propose_grasps.assert_called_once_with(cloud) + assert [(candidate.pose.position.x, candidate.score) for candidate in candidates] == [ + (0.4, 0.8) + ] + + @pytest.mark.parametrize("cloud_available", [False, True]) + def test_provider_rejects_missing_or_stale_cloud( + self, + module: PickAndPlaceModule, + mocker: MockerFixture, + cloud_available: bool, + ) -> None: + scene = mocker.Mock() + scene.get_object_pointcloud_by_object_id.return_value = ( + _pointcloud(timestamp=1.0) if cloud_available else None + ) + module._object_scene = scene + module._grasp_generator = mocker.Mock() + mocker.patch("dimos.manipulation.pick_and_place_module.time.time", return_value=100.0) + + with pytest.raises(RuntimeError, match="point cloud"): + module._provider_candidates( + _make_det_object(), SimpleNamespace(proposal_source="grasp_provider") + ) + + @pytest.mark.parametrize( + ("cloud_frame", "proposal_frame"), + [("camera", "world"), ("world", "camera")], + ) + def test_provider_rejects_frame_mismatch( + self, + module: PickAndPlaceModule, + mocker: MockerFixture, + cloud_frame: str, + proposal_frame: str, + ) -> None: + now = 100.0 + scene = mocker.Mock() + scene.get_object_pointcloud_by_object_id.return_value = _pointcloud( + cloud_frame, timestamp=now + ) + generator = mocker.Mock() + generator.propose_grasps.return_value = GraspCandidateArray( + Header(now, proposal_frame), [_candidate(0.4, 0.8)] + ) + module._object_scene = scene + module._grasp_generator = generator + mocker.patch("dimos.manipulation.pick_and_place_module.time.time", return_value=now) + + with pytest.raises(RuntimeError, match="frame"): + module._provider_candidates( + _make_det_object(), SimpleNamespace(proposal_source="grasp_provider") + ) + + def test_provider_preserves_stable_order_for_equal_scores( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + now = 100.0 + scene = mocker.Mock() + scene.get_object_pointcloud_by_object_id.return_value = _pointcloud(timestamp=now) + generator = mocker.Mock() + generator.propose_grasps.return_value = GraspCandidateArray( + Header(now, "world"), + [_candidate(0.1, 0.5), _candidate(0.2, 0.7), _candidate(0.3, 0.7)], + ) + module._object_scene = scene + module._grasp_generator = generator + mocker.patch("dimos.manipulation.pick_and_place_module.time.time", return_value=now) + + candidates = module._provider_candidates( + _make_det_object(), SimpleNamespace(proposal_source="grasp_provider") + ) + + assert [candidate.pose.position.x for candidate in candidates] == [0.2, 0.3, 0.1] + + def test_explicit_heuristic_fallback_identifies_source( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + module.config.heuristic_grasp_fallback = True + transaction = SimpleNamespace(proposal_source="grasp_provider") + pose = Pose(0.4, 0.0, 0.2) + mocker.patch.object(module, "_generate_grasps_for_pick", return_value=[pose]) + + candidates = module._provider_candidates(_make_det_object(), transaction) + + assert transaction.proposal_source == "heuristic" + assert [(candidate.pose, candidate.score) for candidate in candidates] == [(pose, 0.0)] + + def test_provider_is_required_when_fallback_is_disabled( + self, module: PickAndPlaceModule + ) -> None: + with pytest.raises(RuntimeError, match="fallback is disabled"): + module._provider_candidates( + _make_det_object(), SimpleNamespace(proposal_source="grasp_provider") + ) + + def test_selection_skips_higher_scored_infeasible_candidate( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + endpoint = JointState(name=["arm/joint1"], position=[0.1]) + plan_sequence = mocker.patch.object( + module, + "_check_connected_pose_sequence", + side_effect=[(0, None), (None, endpoint)], + ) + plan_motion = mocker.patch.object(module, "plan_to_pose") + command_gripper = mocker.patch.object(module, "_set_gripper_position") + transaction = SimpleNamespace(rejections=Counter()) + + selected = module._select_feasible_grasp( + [_candidate(0.4, 0.9), _candidate(0.5, 0.8)], + "arm", + 0.1, + transaction, + ) + + assert selected.rank == 2 + assert selected.candidate.score == 0.8 + assert plan_sequence.call_count == 2 + assert transaction.rejections == {"pre_grasp_infeasible": 1} + plan_motion.assert_not_called() + command_gripper.assert_not_called() + + @pytest.mark.parametrize( + ("failed_index", "expected_rejection"), + [ + (0, "pre_grasp_infeasible"), + (1, "grasp_infeasible"), + (2, "retreat_infeasible"), + ], + ) + def test_selection_reports_failed_connected_segment( + self, + module: PickAndPlaceModule, + mocker: MockerFixture, + failed_index: int, + expected_rejection: str, + ) -> None: + mocker.patch.object( + module, + "_check_connected_pose_sequence", + return_value=(failed_index, None), + ) + transaction = SimpleNamespace(rejections=Counter()) + + with pytest.raises(RuntimeError, match="No feasible grasp among 1"): + module._select_feasible_grasp([_candidate(0.4, 0.9)], "arm", 0.1, transaction) + + assert transaction.rejections == {expected_rejection: 1} + + def test_selection_rejects_malformed_candidate_and_honors_limit( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + module.config.max_grasp_candidates_to_check = 1 + invalid = _candidate(0.4, 0.9) + invalid.pose.orientation.w = 0.0 + plan_sequence = mocker.patch.object(module, "_check_connected_pose_sequence") + transaction = SimpleNamespace(rejections=Counter()) + + with pytest.raises(RuntimeError, match="No feasible grasp among 1"): + module._select_feasible_grasp([invalid, _candidate(0.5, 0.8)], "arm", 0.1, transaction) + + plan_sequence.assert_not_called() + assert transaction.rejections == {"invalid": 1} + + +class TestPickTransaction: + def _arrange_success( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> tuple[GraspCandidate, SimpleNamespace]: + detection = _make_det_object() + candidate = _candidate(0.4, 0.9) + selected = _FeasibleGrasp(candidate, 1, Pose(0.4, 0.0, 0.3), Pose(0.4, 0.0, 0.3)) + robot_config = SimpleNamespace(pre_grasp_offset=0.1) + mocker.patch.object( + module, "_get_robot", return_value=("arm", "robot-id", robot_config, None) + ) + mocker.patch.object(module, "_require_pick_object", return_value=detection) + mocker.patch.object(module, "_provider_candidates", return_value=[candidate]) + mocker.patch.object(module, "_select_feasible_grasp", return_value=selected) + mocker.patch.object(module, "_safety_lift_pose", return_value=None) + mocker.patch.object(module, "_lift_if_low", return_value=SkillResult.ok()) + mocker.patch.object(module, "plan_to_pose", return_value=True) + mocker.patch.object(module, "_preview_execute_wait", return_value=SkillResult.ok()) + mocker.patch.object(module, "_set_gripper_position", return_value=True) + mocker.patch.object( + module, + "_verify_grasp", + return_value=_GraspVerification(True, 0.1, "verified"), + ) + suppression = SimpleNamespace(cleanup_error=None) + world = mocker.Mock() + world.suppress_object_obstacle.return_value = nullcontext(suppression) + module._world_monitor = world + return candidate, suppression + + def test_success_executes_ordered_pick_and_records_metadata( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + candidate, _ = self._arrange_success(module, mocker) + + result = module.pick("cup", object_id="abc12345") + + assert result.is_success() + assert result.metadata["candidate_rank"] == 1 + assert result.metadata["candidate_score"] == 0.9 + assert module._last_pick_pose is candidate.pose + assert module._set_gripper_position.call_args_list == [ + mocker.call(0.85, "arm"), + mocker.call(0.0, "arm"), + ] + assert module.plan_to_pose.call_args_list == [ + mocker.call(Pose(0.4, 0.0, 0.3), "arm"), + mocker.call(candidate.pose, "arm"), + mocker.call(Pose(0.4, 0.0, 0.3), "arm"), + ] + + def test_no_safety_lift_validates_candidates_from_current_state( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + self._arrange_success(module, mocker) + check_sequence = mocker.patch.object(module, "_check_connected_pose_sequence") + + result = module.pick("cup", object_id="abc12345") + + assert result.is_success() + check_sequence.assert_not_called() + assert module._select_feasible_grasp.call_args.args[4] is None + + def test_safety_lift_endpoint_is_shared_with_candidate_validation( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + self._arrange_success(module, mocker) + lift_pose = Pose(0.2, 0.0, 0.1) + lift_endpoint = JointState(name=["arm/joint1"], position=[0.2]) + module._safety_lift_pose.return_value = lift_pose + check_sequence = mocker.patch.object( + module, + "_check_connected_pose_sequence", + return_value=(None, lift_endpoint), + ) + + result = module.pick("cup", object_id="abc12345") + + assert result.is_success() + check_sequence.assert_called_once_with((lift_pose,), "arm") + assert module._select_feasible_grasp.call_args.args[4] is lift_endpoint + + def test_safety_lift_planning_failure_aborts_prepare_without_candidate_rejection( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + self._arrange_success(module, mocker) + lift_pose = Pose(0.2, 0.0, 0.1) + module._safety_lift_pose.return_value = lift_pose + check_sequence = mocker.patch.object( + module, + "_check_connected_pose_sequence", + return_value=(0, None), + ) + + result = module.pick("cup", object_id="abc12345") + + assert result.error_code == "PLANNING_FAILED" + assert result.metadata["phase"] == "PREPARE" + assert result.metadata["rejections"] == {} + check_sequence.assert_called_once_with((lift_pose,), "arm") + module._select_feasible_grasp.assert_not_called() + module._lift_if_low.assert_not_called() + module.plan_to_pose.assert_not_called() + module._set_gripper_position.assert_not_called() + + def test_retreat_failure_keeps_gripper_closed( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + self._arrange_success(module, mocker) + module.plan_to_pose.side_effect = [True, True, False] + + result = module.pick("cup", object_id="abc12345") + + assert result.error_code == "PLANNING_FAILED" + assert result.metadata["object_may_be_held"] is True + assert module._set_gripper_position.call_args_list == [ + mocker.call(0.85, "arm"), + mocker.call(0.0, "arm"), + ] + + def test_concurrent_pick_is_rejected_without_robot_access( + self, + module: PickAndPlaceModule, + mocker: MockerFixture, + ) -> None: + get_robot = mocker.patch.object(module, "_get_robot") + log = mocker.patch("dimos.agents.annotation.logger.info") + module._pick_guard.acquire() + try: + result = module.pick("cup") + finally: + module._pick_guard.release() + + assert result.error_code == "PICK_BUSY" + get_robot.assert_not_called() + log.assert_called_once() + assert log.call_args.args[:3] == ( + "SKILL %s result=%s duration_ms=%.1f", + "pick", + "PICK_BUSY", + ) + + def test_cleanup_failure_does_not_hide_primary_failure( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + _, suppression = self._arrange_success(module, mocker) + suppression.cleanup_error = "restore failed" + module.plan_to_pose.side_effect = [False] + + result = module.pick("cup", object_id="abc12345") + + assert result.error_code == "PLANNING_FAILED" + assert "cleanup: restore failed" in result.message + + def test_cleanup_failure_turns_success_into_scene_failure( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + _, suppression = self._arrange_success(module, mocker) + suppression.cleanup_error = "restore failed" + + result = module.pick("cup", object_id="abc12345") + + assert result.error_code == "WORLD_MONITOR_UNAVAILABLE" + assert "restore failed" in result.message + + @pytest.mark.parametrize( + ("setup", "expected_code", "expected_phase"), + [ + ("prepare", "EXECUTION_FAILED", "PREPARE"), + ("open", "GRIPPER_FAILED", "PREPARE"), + ("approach_planning", "PLANNING_FAILED", "APPROACH"), + ("approach_execution", "EXECUTION_FAILED", "APPROACH"), + ("grasp_planning", "PLANNING_FAILED", "GRASP"), + ("grasp_execution", "EXECUTION_FAILED", "GRASP"), + ("close", "GRIPPER_FAILED", "CLOSE"), + ("verification", "GRASP_VERIFICATION_FAILED", "VERIFY"), + ("retreat_planning", "PLANNING_FAILED", "RETREAT"), + ("retreat_execution", "EXECUTION_FAILED", "RETREAT"), + ], + ) + def test_phase_failures_stop_the_pipeline( + self, + module: PickAndPlaceModule, + mocker: MockerFixture, + setup: str, + expected_code: str, + expected_phase: str, + ) -> None: + self._arrange_success(module, mocker) + if setup == "prepare": + module._lift_if_low.return_value = SkillResult.fail("EXECUTION_FAILED", "lift failed") + elif setup == "open": + module._set_gripper_position.return_value = False + elif setup == "approach_planning": + module.plan_to_pose.side_effect = [False] + elif setup == "approach_execution": + module._preview_execute_wait.side_effect = [ + SkillResult.fail("EXECUTION_FAILED", "rejected") + ] + elif setup == "grasp_planning": + module.plan_to_pose.side_effect = [True, False] + elif setup == "grasp_execution": + module._preview_execute_wait.side_effect = [ + SkillResult.ok(), + SkillResult.fail("EXECUTION_FAILED", "rejected"), + ] + elif setup == "close": + module._set_gripper_position.side_effect = [True, False] + elif setup == "verification": + module._verify_grasp.return_value = _GraspVerification(False, 0.0, "empty close") + elif setup == "retreat_planning": + module.plan_to_pose.side_effect = [True, True, False] + else: + module._preview_execute_wait.side_effect = [ + SkillResult.ok(), + SkillResult.ok(), + SkillResult.fail("EXECUTION_FAILED", "rejected"), + ] + + result = module.pick("cup", object_id="abc12345") + + assert result.error_code == expected_code + assert result.metadata["phase"] == expected_phase + module._select_feasible_grasp.assert_called_once() + + +def test_full_pick_pipeline_uses_real_messages_and_fake_boundary_providers( + module: PickAndPlaceModule, mocker: MockerFixture +) -> None: + now = 100.0 + detection = _make_det_object() + module._detection_snapshot = [detection] + scene = mocker.Mock() + scene.get_object_pointcloud_by_object_id.return_value = _pointcloud(timestamp=now) + generator = mocker.Mock() + generator.propose_grasps.return_value = GraspCandidateArray( + Header(now, "world"), + [_candidate(0.4, 0.9), _candidate(0.5, 0.8)], + ) + module._object_scene = scene + module._grasp_generator = generator + robot_config = SimpleNamespace(pre_grasp_offset=0.1) + mocker.patch.object(module, "_get_robot", return_value=("arm", "robot-id", robot_config, None)) + plan_sequence = mocker.patch.object( + module, + "_check_connected_pose_sequence", + side_effect=[(0, None), (None, JointState())], + ) + mocker.patch.object(module, "_safety_lift_pose", return_value=None) + mocker.patch.object(module, "_lift_if_low", return_value=SkillResult.ok()) + plan = mocker.patch.object(module, "plan_to_pose", return_value=True) + execute = mocker.patch.object(module, "_preview_execute_wait", return_value=SkillResult.ok()) + gripper = mocker.patch.object(module, "_set_gripper_position", return_value=True) + suppression = SimpleNamespace(cleanup_error=None) + world = mocker.Mock() + world.suppress_object_obstacle.return_value = nullcontext(suppression) + module._world_monitor = world + mocker.patch("dimos.manipulation.pick_and_place_module.time.time", return_value=now) + + result = module.pick("cup", object_id="abc12345") + + assert result.is_success() + assert result.metadata["proposal_source"] == "grasp_provider" + assert result.metadata["candidate_rank"] == 2 + assert result.metadata["candidate_score"] == 0.8 + assert result.metadata["rejections"] == {"pre_grasp_infeasible": 1} + scene.get_object_pointcloud_by_object_id.assert_called_once_with("abc12345") + generator.propose_grasps.assert_called_once_with(scene.get_object_pointcloud_by_object_id()) + world.suppress_object_obstacle.assert_called_once_with("abc12345") + assert world.method_calls == [mocker.call.suppress_object_obstacle("abc12345")] + assert plan_sequence.call_count == 2 + assert plan.call_count == 3 + assert execute.call_count == 3 + assert gripper.call_args_list == [mocker.call(0.85, "arm"), mocker.call(0.0, "arm")] + + +class TestGraspVerification: + def test_empty_close_fails_immediately( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + module.config.grasp_verification = GraspVerificationConfig( + enabled=True, + timeout=1.0, + poll_interval=0.1, + held_threshold=0.02, + ) + mocker.patch.object(module, "get_gripper", return_value=0.0) + mocker.patch( + "dimos.manipulation.pick_and_place_module.time.monotonic", + side_effect=[0.0, 0.1], + ) + + result = module._verify_grasp("arm") + + assert result == _GraspVerification(False, 0.0, "gripper reached the empty-closed region") + + def test_held_position_succeeds_after_timeout( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + module.config.grasp_verification = GraspVerificationConfig( + enabled=True, + timeout=1.0, + poll_interval=0.1, + held_threshold=0.02, + ) + mocker.patch.object(module, "get_gripper", return_value=0.1) + mocker.patch( + "dimos.manipulation.pick_and_place_module.time.monotonic", + side_effect=[0.0, 0.1, 1.1], + ) + sleep = mocker.patch("dimos.manipulation.pick_and_place_module.time.sleep") + + result = module._verify_grasp("arm") + + assert result == _GraspVerification(True, 0.1, "grasp verified by gripper closure feedback") + sleep.assert_called_once_with(0.1) + + def test_no_gripper_motion_is_not_misclassified_as_a_grasp( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + module.config.grasp_verification = GraspVerificationConfig( + enabled=True, + timeout=1.0, + poll_interval=0.1, + held_threshold=0.02, + ) + mocker.patch.object(module, "get_gripper", return_value=0.85) + mocker.patch( + "dimos.manipulation.pick_and_place_module.time.monotonic", + side_effect=[0.0, 0.1, 1.1], + ) + mocker.patch("dimos.manipulation.pick_and_place_module.time.sleep") + + result = module._verify_grasp("arm") + + assert result == _GraspVerification(False, 0.85, "gripper did not leave the open position") + + def test_feedback_timeout_is_reported( + self, module: PickAndPlaceModule, mocker: MockerFixture + ) -> None: + module.config.grasp_verification = GraspVerificationConfig( + enabled=True, + timeout=1.0, + poll_interval=0.1, + held_threshold=0.02, + ) + mocker.patch.object(module, "get_gripper", return_value=None) + mocker.patch( + "dimos.manipulation.pick_and_place_module.time.monotonic", + side_effect=[0.0, 0.1, 1.1], + ) + mocker.patch("dimos.manipulation.pick_and_place_module.time.sleep") + + result = module._verify_grasp("arm") + + assert result == _GraspVerification(False, None, "gripper feedback was unavailable") diff --git a/dimos/manipulation/test_picknplace.py b/dimos/manipulation/test_picknplace.py new file mode 100644 index 0000000000..f1d693800f --- /dev/null +++ b/dimos/manipulation/test_picknplace.py @@ -0,0 +1,456 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser +from dimos.core.module import ModuleBase +from dimos.manipulation.blueprints import _picknplace_xarm6_model, _xarm_graspgenx, picknplace +from dimos.manipulation.picknplace import ( + PickNPlaceConfig, + PickNPlaceModule, + _estimate_table_surface, + _table_midpoint_grasp_z, +) +from dimos.manipulation.planning.spec.models import IKResult, IKStatus +from dimos.manipulation.visualization.layers import MeshElement +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray +from dimos.msgs.std_msgs.Header import Header +from dimos.robot.manipulators.xarm.grasp_config import XARM_TCP_TO_GRASP_FRAME + + +def test_picknplace_scans_and_selects_target() -> None: + with patch.object(ModuleBase, "__init__", lambda self, config_args: None): + module = PickNPlaceModule() + module.config = PickNPlaceConfig() + module._visualization = MagicMock() + scene = MagicMock() + detections = MagicMock() + module._scene = scene + obj = MagicMock( + ts=1.0, + frame_id="link_base", + center=Vector3(0.1, 0.2, 0.04), + confidence=0.9, + ) + obj.name = "cup" + obj.size = Vector3(0.4, 0.1, 0.2) + obj.camera_transform = None + obj.image = None + obj.pose.orientation = Quaternion(0.0, 0.0, 0.0, 1.0) + scene.scan_scene.side_effect = lambda: (module._on_objects([obj]), detections)[1] + + with patch("dimos.manipulation.picknplace.to_detection3d_array") as to_detection3d_array: + result = MagicMock() + to_detection3d_array.return_value = result + assert module.scan_scene() is result + to_detection3d_array.assert_called_once_with([obj], frame_id="link_base", ts=1.0) + + assert module.get_scene_info() == [{"number": 1, "name": "cup", "confidence": 0.9}] + goal = module.get_goal_pose(1) + assert goal is not None + assert goal.position == Vector3(0.1, 0.2, 0.100) + assert goal.orientation == Quaternion.from_euler(Vector3(-3.141592653589793, 0.0, 0.0)) + pre_grasp = module.get_pre_grasp_pose() + assert pre_grasp is not None + assert pre_grasp.position == Vector3(0.1, 0.2, 0.200) + + selected = module.select_object(1) + assert selected.is_success() + assert selected.metadata["goal"] == { + "x": 0.1, + "y": 0.2, + "z": 0.1, + "roll": -math.pi, + "pitch": 0.0, + "yaw": 0.0, + } + assert selected.metadata["pre_grasp"] == { + "x": 0.1, + "y": 0.2, + "z": 0.2, + "roll": -math.pi, + "pitch": 0.0, + "yaw": 0.0, + } + assert module._selected_object is obj + + module.scan_scene("water bottle") + scene.set_prompts.assert_called_once_with(["water bottle"]) + + module.config = PickNPlaceConfig(align_grasp_yaw=True) + yaw_aligned_goal = module.get_goal_pose(1) + assert yaw_aligned_goal is not None + expected = Quaternion.from_euler(Vector3(-math.pi, 0.0, 0.0)) + assert yaw_aligned_goal.orientation.angle_to(expected) == pytest.approx(0.0) + + +def test_scan_objects_uses_independent_simple_queries() -> None: + with patch.object(ModuleBase, "__init__", lambda self, config_args: None): + module = PickNPlaceModule() + module._visualization = MagicMock() + module.scan_scene = MagicMock(detections_length=3) + module.get_scene_info = MagicMock(return_value=[]) + module._publish_scene_objects = MagicMock() + + result = module.scan_objects([" wooden block ", "white box", " "]) + + assert result.is_success() + assert result.metadata["queried_names"] == ["wooden block", "white box"] + module.scan_scene.assert_called_once_with(prompts=["wooden block", "white box"]) + + +def test_pick_selected_verifies_gripper_did_not_fully_close() -> None: + with patch.object(ModuleBase, "__init__", lambda self, config_args: None): + module = PickNPlaceModule() + module.config = PickNPlaceConfig(grasp_feedback_delay=0.0) + module._goal_pose = PoseStamped( + position=Vector3(0.3, 0.1, 0.12), + orientation=Quaternion.from_euler(Vector3(-math.pi, 0.0, 0.0)), + ) + module._pre_grasp_pose = PoseStamped( + position=Vector3(0.3, 0.1, 0.22), + orientation=Quaternion.from_euler(Vector3(-math.pi, 0.0, 0.0)), + ) + module._pick_execution = MagicMock() + module._pick_execution.open_gripper.return_value = MagicMock(is_success=lambda: True) + module._pick_execution.move_to_pose.return_value = MagicMock(is_success=lambda: True) + module._pick_execution.close_gripper.return_value = MagicMock(is_success=lambda: True) + module._pick_execution.get_gripper.return_value = 0.0 + + result = module.pick_selected() + + assert not result.is_success() + assert result.error_code == "GRASP_VERIFICATION_FAILED" + assert "empty-closed" in result.message + assert result.metadata["gripper_position"] == 0.0 + assert result.metadata["rescan_required"] is True + assert result.metadata["recovered_to_pre_grasp"] is True + assert module._pick_execution.open_gripper.call_count == 2 + assert module._pick_execution.move_to_pose.call_count == 3 + + +def test_place_selected_uses_remembered_box_and_held_object() -> None: + with patch.object(ModuleBase, "__init__", lambda self, config_args: None): + module = PickNPlaceModule() + module.config = PickNPlaceConfig() + module._open_box = { + "center_x": 0.4, + "center_y": -0.1, + "tabletop_z": 0.1, + "rim_z": 0.18, + "opening_width": 0.18, + "opening_depth": 0.14, + } + module._held_object_size = Vector3(0.04, 0.03, 0.02) + module._pick_execution = MagicMock() + module._pick_execution.move_to_pose.return_value = MagicMock(is_success=lambda: True) + module._pick_execution.open_gripper.return_value = MagicMock(is_success=lambda: True) + module._pick_execution.get_ee_pose.return_value = Pose( + Vector3(0.3, 0.1, 0.22), Quaternion(0.0, 0.0, 0.0, 1.0) + ) + + result = module.place_selected() + + assert result.is_success() + assert module._pick_execution.move_to_pose.call_args_list[0].args[:3] == pytest.approx( + (0.3, 0.1, 0.31) + ) + assert module._pick_execution.move_to_pose.call_args_list[1].args[:3] == pytest.approx( + (0.4, -0.1, 0.31) + ) + assert module._pick_execution.move_to_pose.call_args_list[2].args[:3] == pytest.approx( + (0.4, -0.1, 0.21) + ) + assert result.metadata["object_bottom_clearance"] == pytest.approx(0.02) + assert module._pick_execution.move_to_pose.call_count == 3 + assert module._held_object_size is None + + +def test_picknplace_home_matches_xarm_lifecycle_home() -> None: + assert _picknplace_xarm6_model.home_joints == [ + 0.0, + math.radians(-40.0), + math.radians(-50.0), + 0.0, + math.radians(90.0), + 0.0, + ] + + +def test_picknplace_graspgenx_uses_xarm_tcp_calibration() -> None: + assert _xarm_graspgenx.grasp_frame_to_tcp[2][3] == pytest.approx(0.172) + assert _xarm_graspgenx.grasp_frame_to_tcp[:2] == ((0.0, -1.0, 0.0, 0.0), (1.0, 0.0, 0.0, 0.0)) + assert np.allclose( + np.asarray(_xarm_graspgenx.grasp_frame_to_tcp) @ np.asarray(XARM_TCP_TO_GRASP_FRAME), + np.eye(4), + ) + + +def test_picknplace_yaw_alignment_defaults_to_disabled() -> None: + assert not PickNPlaceConfig().align_grasp_yaw + + +def test_parallel_jaw_yaw_uses_the_nearest_equivalent_orientation() -> None: + assert PickNPlaceModule._closest_parallel_jaw_yaw(-math.pi + 0.02, 0.0) == pytest.approx(0.02) + + +def test_picknplace_blueprint_accepts_short_backend_and_grasp_options() -> None: + config = BlueprintConfigParser(picknplace) + + options = config.parse( + overrides={"osr": {"det": "moondream", "seg": "edgetam"}, "pnp": {"grasp": "graspgenx"}} + ) + + assert options.module_configs["osr"]["det"] == "moondream" + assert options.module_configs["osr"]["seg"] == "edgetam" + assert options.module_configs["pnp"]["grasp"] == "graspgenx" + + +def test_table_surface_estimate_ignores_objects_above_the_table() -> None: + x, y = np.meshgrid(np.linspace(0.2, 0.8, 20), np.linspace(-0.4, 0.4, 20)) + table = np.column_stack((x.ravel(), y.ravel(), np.full(x.size, 0.35))) + object_points = np.array([[0.5, 0.0, 0.55], [0.51, 0.0, 0.57], [0.5, 0.01, 0.56]]) + + estimate = _estimate_table_surface(np.vstack((table, object_points))) + + assert estimate is not None + assert estimate["tabletop_z"] == pytest.approx(0.35, abs=0.01) + assert estimate["width"] >= 0.8 + assert estimate["depth"] >= 1.0 + + +def test_table_midpoint_grasp_uses_observed_object_height() -> None: + points = np.array( + [ + [0.2, 0.1, 0.117], + [0.2, 0.1, 0.119], + [0.2, 0.1, 0.120], + [0.2, 0.1, 0.121], + [0.2, 0.1, 0.120], + [0.2, 0.1, 0.119], + [0.2, 0.1, 0.120], + [0.2, 0.1, 0.119], + [0.2, 0.1, 0.120], + [0.2, 0.1, 0.120], + ] + ) + + assert _table_midpoint_grasp_z(points, 0.100, 0.120) == pytest.approx(0.110275) + assert _table_midpoint_grasp_z(points[:9], 0.100, 0.120) == pytest.approx(0.120) + + +def test_table_surface_estimate_displays_filled_tabletop() -> None: + with patch.object(ModuleBase, "__init__", lambda self, config_args: None): + module = PickNPlaceModule() + x, y = np.meshgrid(np.linspace(0.2, 0.8, 20), np.linspace(-0.4, 0.4, 20)) + scene_cloud = MagicMock() + scene_cloud.points_f32.return_value = np.column_stack( + (x.ravel(), y.ravel(), np.full(x.size, 0.35)) + ) + module._scene = MagicMock(get_full_scene_pointcloud=MagicMock(return_value=scene_cloud)) + module._visualization = MagicMock() + + assert module.estimate_table_surface() is not None + + layer = module._visualization.set_visualization_layer.call_args.args[0] + assert isinstance(layer.elements[0], MeshElement) + assert layer.elements[0].opacity == pytest.approx(1.0) + np.testing.assert_array_equal(layer.elements[0].triangles, [[0, 1, 2], [0, 2, 3]]) + + +def test_estimate_table_runs_a_fresh_scan_before_fitting() -> None: + with patch.object(ModuleBase, "__init__", lambda self, config_args: None): + module = PickNPlaceModule() + module.scan_scene = MagicMock() + module.estimate_table_surface = MagicMock( + return_value={ + "center_x": 0.5, + "center_y": 0.0, + "tabletop_z": 0.35, + "width": 0.8, + "depth": 1.0, + } + ) + + result = module.estimate_table() + + assert result.is_success() + module.scan_scene.assert_called_once_with() + + +def test_install_open_box_is_display_only() -> None: + with patch.object(ModuleBase, "__init__", lambda self, config_args: None): + module = PickNPlaceModule() + obj = MagicMock( + center=Vector3(0.4, 0.1, 0.14), + size=Vector3(0.20, 0.16, 0.08), + ) + obj.pose.orientation = Quaternion.from_euler(Vector3(0.0, 0.0, 0.0)) + obj.pointcloud.points_f32.return_value = np.asarray([[0.3, 0.1, 0.18]] * 10, dtype=np.float32) + module._latest_objects = (obj,) + module._tabletop_z = 0.10 + module._obstacle_world = MagicMock() + module._obstacle_world.update_obstacle.return_value = False + module._obstacle_world.add_obstacle.side_effect = lambda name, *_: name + module._visualization = MagicMock() + + result = module.install_open_box(1, wall_thickness=0.01) + + assert result.is_success() + assert result.metadata["opening_width"] == pytest.approx(0.18) + assert result.metadata["opening_depth"] == pytest.approx(0.14) + module._obstacle_world.add_obstacle.assert_not_called() + module._obstacle_world.update_obstacle.assert_not_called() + layer = module._visualization.set_visualization_layer.call_args.args[0] + assert layer.id == "picknplace/open-box" + assert layer.elements[0].id == "box-envelope" + + +def test_picknplace_uses_top_graspgenx_candidate() -> None: + with patch.object(ModuleBase, "__init__", lambda self, config_args: None): + module = PickNPlaceModule() + module.config = PickNPlaceConfig(grasp_strategy="graspgenx") + obj = MagicMock( + ts=1.0, + frame_id="link_base", + center=Vector3(0.1, 0.2, 0.3), + pointcloud=MagicMock(), + ) + obj.camera_transform = None + obj.image = None + obj.pose.orientation = Quaternion(0.0, 0.0, 0.0, 1.0) + obj.pointcloud.points_f32.return_value = np.asarray([[0.4, 0.5, 0.6]], dtype=np.float32) + module._latest_objects = (obj,) + candidate = GraspCandidate( + Pose( + Vector3(0.4, 0.5, 0.6), + Quaternion.from_euler(Vector3(0.0, math.pi / 2.0, 0.0)), + ), + score=0.9, + ) + second_candidate = GraspCandidate( + Pose(Vector3(0.2, 0.3, 0.4), Quaternion()), + score=0.8, + ) + module._grasp_generator = MagicMock( + propose_grasps=MagicMock( + return_value=GraspCandidateArray( + Header(2.0, "link_base"), [candidate, second_candidate] + ) + ) + ) + module._grasp_filter = MagicMock( + inverse_kinematics_single=MagicMock(return_value=IKResult(IKStatus.SUCCESS)) + ) + module.graspgenx_candidates = MagicMock() + module._visualization = MagicMock() + + goal = module.get_goal_pose(1) + + assert goal is not None + assert goal.ts == 2.0 + assert goal.frame_id == "link_base" + assert goal.position == candidate.pose.position + assert goal.orientation == candidate.pose.orientation + assert module.get_grasp_candidates().candidates == [candidate, second_candidate] + module.graspgenx_candidates.publish.assert_called_once_with(module.get_grasp_candidates()) + pre_grasp = module.get_pre_grasp_pose() + assert pre_grasp is not None + assert pre_grasp.position.x == pytest.approx(goal.position.x - 0.1) + assert pre_grasp.position.z == pytest.approx(goal.position.z) + layer = module._visualization.set_visualization_layer.call_args.args[0] + assert layer.id == "picknplace/selection" + assert layer.elements[0].points.shape[1] == 3 + assert layer.elements[1].line_width is None + selected_goal = module.select_grasp_candidate(1) + assert selected_goal is not None + assert selected_goal.position == second_candidate.pose.position + assert module.get_grasp_candidates().selected_index == 1 + pre_grasp = module.get_pre_grasp_pose() + assert pre_grasp is not None + assert pre_grasp.position.x == pytest.approx(selected_goal.position.x) + assert pre_grasp.position.z == pytest.approx(selected_goal.position.z - 0.1) + + +def test_picknplace_excludes_collision_or_ik_infeasible_graspgenx_candidates() -> None: + with patch.object(ModuleBase, "__init__", lambda self, config_args: None): + module = PickNPlaceModule() + module.config = PickNPlaceConfig(grasp_strategy="graspgenx") + safe = GraspCandidate(Pose(Vector3(0.4, 0.5, 0.6), Quaternion()), score=0.9) + unsafe = GraspCandidate(Pose(Vector3(0.2, 0.3, 0.4), Quaternion()), score=0.8) + module._grasp_filter = MagicMock( + inverse_kinematics_single=MagicMock( + side_effect=[IKResult(IKStatus.SUCCESS), IKResult(IKStatus.NO_SOLUTION)] + ) + ) + + filtered = module._filter_graspgenx_candidates( + GraspCandidateArray(Header(2.0, "link_base"), [safe, unsafe]) + ) + + assert filtered.candidates == [safe] + module._grasp_filter.inverse_kinematics_single.assert_any_call( + safe.pose, "arm", check_collision=True + ) + module._grasp_filter.inverse_kinematics_single.assert_any_call( + unsafe.pose, "arm", check_collision=True + ) + + +def test_picknplace_clears_candidates_when_no_graspgenx_proposal_is_safe() -> None: + with patch.object(ModuleBase, "__init__", lambda self, config_args: None): + module = PickNPlaceModule() + module.config = PickNPlaceConfig(grasp_strategy="graspgenx") + obj = MagicMock( + ts=1.0, + frame_id="link_base", + center=Vector3(0.1, 0.2, 0.3), + pointcloud=MagicMock(), + ) + obj.camera_transform = None + obj.image = None + obj.pose.orientation = Quaternion() + module._latest_objects = (obj,) + unsafe = GraspCandidate(Pose(Vector3(0.2, 0.3, 0.4), Quaternion()), score=0.8) + module._grasp_generator = MagicMock( + propose_grasps=MagicMock( + return_value=GraspCandidateArray(Header(2.0, "link_base"), [unsafe]) + ) + ) + module._grasp_filter = MagicMock( + inverse_kinematics_single=MagicMock(return_value=IKResult(IKStatus.NO_SOLUTION)) + ) + module.graspgenx_candidates = MagicMock() + + assert module.get_goal_pose(1) is None + module.graspgenx_candidates.publish.assert_called_once() + assert module.graspgenx_candidates.publish.call_args.args[0].candidates == [] + + +def test_picknplace_returns_empty_candidates_for_obb_grasps() -> None: + with patch.object(ModuleBase, "__init__", lambda self, config_args: None): + module = PickNPlaceModule() + module.config = PickNPlaceConfig() + + assert module.get_grasp_candidates().candidates == [] diff --git a/dimos/manipulation/test_pnpconsole.py b/dimos/manipulation/test_pnpconsole.py new file mode 100644 index 0000000000..b1d302b73e --- /dev/null +++ b/dimos/manipulation/test_pnpconsole.py @@ -0,0 +1,206 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock + +from dimos.manipulation import pnpconsole +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray + + +def test_client_scans_scene_and_quits(monkeypatch) -> None: # type: ignore[no-untyped-def] + pnp = MagicMock() + pnp.scan_scene.return_value = MagicMock(detections_length=3) + app = MagicMock(pnp=pnp) + monkeypatch.setattr(pnpconsole.Dimos, "connect", lambda: app) + choices = iter(["1", "", "q"]) + monkeypatch.setattr("builtins.input", lambda _prompt: next(choices)) + + pnpconsole.main() + + pnp.scan_scene.assert_called_once_with(None) + + +def test_client_scans_scene_with_text_prompt(monkeypatch) -> None: # type: ignore[no-untyped-def] + pnp = MagicMock() + pnp.scan_scene.return_value = MagicMock(detections_length=1) + app = MagicMock(pnp=pnp) + monkeypatch.setattr(pnpconsole.Dimos, "connect", lambda: app) + choices = iter(["1", "water bottle", "q"]) + monkeypatch.setattr("builtins.input", lambda _prompt: next(choices)) + + pnpconsole.main() + + pnp.scan_scene.assert_called_once_with("water bottle") + + +def test_client_describes_current_scene(monkeypatch) -> None: # type: ignore[no-untyped-def] + pnp = MagicMock() + pnp.describe_scene.return_value = "A blue block is on the table." + app = MagicMock(pnp=pnp) + monkeypatch.setattr(pnpconsole.Dimos, "connect", lambda: app) + choices = iter(["16", "", "q"]) + monkeypatch.setattr("builtins.input", lambda _prompt: next(choices)) + + pnpconsole.main() + + pnp.describe_scene.assert_called_once_with("What objects are visible on the table?") + + +def test_client_does_not_execute_without_a_plan(monkeypatch) -> None: # type: ignore[no-untyped-def] + app = MagicMock() + manipulation = app.ManipulationModule + monkeypatch.setattr(pnpconsole.Dimos, "connect", lambda: app) + choices = iter(["5", "q"]) + monkeypatch.setattr("builtins.input", lambda _prompt: next(choices)) + + pnpconsole.main() + + manipulation.execute_and_wait.assert_not_called() + + +def test_client_goes_home(monkeypatch) -> None: # type: ignore[no-untyped-def] + app = MagicMock() + manipulation = app.ManipulationModule + monkeypatch.setattr(pnpconsole.Dimos, "connect", lambda: app) + choices = iter(["13", "q"]) + monkeypatch.setattr("builtins.input", lambda _prompt: next(choices)) + + pnpconsole.main() + + manipulation.go_home.assert_called_once_with("arm") + + +def test_client_does_not_execute_descent_without_a_plan(monkeypatch) -> None: # type: ignore[no-untyped-def] + app = MagicMock() + manipulation = app.ManipulationModule + monkeypatch.setattr(pnpconsole.Dimos, "connect", lambda: app) + choices = iter(["7", "q"]) + monkeypatch.setattr("builtins.input", lambda _prompt: next(choices)) + + pnpconsole.main() + + manipulation.execute_and_wait.assert_not_called() + + +def test_client_previews_descent_before_explicit_execution(monkeypatch) -> None: # type: ignore[no-untyped-def] + goal = PoseStamped(position=Vector3(0.1, 0.2, 0.3)) + pre_grasp = PoseStamped(position=Vector3(0.1, 0.2, 0.2)) + pnp = MagicMock() + pnp.get_goal_pose.return_value = goal + pnp.get_grasp_candidates.return_value = GraspCandidateArray() + pnp.get_pre_grasp_pose.return_value = pre_grasp + app = MagicMock(pnp=pnp) + manipulation = app.ManipulationModule + manipulation.plan_to_pose.return_value = True + manipulation.execute_and_wait.return_value = True + manipulation.get_ee_pose.return_value = pre_grasp + manipulation.plan_cartesian_targets.return_value = True + monkeypatch.setattr(pnpconsole.Dimos, "connect", lambda: app) + choices = iter(["3", "1", "4", "5", "6", "q"]) + monkeypatch.setattr("builtins.input", lambda _prompt: next(choices)) + + pnpconsole.main() + + assert manipulation.execute_and_wait.call_count == 1 + manipulation.preview_plan.assert_called_with(duration=2.0) + + +def test_client_runs_grasp_and_lift_without_preview(monkeypatch) -> None: # type: ignore[no-untyped-def] + goal = PoseStamped(position=Vector3(0.1, 0.2, 0.3)) + pre_grasp = PoseStamped(position=Vector3(0.1, 0.2, 0.2)) + pnp = MagicMock() + pnp.get_goal_pose.return_value = goal + pnp.get_grasp_candidates.return_value = GraspCandidateArray() + pnp.get_pre_grasp_pose.return_value = pre_grasp + app = MagicMock(pnp=pnp) + manipulation = app.ManipulationModule + manipulation.plan_to_pose.return_value = True + manipulation.plan_cartesian_targets.return_value = True + manipulation.execute_and_wait.return_value = True + manipulation.get_ee_pose.return_value = pre_grasp + manipulation.close_gripper.return_value.is_success.return_value = True + monkeypatch.setattr(pnpconsole.Dimos, "connect", lambda: app) + monkeypatch.setattr(pnpconsole.time, "sleep", lambda _: None) + choices = iter(["3", "1", "4", "5", "15", "q"]) + monkeypatch.setattr("builtins.input", lambda _prompt: next(choices)) + + pnpconsole.main() + + assert manipulation.execute_and_wait.call_count == 3 + assert manipulation.plan_cartesian_targets.call_count == 2 + manipulation.close_gripper.assert_called_once_with("arm") + manipulation.preview_plan.assert_called_once_with(duration=2.0) + + +def test_client_installs_table_collision_with_recommended_clearance(monkeypatch) -> None: # type: ignore[no-untyped-def] + pnp = MagicMock() + pnp.estimate_table_surface.return_value = { + "center_x": 0.5, + "center_y": 0.0, + "tabletop_z": 0.35, + "width": 0.8, + "depth": 1.0, + } + app = MagicMock(pnp=pnp) + manipulation = app.ManipulationModule + monkeypatch.setattr(pnpconsole.Dimos, "connect", lambda: app) + choices = iter(["14", "", "q"]) + monkeypatch.setattr("builtins.input", lambda _prompt: next(choices)) + + pnpconsole.main() + + pnp.scan_scene.assert_called_once_with() + manipulation.set_table_collision.assert_called_once_with( + 0.5, 0.0, 0.35, 0.8, 1.0, safety_margin=0.01 + ) + + +def test_client_accepts_zero_table_clearance(monkeypatch) -> None: # type: ignore[no-untyped-def] + pnp = MagicMock() + pnp.estimate_table_surface.return_value = { + "center_x": 0.5, + "center_y": 0.0, + "tabletop_z": 0.35, + "width": 0.8, + "depth": 1.0, + } + app = MagicMock(pnp=pnp) + manipulation = app.ManipulationModule + monkeypatch.setattr(pnpconsole.Dimos, "connect", lambda: app) + choices = iter(["14", "0", "q"]) + monkeypatch.setattr("builtins.input", lambda _prompt: next(choices)) + + pnpconsole.main() + + manipulation.set_table_collision.assert_called_once_with( + 0.5, 0.0, 0.35, 0.8, 1.0, safety_margin=0.0 + ) + + +def test_preview_plays_once_slowly() -> None: + manipulation = MagicMock() + + pnpconsole._preview(manipulation) + + manipulation.preview_plan.assert_called_once_with(duration=2.0) + + +def test_grasp_rank_accepts_default_and_valid_selection(monkeypatch) -> None: # type: ignore[no-untyped-def] + choices = iter(["", "7"]) + monkeypatch.setattr("builtins.input", lambda _prompt: next(choices)) + + assert pnpconsole._grasp_rank(10) == 0 + assert pnpconsole._grasp_rank(10) == 7 diff --git a/dimos/manipulation/test_roboplan_world.py b/dimos/manipulation/test_roboplan_world.py index 57e437041f..6c5f80e624 100644 --- a/dimos/manipulation/test_roboplan_world.py +++ b/dimos/manipulation/test_roboplan_world.py @@ -1357,6 +1357,26 @@ def test_native_planner_names_path_from_robot_config_when_start_is_unnamed( assert [state.name for state in result.path] == [["joint1", "joint2"]] * 3 +def test_native_planner_accepts_hypothetical_start_without_mutating_live_state( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + world, robot_id = _make_world(fake_roboplan, robot_config) + + result = world.plan_joint_path( + world, + robot_id, + JointState(name=["joint1", "joint2"], position=[0.2, -0.1]), + JointState(name=["joint1", "joint2"], position=[0.4, 0.3]), + timeout=1.0, + ) + + assert result.status == PlanningStatus.SUCCESS + assert result.path[0].position == pytest.approx([0.2, -0.1]) + assert result.path[-1].position == pytest.approx([0.4, 0.3]) + live_state = world.get_joint_state(world.get_live_context(), robot_id) + assert live_state.position == pytest.approx([0.0, 0.0]) + + def test_native_selected_planner_returns_global_selected_joint_names( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: @@ -1803,7 +1823,7 @@ def test_native_planner_preserves_other_robot_and_auxiliary_joint_state( robot_config: RobotModelConfig, mocker: MockerFixture, ) -> None: - world, _, second_id, second_config = _make_two_robot_world(fake_roboplan, robot_config) + world, first_id, second_id, second_config = _make_two_robot_world(fake_roboplan, robot_config) world.sync_from_joint_state( second_id, JointState(name=["joint1", "joint2"], position=[0.3, 0.1]), @@ -1831,15 +1851,20 @@ def capture_scene_state( result = world.plan_selected_joint_path( world, selection, - JointState(name=list(selection.joint_names), position=[0.0, 0.0]), + JointState(name=list(selection.joint_names), position=[0.15, -0.1]), JointState(name=list(selection.joint_names), position=[0.2, 0.1]), ) assert result.status == PlanningStatus.SUCCESS + assert result.path[0].position == pytest.approx([0.15, -0.1]) + assert observed_positions["arm__joint1"] == pytest.approx(0.0) + assert observed_positions["arm__joint2"] == pytest.approx(0.0) assert observed_positions["right__joint1"] == pytest.approx(0.3) assert observed_positions["right__joint2"] == pytest.approx(0.1) assert observed_positions["arm__joint3"] == pytest.approx(0.0) assert observed_positions["right__joint3"] == pytest.approx(0.0) + live_state = world.get_joint_state(world.get_live_context(), first_id) + assert live_state.position == pytest.approx([0.0, 0.0]) def test_native_planner_waits_for_every_robot_state( diff --git a/dimos/manipulation/test_table_collision.py b/dimos/manipulation/test_table_collision.py new file mode 100644 index 0000000000..e09ab40622 --- /dev/null +++ b/dimos/manipulation/test_table_collision.py @@ -0,0 +1,34 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock + +import pytest + +from dimos.manipulation.manipulation_module import ManipulationModule + + +def test_table_collision_is_a_conservative_slab() -> None: + module = object.__new__(ManipulationModule) + monitor = MagicMock() + monitor.update_obstacle.return_value = False + monitor.add_obstacle.return_value = "calibrated-table" + module._world_monitor = monitor + + assert module.set_table_collision(0.5, 0.0, 0.35, 0.8, 1.0) + + obstacle = monitor.add_obstacle.call_args.args[0] + assert obstacle.name == "calibrated-table" + assert obstacle.dimensions == (0.8, 1.0, 0.2) + assert obstacle.pose.position.z == pytest.approx(0.25) diff --git a/dimos/manipulation/visualization/layers.py b/dimos/manipulation/visualization/layers.py new file mode 100644 index 0000000000..e85b1a19f2 --- /dev/null +++ b/dimos/manipulation/visualization/layers.py @@ -0,0 +1,226 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Backend-neutral, display-only manipulation visualization layers.""" + +from __future__ import annotations + +from dataclasses import dataclass +import math +import re +from typing import TypeAlias + +import numpy as np +from numpy.typing import NDArray + +_ID_SEGMENT = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*$") + + +def _validate_id(value: str, *, hierarchical: bool) -> str: + if not isinstance(value, str) or not value: + raise ValueError("visualization ID must be a nonempty string") + segments = value.split("/") if hierarchical else [value] + if any(not _ID_SEGMENT.fullmatch(segment) for segment in segments): + kind = "layer" if hierarchical else "element" + raise ValueError(f"{kind} ID contains an invalid segment: {value!r}") + return value + + +def _snapshot_positions(value: NDArray[np.generic], *, name: str) -> NDArray[np.float32]: + result = np.array(value, dtype=np.float32, copy=True) + if result.ndim != 2 or result.shape[1:] != (3,): + raise ValueError(f"{name} must have shape (N, 3)") + if not np.all(np.isfinite(result)): + raise ValueError(f"{name} must contain only finite values") + result.setflags(write=False) + return result + + +def _snapshot_colors( + value: NDArray[np.generic] | None, + *, + count: int, + allow_uniform: bool, +) -> NDArray[np.uint8] | None: + if value is None: + return None + source = np.asarray(value) + valid_shapes: set[tuple[int, ...]] = {(count, 3)} + if allow_uniform: + valid_shapes.add((3,)) + if source.shape not in valid_shapes: + expected = "(3,) or (N, 3)" if allow_uniform else "(N, 3)" + raise ValueError(f"colors must have shape {expected}") + if not np.issubdtype(source.dtype, np.number): + raise ValueError("colors must be numeric RGB values") + numeric = np.asarray(source, dtype=np.float64) + if not np.all(np.isfinite(numeric)): + raise ValueError("colors must contain only finite values") + if np.issubdtype(source.dtype, np.floating) and np.all((numeric >= 0.0) & (numeric <= 1.0)): + numeric = np.rint(numeric * 255.0) + if np.any(numeric < 0.0) or np.any(numeric > 255.0): + raise ValueError("colors must be in [0, 1] or [0, 255]") + if not np.all(numeric == np.rint(numeric)): + raise ValueError("colors above 1 must be integer RGB values") + result = np.array(numeric, dtype=np.uint8, copy=True) + result.setflags(write=False) + return result + + +def _validate_size(value: float | None, *, name: str) -> float | None: + if value is None: + return None + result = float(value) + if not math.isfinite(result) or result <= 0.0: + raise ValueError(f"{name} must be finite and positive") + return result + + +def _snapshot_triangles(value: NDArray[np.generic], *, vertex_count: int) -> NDArray[np.int32]: + source = np.asarray(value) + if source.ndim != 2 or source.shape[1:] != (3,): + raise ValueError("triangles must have shape (M, 3)") + if not np.issubdtype(source.dtype, np.number): + raise ValueError("triangles must contain integer indices") + numeric = np.asarray(source, dtype=np.float64) + if not np.all(np.isfinite(numeric)) or not np.all(numeric == np.rint(numeric)): + raise ValueError("triangles must contain finite integer indices") + if np.any(numeric < 0) or (numeric.size and np.any(numeric >= vertex_count)): + raise ValueError("triangles contain an out-of-range vertex index") + result = np.array(numeric, dtype=np.int32, copy=True) + result.setflags(write=False) + return result + + +def _validate_opacity(value: float) -> float: + result = float(value) + if not math.isfinite(result) or not 0.0 < result <= 1.0: + raise ValueError("opacity must be finite and in (0, 1]") + return result + + +@dataclass(frozen=True) +class PointCloudElement: + """A generic colored point cloud with no planning authority.""" + + id: str + points: NDArray[np.generic] + colors: NDArray[np.generic] | None = None + point_size: float | None = None + + def __post_init__(self) -> None: + object.__setattr__(self, "id", _validate_id(self.id, hierarchical=False)) + points = _snapshot_positions(self.points, name="points") + object.__setattr__(self, "points", points) + object.__setattr__( + self, + "colors", + _snapshot_colors(self.colors, count=len(points), allow_uniform=False), + ) + object.__setattr__(self, "point_size", _validate_size(self.point_size, name="point_size")) + + +@dataclass(frozen=True) +class LineSetElement: + """Indexed line geometry with optional uniform or per-line RGB.""" + + id: str + vertices: NDArray[np.generic] + edges: NDArray[np.generic] + colors: NDArray[np.generic] | None = None + line_width: float | None = None + + def __post_init__(self) -> None: + object.__setattr__(self, "id", _validate_id(self.id, hierarchical=False)) + vertices = _snapshot_positions(self.vertices, name="vertices") + object.__setattr__(self, "vertices", vertices) + + source_edges = np.asarray(self.edges) + if source_edges.ndim != 2 or source_edges.shape[1:] != (2,): + raise ValueError("edges must have shape (M, 2)") + if not np.issubdtype(source_edges.dtype, np.number): + raise ValueError("edges must contain integer indices") + numeric_edges = np.asarray(source_edges, dtype=np.float64) + if not np.all(np.isfinite(numeric_edges)) or not np.all( + numeric_edges == np.rint(numeric_edges) + ): + raise ValueError("edges must contain finite integer indices") + if np.any(numeric_edges < 0) or ( + numeric_edges.size and np.any(numeric_edges >= len(vertices)) + ): + raise ValueError("edges contain an out-of-range vertex index") + edges = np.array(numeric_edges, dtype=np.int32, copy=True) + edges.setflags(write=False) + object.__setattr__(self, "edges", edges) + object.__setattr__( + self, + "colors", + _snapshot_colors(self.colors, count=len(edges), allow_uniform=True), + ) + object.__setattr__(self, "line_width", _validate_size(self.line_width, name="line_width")) + + +@dataclass(frozen=True) +class MeshElement: + """Indexed triangle mesh with a uniform RGB color and opacity.""" + + id: str + vertices: NDArray[np.generic] + triangles: NDArray[np.generic] + color: NDArray[np.generic] + opacity: float = 1.0 + + def __post_init__(self) -> None: + object.__setattr__(self, "id", _validate_id(self.id, hierarchical=False)) + vertices = _snapshot_positions(self.vertices, name="vertices") + object.__setattr__(self, "vertices", vertices) + object.__setattr__( + self, + "triangles", + _snapshot_triangles(self.triangles, vertex_count=len(vertices)), + ) + color = _snapshot_colors(self.color, count=1, allow_uniform=True) + if color is None or color.ndim != 1: + raise ValueError("color must have shape (3,)") + object.__setattr__(self, "color", color) + object.__setattr__(self, "opacity", _validate_opacity(self.opacity)) + + +VisualizationElement: TypeAlias = PointCloudElement | LineSetElement | MeshElement + + +@dataclass(frozen=True) +class VisualizationLayer: + """A complete, owner-scoped set of display-only visual elements.""" + + id: str + frame_id: str + elements: tuple[VisualizationElement, ...] + default_visible: bool = True + + def __post_init__(self) -> None: + object.__setattr__(self, "id", _validate_id(self.id, hierarchical=True)) + if not isinstance(self.frame_id, str) or not self.frame_id.strip(): + raise ValueError("frame_id must be a nonempty string") + object.__setattr__(self, "frame_id", self.frame_id.strip()) + elements = tuple(self.elements) + if any( + not isinstance(item, (PointCloudElement, LineSetElement, MeshElement)) + for item in elements + ): + raise TypeError("elements must be point-cloud, line-set, or mesh elements") + ids = [item.id for item in elements] + if len(ids) != len(set(ids)): + raise ValueError("element IDs must be unique within a layer") + object.__setattr__(self, "elements", elements) diff --git a/dimos/manipulation/visualization/pose_overlay.py b/dimos/manipulation/visualization/pose_overlay.py new file mode 100644 index 0000000000..3b0bb97f34 --- /dev/null +++ b/dimos/manipulation/visualization/pose_overlay.py @@ -0,0 +1,68 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Project 3D poses onto camera images.""" + +import cv2 + +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo +from dimos.msgs.sensor_msgs.Image import Image + + +def draw_pose_axes( + image: Image, + pose: PoseStamped, + camera_from_pose_frame: Transform, + camera_info: CameraInfo, + axis_length_m: float = 0.05, +) -> Image | None: + """Draw a pose midpoint and projected RGB coordinate axes onto an image.""" + center_px = _project_point( + _transform_point(camera_from_pose_frame, pose.position), + camera_info, + ) + if center_px is None: + return None + + overlay = image.to_opencv().copy() + axes = ( + (Vector3(1.0, 0.0, 0.0), (0, 0, 255), "X"), + (Vector3(0.0, 1.0, 0.0), (0, 255, 0), "Y"), + (Vector3(0.0, 0.0, 1.0), (255, 0, 0), "Z"), + ) + for axis, color, label in axes: + endpoint = pose.position + pose.orientation.rotate_vector(axis * axis_length_m) + endpoint_px = _project_point( + _transform_point(camera_from_pose_frame, endpoint), camera_info + ) + if endpoint_px is None: + continue + cv2.arrowedLine(overlay, center_px, endpoint_px, color, 2, tipLength=0.2) + cv2.putText(overlay, label, endpoint_px, cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) + cv2.circle(overlay, center_px, 4, (255, 255, 255), -1) + return Image.from_opencv(overlay, frame_id=image.frame_id, ts=image.ts) + + +def _transform_point(transform: Transform, point: Vector3) -> Vector3: + return transform.rotation.rotate_vector(point) + transform.translation + + +def _project_point(point: Vector3, camera_info: CameraInfo) -> tuple[int, int] | None: + if point.z <= 0: + return None + fx, fy, cx, cy = camera_info.K[0], camera_info.K[4], camera_info.K[2], camera_info.K[5] + return (round(fx * point.x / point.z + cx), round(fy * point.y / point.z + cy)) diff --git a/dimos/manipulation/visualization/rerun.py b/dimos/manipulation/visualization/rerun.py new file mode 100644 index 0000000000..61c4559ec2 --- /dev/null +++ b/dimos/manipulation/visualization/rerun.py @@ -0,0 +1,157 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Rerun configuration for the pick-and-place workflow.""" + +from functools import partial +from typing import Any, cast + +import rerun.blueprint as rrb + +from dimos.robot.manipulators.xarm.grasp_config import XARM_TCP_TO_GRASP_FRAME + + +def picknplace_rerun_config() -> dict[str, Any]: + """Return the Rerun layout and message conversions for pick and place.""" + return { + "blueprint": _blueprint, + "topic_to_entity": _topic_to_entity, + "visual_override": { + "world/color_camera": partial( + _camera_info_to_rerun, + image_topic="world/color_camera/color_image", + ), + "world/pointcloud": _pointcloud_to_rerun, + "world/graspgenx_candidates": _graspgenx_candidates_to_rerun, + "world/detections_3d": None, + "world/depth_camera": None, + "world/depth_camera/depth_image": None, + }, + } + + +def _blueprint() -> rrb.Blueprint: + return rrb.Blueprint( + rrb.Horizontal( + rrb.Vertical( + rrb.Spatial2DView(origin="world/annotated_image", name="Object Segmentation"), + rrb.Spatial2DView(origin="world/basic_grasp_overlay", name="Grasp Pose"), + rrb.Spatial2DView(origin="world/color_camera/color_image", name="RGB"), + ), + rrb.Spatial3DView(origin="world", name="Filtered Objects"), + ) + ) + + +def _topic_to_entity(topic: Any) -> str: + topic_name = str(getattr(topic, "name", topic)).split("#", 1)[0] + entities = { + "/color_image": "world/color_camera/color_image", + "/camera_info": "world/color_camera", + "/depth_image": "world/depth_camera/depth_image", + "/depth_camera_info": "world/depth_camera", + "/basic_grasp_overlay": "world/basic_grasp_overlay", + "/graspgenx_candidates": "world/graspgenx_candidates", + "/detections_3d": "world/detections_3d", + "/pointcloud": "world/pointcloud", + } + for suffix, entity in entities.items(): + if topic_name == suffix or topic_name.endswith(suffix): + return entity + return f"world/{topic_name.lstrip('/')}" + + +def _camera_info_to_rerun(msg: Any, image_topic: str) -> list[tuple[str, Any]]: + return cast( + "list[tuple[str, Any]]", + msg.to_rerun(image_topic=image_topic, optical_frame=getattr(msg, "frame_id", None)), + ) + + +def _pointcloud_to_rerun(msg: Any) -> Any: + return msg.to_rerun(voxel_size=0.001, mode="points") + + +def _graspgenx_candidates_to_rerun(msg: Any) -> list[tuple[str, Any]]: + """Render calibrated xArm TCP grasp candidates and their gripper geometry.""" + import rerun as rr + + root = "world/graspgenx_candidates" + data: list[tuple[str, Any]] = [(root, rr.Clear(recursive=True))] + frame_id = msg.header.frame_id + if frame_id: + data.append((root, rr.Transform3D(parent_frame=f"tf#/{frame_id}"))) + for rank, candidate in enumerate(msg.candidates[:10]): + pose = candidate.pose + path = f"{root}/{rank:02d}" + selected = rank == msg.selected_index + gripper_color = [255, 255, 0] if selected else [100, 190, 255] + data.extend( + [ + ( + path, + rr.Transform3D( + translation=pose.position.as_tuple, + rotation=rr.Quaternion(xyzw=pose.orientation.to_tuple()), + ), + ), + ( + f"{path}/tcp_axes", + rr.Arrows3D( + origins=[[0.0, 0.0, 0.0]] * 3, + vectors=[ + [0.04, 0.0, 0.0], + [0.0, 0.04, 0.0], + [0.0, 0.0, 0.04], + ], + colors=[[255, 0, 0], [0, 255, 0], [0, 128, 255]], + radii=[0.0015] * 3, + ), + ), + ( + f"{path}/gripper_base", + rr.Transform3D( + translation=[0.0, 0.0, XARM_TCP_TO_GRASP_FRAME[2][3]], + rotation=rr.Quaternion(xyzw=[0.0, 0.0, -0.70710678, 0.70710678]), + ), + ), + ( + f"{path}/gripper_base/jaws", + # The model sweep geometry is in the gripper-base frame: + # local X closes the jaws and local +Z approaches the object. + rr.LineStrips3D( + strips=[ + [[-0.0425, 0.0, 0.095], [-0.0425, 0.0, 0.162]], + [[0.0425, 0.0, 0.095], [0.0425, 0.0, 0.162]], + [[-0.0425, 0.0, 0.095], [0.0425, 0.0, 0.095]], + ], + colors=[gripper_color] * 3, + radii=[0.0015] * 3, + ), + ), + ] + ) + if selected: + data.append( + ( + f"{path}/selected", + rr.Points3D( + positions=[[0.0, 0.0, 0.0]], + labels=[f"SELECTED #{rank} score={candidate.score:.3f}"], + colors=[[255, 255, 0]], + radii=[0.008], + ), + ) + ) + return data diff --git a/dimos/manipulation/visualization/test_factory.py b/dimos/manipulation/visualization/test_factory.py index 9f952f0deb..a53913ce74 100644 --- a/dimos/manipulation/visualization/test_factory.py +++ b/dimos/manipulation/visualization/test_factory.py @@ -41,6 +41,7 @@ NoManipulationVisualizationConfig, ) from dimos.manipulation.visualization.factory import create_manipulation_visualization +from dimos.manipulation.visualization.layers import VisualizationLayer from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState @@ -83,6 +84,12 @@ def remove_vis_obstacle(self, obstacle_id: str) -> None: def clear_vis_obstacles(self) -> None: return None + def set_layer(self, layer: VisualizationLayer) -> None: + return None + + def clear_layer(self, layer_id: str) -> None: + return None + class FakeWorld: def add_robot(self, config: RobotModelConfig) -> WorldRobotID: @@ -237,6 +244,12 @@ def remove_vis_obstacle(self, obstacle_id: str) -> None: def clear_vis_obstacles(self) -> None: self.visualization_calls.append(("clear_vis_obstacles",)) + def set_layer(self, layer: VisualizationLayer) -> None: + self.visualization_calls.append(("set_layer", layer)) + + def clear_layer(self, layer_id: str) -> None: + self.visualization_calls.append(("clear_layer", layer_id)) + def test_config_defaults_to_no_visualization() -> None: config = ManipulationModuleConfig() @@ -307,6 +320,7 @@ def test_create_visualization_meshcat_accepts_structural_world() -> None: pose=PoseStamped(), dimensions=(1.0, 1.0, 1.0), ) + layer = VisualizationLayer("debug/cloud", "world", ()) visualization.initialize(session) assert visualization.get_visualization_url() == "meshcat://test" visualization.update_state(frame) @@ -316,6 +330,8 @@ def test_create_visualization_meshcat_accepts_structural_world() -> None: visualization.add_vis_obstacle("box", obstacle) visualization.remove_vis_obstacle("box") visualization.clear_vis_obstacles() + visualization.set_layer(layer) + visualization.clear_layer(layer.id) assert fake_world.visualization_calls == [ ("initialize", session), ("get_visualization_url",), @@ -326,6 +342,8 @@ def test_create_visualization_meshcat_accepts_structural_world() -> None: ("add_vis_obstacle", "box", obstacle), ("remove_vis_obstacle", "box"), ("clear_vis_obstacles",), + ("set_layer", layer), + ("clear_layer", layer.id), ] assert fake_world.native_calls == [] diff --git a/dimos/manipulation/visualization/test_layers.py b/dimos/manipulation/visualization/test_layers.py new file mode 100644 index 0000000000..4103b15f1e --- /dev/null +++ b/dimos/manipulation/visualization/test_layers.py @@ -0,0 +1,186 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for backend-neutral visualization layers.""" + +from dataclasses import FrozenInstanceError + +import numpy as np +import pytest + +from dimos.manipulation.visualization.layers import ( + LineSetElement, + MeshElement, + PointCloudElement, + VisualizationLayer, +) + + +def test_point_cloud_snapshots_positions_and_normalizes_colors() -> None: + points = np.asarray([[1.0, 2.0, 3.0]], dtype=np.float64) + colors = np.asarray([[0.0, 0.5, 1.0]], dtype=np.float32) + + element = PointCloudElement("object", points, colors, point_size=0.005) + points[:] = 9.0 + colors[:] = 0.0 + + np.testing.assert_array_equal(element.points, [[1.0, 2.0, 3.0]]) + np.testing.assert_array_equal(element.colors, [[0, 128, 255]]) + assert element.points.dtype == np.float32 + assert element.points.flags.writeable is False + assert element.colors is not None and element.colors.flags.writeable is False + with pytest.raises(ValueError, match="read-only"): + element.points[0, 0] = 4.0 + with pytest.raises(FrozenInstanceError): + element.point_size = 1.0 # type: ignore[misc] + + +def test_line_set_accepts_uniform_and_per_line_colors() -> None: + vertices = np.asarray([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0]]) + edges = np.asarray([[0, 1], [1, 2]]) + + uniform = LineSetElement("uniform", vertices, edges, colors=np.asarray([255, 0, 0])) + per_line = LineSetElement( + "rank-1", + vertices, + edges, + colors=np.asarray([[0, 255, 0], [255, 128, 0]]), + line_width=2.0, + ) + + np.testing.assert_array_equal(uniform.colors, [255, 0, 0]) + np.testing.assert_array_equal(per_line.colors, [[0, 255, 0], [255, 128, 0]]) + assert per_line.edges.dtype == np.int32 + assert per_line.edges.flags.writeable is False + + +def test_mesh_snapshots_triangles_color_and_opacity() -> None: + vertices = np.asarray([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + triangles = np.asarray([[0, 1, 2]]) + + element = MeshElement("surface", vertices, triangles, np.asarray([0.0, 0.5, 1.0]), 0.65) + vertices[:] = 9.0 + triangles[:] = 0 + + np.testing.assert_array_equal(element.vertices[1], [1.0, 0.0, 0.0]) + np.testing.assert_array_equal(element.triangles, [[0, 1, 2]]) + np.testing.assert_array_equal(element.color, [0, 128, 255]) + assert element.opacity == pytest.approx(0.65) + + +@pytest.mark.parametrize( + ("triangles", "color", "opacity", "message"), + [ + (np.asarray([[0, 1]]), np.asarray([0, 0, 0]), 1.0, "shape"), + (np.asarray([[0, 1, 3]]), np.asarray([0, 0, 0]), 1.0, "out-of-range"), + (np.asarray([[0, 1, 2]]), np.asarray([[0, 0, 0]]), 1.0, "color"), + (np.asarray([[0, 1, 2]]), np.asarray([0, 0, 0]), 0.0, "opacity"), + ], +) +def test_mesh_rejects_invalid_geometry( + triangles: np.ndarray, color: np.ndarray, opacity: float, message: str +) -> None: + with pytest.raises(ValueError, match=message): + MeshElement("surface", np.zeros((3, 3)), triangles, color, opacity) + + +@pytest.mark.parametrize("value", ["", "/grasp", "grasp/", "grasp//cloud", "grasp cloud"]) +def test_layer_rejects_invalid_id(value: str) -> None: + with pytest.raises(ValueError, match="layer ID|visualization ID"): + VisualizationLayer(value, "world", ()) + + +@pytest.mark.parametrize("value", ["", "rank/1", "rank 1"]) +def test_element_rejects_invalid_id(value: str) -> None: + with pytest.raises(ValueError, match="element ID|visualization ID"): + PointCloudElement(value, np.empty((0, 3))) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"points": np.zeros((3,))}, "shape"), + ({"points": np.asarray([[np.nan, 0.0, 0.0]])}, "finite"), + ( + { + "points": np.zeros((2, 3)), + "colors": np.zeros((1, 3)), + }, + "colors", + ), + ( + { + "points": np.zeros((1, 3)), + "colors": np.asarray([[256, 0, 0]]), + }, + "colors", + ), + ({"points": np.zeros((1, 3)), "point_size": 0.0}, "positive"), + ], +) +def test_point_cloud_rejects_invalid_geometry(kwargs: dict[str, object], message: str) -> None: + with pytest.raises(ValueError, match=message): + PointCloudElement("cloud", **kwargs) # type: ignore[arg-type] + + +@pytest.mark.parametrize( + ("edges", "message"), + [ + (np.zeros((2, 3)), "shape"), + (np.asarray([[0.5, 1.0]]), "integer"), + (np.asarray([[-1, 0]]), "out-of-range"), + (np.asarray([[0, 2]]), "out-of-range"), + ], +) +def test_line_set_rejects_invalid_edges(edges: np.ndarray, message: str) -> None: + with pytest.raises(ValueError, match=message): + LineSetElement( + "lines", + np.asarray([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]), + edges, + ) + + +def test_line_set_rejects_invalid_appearance() -> None: + with pytest.raises(ValueError, match="colors"): + LineSetElement( + "lines", + np.asarray([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]), + np.asarray([[0, 1]]), + colors=np.asarray([[0, 0, 0], [255, 255, 255]]), + ) + with pytest.raises(ValueError, match="positive"): + LineSetElement( + "lines", + np.asarray([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]), + np.asarray([[0, 1]]), + line_width=float("nan"), + ) + + +def test_layer_rejects_empty_frame_and_duplicate_elements() -> None: + element = PointCloudElement("object", np.empty((0, 3))) + with pytest.raises(ValueError, match="frame_id"): + VisualizationLayer("grasp/object-cloud", " ", (element,)) + with pytest.raises(ValueError, match="unique"): + VisualizationLayer("grasp/object-cloud", "world", (element, element)) + + +def test_layer_owns_element_tuple() -> None: + source = [PointCloudElement("object", np.empty((0, 3)))] + + layer = VisualizationLayer("grasp/object-cloud", "world", source) # type: ignore[arg-type] + source.clear() + + assert [item.id for item in layer.elements] == ["object"] diff --git a/dimos/manipulation/visualization/test_rerun.py b/dimos/manipulation/visualization/test_rerun.py new file mode 100644 index 0000000000..dc37761bdb --- /dev/null +++ b/dimos/manipulation/visualization/test_rerun.py @@ -0,0 +1,48 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dimos.manipulation.visualization.rerun import ( + _graspgenx_candidates_to_rerun, + _topic_to_entity, +) +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate +from dimos.msgs.manipulation_msgs.GraspCandidateArray import GraspCandidateArray +from dimos.msgs.std_msgs.Header import Header + + +def test_qualified_grasp_candidate_topic_uses_candidate_entity() -> None: + assert ( + _topic_to_entity("dimos/PickNPlaceModule/graspgenx_candidates") + == "world/graspgenx_candidates" + ) + + +def test_top_grasp_candidate_has_selected_rerun_marker() -> None: + candidates = GraspCandidateArray( + Header(0.0, "link_base"), + [ + GraspCandidate(Pose(Vector3(0.1, 0.2, 0.3)), 0.9), + GraspCandidate(Pose(Vector3(0.2, 0.3, 0.4)), 0.8), + ], + selected_index=1, + ) + + paths = [path for path, _ in _graspgenx_candidates_to_rerun(candidates)] + + assert "world/graspgenx_candidates/00/selected" not in paths + assert "world/graspgenx_candidates/01/selected" in paths + assert "world/graspgenx_candidates/01/gripper_base" in paths + assert "world/graspgenx_candidates/01/gripper_base/jaws" in paths diff --git a/dimos/manipulation/visualization/viser/layers.py b/dimos/manipulation/visualization/viser/layers.py new file mode 100644 index 0000000000..5152cc2c8a --- /dev/null +++ b/dimos/manipulation/visualization/viser/layers.py @@ -0,0 +1,277 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Latest-wins Viser registry for generic visualization layers.""" + +from __future__ import annotations + +from dataclasses import dataclass +from threading import Condition, Thread +import time +from typing import Any + +from dimos.manipulation.visualization.layers import VisualizationLayer +from dimos.manipulation.visualization.viser.scene import ViserManipulationScene +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +def _display_name(segment: str) -> str: + return segment.replace("-", " ").replace("_", " ").title() + + +@dataclass +class _LayerState: + visible: bool + generation: int = 0 + warning: str | None = None + + +@dataclass(frozen=True) +class _LayerOperation: + sequence: int + layer: VisualizationLayer | None + + +class ViserLayerManager: + """Own Viser layer state, controls, and asynchronous scene reconciliation.""" + + def __init__(self, server: Any, scene: ViserManipulationScene) -> None: + self._server = server + self._scene = scene + self._condition = Condition() + self._states: dict[str, _LayerState] = {} + self._pending: dict[str, _LayerOperation] = {} + self._sequence = 0 + self._processing = False + self._closed = False + self._leaf_handles: dict[str, Any] = {} + self._group_folders: dict[str, Any] = {} + self._group_handles: dict[str, Any] = {} + self._control_sync_depth = 0 + self._root_folder = server.gui.add_folder("Visualization Layers", expand_by_default=True) + self._worker = Thread( + target=self._run, + name="viser-visualization-layers", + daemon=True, + ) + self._worker.start() + + @property + def layer_ids(self) -> tuple[str, ...]: + with self._condition: + return tuple(sorted(self._states)) + + def visibility(self, layer_id: str) -> bool | None: + with self._condition: + state = self._states.get(layer_id) + return None if state is None else state.visible + + def warning(self, layer_id: str) -> str | None: + with self._condition: + state = self._states.get(layer_id) + return None if state is None else state.warning + + def set_layer(self, layer: VisualizationLayer) -> None: + """Queue a complete replacement without waiting for scene rendering.""" + with self._condition: + if self._closed: + return + if layer.id not in self._states: + self._states[layer.id] = _LayerState(visible=layer.default_visible) + self._rebuild_controls() + self._sequence += 1 + self._pending[layer.id] = _LayerOperation(self._sequence, layer) + self._condition.notify() + + def clear_layer(self, layer_id: str) -> None: + """Queue a clear for a known layer while retaining its viewer state.""" + with self._condition: + if self._closed or layer_id not in self._states: + return + self._sequence += 1 + self._pending[layer_id] = _LayerOperation(self._sequence, None) + self._condition.notify() + + def set_visible(self, layer_id: str, visible: bool) -> None: + with self._condition: + state = self._states.get(layer_id) + if self._closed or state is None: + return + state.visible = bool(visible) + handle = self._leaf_handles.get(layer_id) + if handle is not None: + self._set_control_value(handle, state.visible) + self._sync_parent_controls(layer_id) + self._scene.set_visualization_layer_visible(layer_id, bool(visible)) + + def set_group_visible(self, group_id: str, visible: bool) -> None: + descendants = [ + layer_id for layer_id in self.layer_ids if layer_id.startswith(f"{group_id}/") + ] + for layer_id in descendants: + self.set_visible(layer_id, visible) + + def wait_idle(self, timeout: float = 2.0) -> bool: + deadline = time.monotonic() + timeout + with self._condition: + while self._pending or self._processing: + remaining = deadline - time.monotonic() + if remaining <= 0.0: + return False + self._condition.wait(remaining) + return True + + def close(self) -> None: + with self._condition: + if self._closed: + return + self._closed = True + self._pending.clear() + self._condition.notify_all() + self._worker.join(timeout=2.0) + self._scene.clear_visualization_layers() + for handle in ( + *self._leaf_handles.values(), + *self._group_handles.values(), + *reversed(self._group_folders.values()), + self._root_folder, + ): + remove = getattr(handle, "remove", None) + if callable(remove): + remove() + self._leaf_handles.clear() + self._group_handles.clear() + self._group_folders.clear() + + def _run(self) -> None: + while True: + with self._condition: + while not self._pending and not self._closed: + self._condition.wait() + if self._closed: + self._processing = False + self._condition.notify_all() + return + layer_id = next(iter(self._pending)) + operation = self._pending.pop(layer_id) + state = self._states[layer_id] + state.generation += 1 + generation = state.generation + visible = state.visible + self._processing = True + try: + if operation.layer is None: + self._scene.clear_visualization_layer(layer_id) + else: + self._scene.replace_visualization_layer( + operation.layer, + generation=generation, + visible=visible, + ) + warning = None + except Exception as error: + warning = str(error) + logger.warning( + "Visualization layer update failed for '%s': %s", + layer_id, + error, + exc_info=True, + ) + with self._condition: + state.warning = warning + self._processing = False + self._condition.notify_all() + + def _add_layer_controls(self, layer_id: str) -> None: + segments = layer_id.split("/") + parent = self._root_folder + for index, segment in enumerate(segments[:-1], start=1): + group_id = "/".join(segments[:index]) + folder = self._group_folders.get(group_id) + if folder is None: + with parent: + folder = self._server.gui.add_folder( + _display_name(segment), expand_by_default=True + ) + self._group_folders[group_id] = folder + with folder: + group_handle = self._server.gui.add_checkbox("All", initial_value=True) + group_handle.on_update( + lambda event, selected_group=group_id: self._on_group_update( + selected_group, + bool(event.target.value), + ) + ) + self._group_handles[group_id] = group_handle + parent = folder + state = self._states[layer_id] + with parent: + leaf = self._server.gui.add_checkbox( + _display_name(segments[-1]), initial_value=state.visible + ) + leaf.on_update( + lambda event, selected_layer=layer_id: self._on_leaf_update( + selected_layer, + bool(event.target.value), + ) + ) + self._leaf_handles[layer_id] = leaf + self._sync_parent_controls(layer_id) + + def _rebuild_controls(self) -> None: + for handle in ( + *self._leaf_handles.values(), + *self._group_handles.values(), + *reversed(self._group_folders.values()), + ): + remove = getattr(handle, "remove", None) + if callable(remove): + remove() + self._leaf_handles.clear() + self._group_handles.clear() + self._group_folders.clear() + for layer_id in sorted(self._states): + self._add_layer_controls(layer_id) + + def _sync_parent_controls(self, layer_id: str) -> None: + segments = layer_id.split("/") + for index in range(1, len(segments)): + group_id = "/".join(segments[:index]) + handle = self._group_handles.get(group_id) + if handle is None: + continue + descendants = [ + state.visible + for candidate, state in self._states.items() + if candidate.startswith(f"{group_id}/") + ] + self._set_control_value(handle, bool(descendants) and all(descendants)) + + def _set_control_value(self, handle: Any, value: bool) -> None: + """Update Viser state without treating its callback as a user action.""" + self._control_sync_depth += 1 + try: + handle.value = value + finally: + self._control_sync_depth -= 1 + + def _on_leaf_update(self, layer_id: str, visible: bool) -> None: + if self._control_sync_depth == 0: + self.set_visible(layer_id, visible) + + def _on_group_update(self, group_id: str, visible: bool) -> None: + if self._control_sync_depth == 0: + self.set_group_visible(group_id, visible) diff --git a/dimos/manipulation/visualization/viser/scene.py b/dimos/manipulation/visualization/viser/scene.py index 092ba61a92..9ad274ba79 100644 --- a/dimos/manipulation/visualization/viser/scene.py +++ b/dimos/manipulation/visualization/viser/scene.py @@ -29,6 +29,7 @@ import xml.etree.ElementTree as ET import numpy as np +from numpy.typing import NDArray import trimesh from yourdfpy import URDF # type: ignore[import-untyped] @@ -37,6 +38,12 @@ from dimos.manipulation.planning.spec.enums import ObstacleType from dimos.manipulation.planning.spec.models import DEFAULT_OBSTACLE_RGBA, Obstacle from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake +from dimos.manipulation.visualization.layers import ( + LineSetElement, + MeshElement, + PointCloudElement, + VisualizationLayer, +) from dimos.manipulation.visualization.viser.animation import ( GroupPreviewAnimation, PreviewFrame, @@ -102,6 +109,12 @@ OBSTACLE_FALLBACK_COLOR = (55, 190, 210) OBSTACLE_FALLBACK_OPACITY = 0.55 OBSTACLE_PROXY_COLOR = (255, 45, 25) +VISUALIZATION_LAYER_NAMESPACE = "/manipulation/layers" +VISUALIZATION_POINT_CAP = 20_000 +VISUALIZATION_DEFAULT_POINT_SIZE = 0.005 +VISUALIZATION_DEFAULT_POINT_COLOR = (0, 204, 204) +VISUALIZATION_DEFAULT_LINE_COLOR = (255, 255, 255) +VISUALIZATION_DEFAULT_LINE_WIDTH = 1.0 class RobotDisplayMode(StrEnum): @@ -147,6 +160,8 @@ def __init__( self._obstacles_visible = True self._obstacle_gui_handles: list[object] = [] self._obstacle_warning_handle: Any | None = None + self._visualization_layer_handles: dict[str, list[Any]] = {} + self._visualization_layer_visibility: dict[str, bool] = {} self._closed = False self._ensure_obstacle_control() self._ensure_reference_grid() @@ -161,6 +176,130 @@ def set_obstacles_visible(self, visible: bool) -> None: for handle in handles: self._set_handle_visibility(handle, self._obstacles_visible) + @staticmethod + def _visualization_path_segment(value: str) -> str: + return f"id-{value.encode('utf-8').hex()}" + + def replace_visualization_layer( + self, + layer: VisualizationLayer, + *, + generation: int, + visible: bool, + ) -> None: + """Atomically replace a display-only layer with one complete generation.""" + if layer.frame_id != "world": + raise ValueError( + f"Viser visualization layer '{layer.id}' requires frame 'world', " + f"got '{layer.frame_id}'" + ) + base_path = ( + f"{VISUALIZATION_LAYER_NAMESPACE}/" + f"{self._visualization_path_segment(layer.id)}/generation-{generation}" + ) + pending: list[Any] = [] + with self._scene_lock: + if self._closed: + raise RuntimeError("Viser scene is closed") + try: + for element in layer.elements: + path = f"{base_path}/{self._visualization_path_segment(element.id)}" + if isinstance(element, PointCloudElement): + handle = self._render_point_cloud_element(path, element) + elif isinstance(element, LineSetElement): + handle = self._render_line_set_element(path, element) + elif isinstance(element, MeshElement): + handle = self._render_mesh_element(path, element) + else: + raise TypeError(f"unsupported visualization element: {type(element)!r}") + if handle is not None: + self._set_handle_visibility(handle, False) + pending.append(handle) + except Exception: + for handle in pending: + self._remove_scene_handle(handle) + raise + + previous = self._visualization_layer_handles.get(layer.id, []) + for handle in pending: + self._set_handle_visibility(handle, visible) + self._visualization_layer_handles[layer.id] = pending + self._visualization_layer_visibility[layer.id] = visible + for handle in previous: + self._remove_scene_handle(handle) + + def clear_visualization_layer(self, layer_id: str) -> None: + """Remove one layer's handles while retaining its visibility.""" + with self._scene_lock: + for handle in self._visualization_layer_handles.pop(layer_id, []): + self._remove_scene_handle(handle) + + def clear_visualization_layers(self) -> None: + """Remove every generic layer handle.""" + with self._scene_lock: + for layer_id in list(self._visualization_layer_handles): + self.clear_visualization_layer(layer_id) + + def set_visualization_layer_visible(self, layer_id: str, visible: bool) -> None: + """Apply viewer-owned visibility to a layer's current generation.""" + with self._scene_lock: + self._visualization_layer_visibility[layer_id] = bool(visible) + for handle in self._visualization_layer_handles.get(layer_id, []): + self._set_handle_visibility(handle, bool(visible)) + + def _render_point_cloud_element(self, path: str, element: PointCloudElement) -> Any | None: + if len(element.points) == 0: + return None + stride = max(1, math.ceil(len(element.points) / VISUALIZATION_POINT_CAP)) + points = element.points[::stride] + colors: NDArray[np.uint8] | tuple[int, int, int] + if element.colors is None: + colors = VISUALIZATION_DEFAULT_POINT_COLOR + else: + colors = np.asarray(element.colors[::stride], dtype=np.uint8) + return self.server.scene.add_point_cloud( + path, + points=points, + colors=colors, + point_size=element.point_size or VISUALIZATION_DEFAULT_POINT_SIZE, + point_shape="circle", + visible=False, + ) + + def _render_line_set_element(self, path: str, element: LineSetElement) -> Any | None: + if len(element.edges) == 0: + return None + points = element.vertices[element.edges] + colors: NDArray[np.uint8] | tuple[int, int, int] + if element.colors is None: + colors = VISUALIZATION_DEFAULT_LINE_COLOR + elif element.colors.ndim == 1: + colors = np.asarray(element.colors, dtype=np.uint8) + else: + colors = np.asarray( + np.repeat(element.colors[:, np.newaxis, :], 2, axis=1), + dtype=np.uint8, + ) + return self.server.scene.add_line_segments( + path, + points=points, + colors=colors, + line_width=element.line_width or VISUALIZATION_DEFAULT_LINE_WIDTH, + visible=False, + ) + + def _render_mesh_element(self, path: str, element: MeshElement) -> Any | None: + if len(element.triangles) == 0: + return None + return self.server.scene.add_mesh_simple( + path, + vertices=element.vertices, + faces=element.triangles, + color=tuple(int(value) for value in element.color), + opacity=element.opacity, + visible=False, + ) + def add_vis_obstacle(self, obstacle_id: str, obstacle: Obstacle) -> None: """Render one accepted planner obstacle under the local obstacle namespace.""" with self._scene_lock: @@ -665,6 +804,8 @@ def close(self) -> None: self._obstacle_handles.clear() self._obstacles.clear() self._obstacle_render_failures.clear() + self.clear_visualization_layers() + self._visualization_layer_visibility.clear() for key in list(self._handles): self._remove_handle(key) if self._grid_handle is not None: @@ -676,7 +817,8 @@ def close(self) -> None: self._remove_scene_handle(frame) for urdf in self._collision_fallback_urdfs.values(): self._remove_scene_handle(urdf) - for handle in self._obstacle_gui_handles: + # Viser folders own their children, so remove children before folders. + for handle in reversed(self._obstacle_gui_handles): self._remove_scene_handle(handle) self._obstacle_gui_handles.clear() self._urdfs.clear() diff --git a/dimos/manipulation/visualization/viser/test_layers.py b/dimos/manipulation/visualization/viser/test_layers.py new file mode 100644 index 0000000000..102137fe0f --- /dev/null +++ b/dimos/manipulation/visualization/viser/test_layers.py @@ -0,0 +1,441 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hermetic Viser tests for generic visualization layers.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator +from dataclasses import dataclass, field +from threading import Event +from types import SimpleNamespace + +import numpy as np +import pytest +from pytest_mock import MockerFixture + +pytest.importorskip("viser", reason="Viser optional dependency is not installed") + +from dimos.manipulation.visualization.layers import ( + LineSetElement, + MeshElement, + PointCloudElement, + VisualizationLayer, +) +from dimos.manipulation.visualization.viser.layers import ViserLayerManager +from dimos.manipulation.visualization.viser.scene import ViserManipulationScene + + +@dataclass +class Handle: + name: str + visible: bool = True + value: bool = True + callback: Callable[[object], None] | None = None + removed: bool = False + callback_on_assignment: bool = False + _initialized: bool = field(default=False, init=False) + + def __post_init__(self) -> None: + self._initialized = True + + def __setattr__(self, name: str, value: object) -> None: + previous = getattr(self, name, None) + object.__setattr__(self, name, value) + if ( + name == "value" + and getattr(self, "_initialized", False) + and getattr(self, "callback_on_assignment", False) + and previous != value + and self.callback is not None + ): + self.callback(SimpleNamespace(target=self)) + + def on_update(self, callback: Callable[[object], None]) -> None: + self.callback = callback + + def remove(self) -> None: + self.removed = True + + def trigger(self, value: bool) -> None: + self.value = value + if not self.callback_on_assignment: + assert self.callback is not None + self.callback(SimpleNamespace(target=self)) + + +class Folder(Handle): + def __enter__(self) -> Folder: + return self + + def __exit__(self, *_args: object) -> bool: + return False + + +class Gui: + def __init__(self, *, callback_on_assignment: bool = False) -> None: + self.folders: list[Folder] = [] + self.checkboxes: list[Handle] = [] + self.callback_on_assignment = callback_on_assignment + + def add_folder(self, label: str, **_kwargs: object) -> Folder: + handle = Folder(label) + self.folders.append(handle) + return handle + + def add_checkbox(self, label: str, *, initial_value: bool) -> Handle: + handle = Handle( + label, + value=initial_value, + callback_on_assignment=self.callback_on_assignment, + ) + self.checkboxes.append(handle) + return handle + + +class SceneApi: + def __init__(self) -> None: + self.handles: list[Handle] = [] + self.calls: list[tuple[str, str, dict[str, object]]] = [] + self.fail_on_name: str | None = None + + def add_grid(self, name: str, **_kwargs: object) -> Handle: + return self._add("grid", name, {}) + + def add_point_cloud(self, name: str, **kwargs: object) -> Handle: + return self._add("point_cloud", name, kwargs) + + def add_line_segments(self, name: str, **kwargs: object) -> Handle: + return self._add("line_segments", name, kwargs) + + def add_mesh_simple( + self, name: str, vertices: np.ndarray, faces: np.ndarray, **kwargs: object + ) -> Handle: + return self._add("mesh", name, {"vertices": vertices, "faces": faces, **kwargs}) + + def _add(self, kind: str, name: str, kwargs: dict[str, object]) -> Handle: + if self.fail_on_name is not None and self.fail_on_name in name: + raise RuntimeError("injected render failure") + handle = Handle(name, visible=bool(kwargs.get("visible", True))) + for key, value in kwargs.items(): + setattr(handle, key, value) + self.handles.append(handle) + self.calls.append((kind, name, kwargs)) + return handle + + +class Server: + def __init__(self, *, callback_on_assignment: bool = False) -> None: + self.gui = Gui(callback_on_assignment=callback_on_assignment) + self.scene = SceneApi() + + +class Urdf: + pass + + +@pytest.fixture +def scene() -> Iterator[ViserManipulationScene]: + value = ViserManipulationScene(Server(), Urdf) # type: ignore[arg-type] + try: + yield value + finally: + value.close() + + +@pytest.fixture +def manager( + scene: ViserManipulationScene, +) -> Iterator[ViserLayerManager]: + value = ViserLayerManager(scene.server, scene) + try: + yield value + finally: + value.close() + + +def point_layer( + value: float = 0.0, + *, + visible: bool = True, + count: int = 1, + colors: np.ndarray | None = None, +) -> VisualizationLayer: + points = np.full((count, 3), value, dtype=np.float32) + return VisualizationLayer( + "grasp/object-cloud", + "world", + (PointCloudElement("object", points, colors),), + default_visible=visible, + ) + + +def test_scene_renders_cloud_with_paired_cap_and_fallback_color( + scene: ViserManipulationScene, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr("dimos.manipulation.visualization.viser.scene.VISUALIZATION_POINT_CAP", 2) + colors = np.asarray( + [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11], [12, 13, 14]], + dtype=np.uint8, + ) + layer = point_layer(count=5, colors=colors) + + scene.replace_visualization_layer(layer, generation=1, visible=True) + cloud_call = next(call for call in scene.server.scene.calls if call[0] == "point_cloud") + + np.testing.assert_array_equal( + cloud_call[2]["points"], + np.asarray([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float32), + ) + np.testing.assert_array_equal(cloud_call[2]["colors"], colors[::3]) + assert cloud_call[2]["point_size"] == pytest.approx(0.005) + assert scene.server.scene.handles[-1].visible is True + + scene.replace_visualization_layer(point_layer(value=1.0), generation=2, visible=True) + fallback = scene.server.scene.calls[-1][2] + assert fallback["colors"] == (0, 204, 204) + + +def test_scene_renders_indexed_line_set_and_encodes_logical_ids( + scene: ViserManipulationScene, +) -> None: + element = LineSetElement( + "rank-1", + np.asarray([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0]]), + np.asarray([[0, 1], [1, 2]]), + colors=np.asarray([[0, 255, 0], [255, 128, 0]]), + line_width=2.5, + ) + layer = VisualizationLayer("grasp/proposals", "world", (element,)) + + scene.replace_visualization_layer(layer, generation=7, visible=True) + kind, name, kwargs = scene.server.scene.calls[-1] + + assert kind == "line_segments" + assert "grasp/proposals" not in name + assert "rank-1" not in name + assert "generation-7" in name + np.testing.assert_array_equal( + kwargs["points"], + np.asarray( + [ + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + [[1.0, 0.0, 0.0], [1.0, 1.0, 0.0]], + ] + ), + ) + np.testing.assert_array_equal( + kwargs["colors"], + np.asarray( + [ + [[0, 255, 0], [0, 255, 0]], + [[255, 128, 0], [255, 128, 0]], + ] + ), + ) + assert kwargs["line_width"] == pytest.approx(2.5) + + +def test_scene_renders_filled_mesh(scene: ViserManipulationScene) -> None: + element = MeshElement( + "tabletop-fill", + np.asarray([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 0.0]]), + np.asarray([[0, 1, 2], [0, 2, 3]]), + np.asarray([80, 180, 255]), + opacity=0.65, + ) + + scene.replace_visualization_layer( + VisualizationLayer("pick/table", "world", (element,)), generation=1, visible=True + ) + + kind, _name, kwargs = scene.server.scene.calls[-1] + assert kind == "mesh" + np.testing.assert_array_equal(kwargs["faces"], [[0, 1, 2], [0, 2, 3]]) + assert kwargs["color"] == (80, 180, 255) + assert kwargs["opacity"] == pytest.approx(0.65) + + +def test_scene_failed_replacement_retains_previous_generation( + scene: ViserManipulationScene, +) -> None: + scene.replace_visualization_layer(point_layer(), generation=1, visible=True) + previous = scene.server.scene.handles[-1] + first = PointCloudElement("first", np.asarray([[1.0, 0.0, 0.0]])) + failing = PointCloudElement("fail", np.asarray([[2.0, 0.0, 0.0]])) + replacement = VisualizationLayer("grasp/object-cloud", "world", (first, failing)) + scene.server.scene.fail_on_name = "6661696c" # "fail" in hexadecimal + + with pytest.raises(RuntimeError, match="injected"): + scene.replace_visualization_layer(replacement, generation=2, visible=True) + + assert previous.removed is False + partial = next(handle for handle in scene.server.scene.handles if "generation-2" in handle.name) + assert partial.removed is True + + +def test_manager_registers_hierarchy_and_preserves_hidden_state( + manager: ViserLayerManager, + scene: ViserManipulationScene, +) -> None: + manager.set_layer(point_layer(visible=False)) + manager.set_layer( + VisualizationLayer( + "grasp/proposals", + "world", + ( + LineSetElement( + "rank-1", + np.asarray([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]), + np.asarray([[0, 1]]), + ), + ), + ) + ) + assert manager.wait_idle() + + assert manager.layer_ids == ("grasp/object-cloud", "grasp/proposals") + assert manager.visibility("grasp/object-cloud") is False + assert [folder.name for folder in scene.server.gui.folders if not folder.removed].count( + "Grasp" + ) == 1 + object_handle = next( + handle + for handle in scene.server.gui.checkboxes + if handle.name == "Object Cloud" and not handle.removed + ) + object_handle.trigger(True) + assert manager.visibility("grasp/object-cloud") is True + + manager.set_visible("grasp/object-cloud", False) + manager.set_layer(point_layer(value=2.0, visible=True)) + assert manager.wait_idle() + assert manager.visibility("grasp/object-cloud") is False + assert scene.server.scene.handles[-1].visible is False + + manager.clear_layer("grasp/object-cloud") + assert manager.wait_idle() + assert "grasp/object-cloud" in manager.layer_ids + assert manager.visibility("grasp/object-cloud") is False + + +def test_manager_parent_toggle_updates_all_descendants( + manager: ViserLayerManager, + scene: ViserManipulationScene, +) -> None: + manager.set_layer(point_layer()) + manager.set_layer(VisualizationLayer("grasp/proposals", "world", ())) + assert manager.wait_idle() + parent = next( + handle + for handle in scene.server.gui.checkboxes + if handle.name == "All" and not handle.removed + ) + + parent.trigger(False) + + assert manager.visibility("grasp/object-cloud") is False + assert manager.visibility("grasp/proposals") is False + + +def test_manager_leaf_toggle_does_not_cascade_through_reactive_parent() -> None: + scene = ViserManipulationScene(Server(callback_on_assignment=True), Urdf) # type: ignore[arg-type] + manager = ViserLayerManager(scene.server, scene) + try: + manager.set_layer(point_layer()) + manager.set_layer(VisualizationLayer("grasp/proposals", "world", ())) + assert manager.wait_idle() + object_handle = next( + handle + for handle in scene.server.gui.checkboxes + if handle.name == "Object Cloud" and not handle.removed + ) + + object_handle.trigger(False) + + assert manager.visibility("grasp/object-cloud") is False + assert manager.visibility("grasp/proposals") is True + finally: + manager.close() + scene.close() + + +def test_manager_latest_pending_operation_wins( + manager: ViserLayerManager, + scene: ViserManipulationScene, + mocker: MockerFixture, +) -> None: + entered = Event() + release = Event() + original = scene.replace_visualization_layer + + def block_first(layer: VisualizationLayer, *, generation: int, visible: bool) -> None: + if not entered.is_set(): + entered.set() + assert release.wait(2.0) + original(layer, generation=generation, visible=visible) + + mocker.patch.object(scene, "replace_visualization_layer", side_effect=block_first) + manager.set_layer(point_layer(value=1.0)) + assert entered.wait(2.0) + manager.set_layer(point_layer(value=2.0)) + manager.clear_layer("grasp/object-cloud") + release.set() + + assert manager.wait_idle() + assert not any( + not handle.removed and "generation-" in handle.name for handle in scene.server.scene.handles + ) + + +def test_manager_failure_is_contained_and_cross_layer_remains_independent( + manager: ViserLayerManager, + scene: ViserManipulationScene, +) -> None: + scene.server.scene.fail_on_name = "6661696c" + manager.set_layer( + VisualizationLayer( + "debug/failing", + "world", + (PointCloudElement("fail", np.asarray([[0.0, 0.0, 0.0]])),), + ) + ) + manager.set_layer(point_layer()) + + assert manager.wait_idle() + assert "injected render failure" in (manager.warning("debug/failing") or "") + assert manager.warning("grasp/object-cloud") is None + assert any( + not handle.removed and "generation-" in handle.name for handle in scene.server.scene.handles + ) + + +def test_scene_rejects_unsupported_frame_without_replacing_current( + scene: ViserManipulationScene, +) -> None: + scene.replace_visualization_layer(point_layer(), generation=1, visible=True) + previous = scene.server.scene.handles[-1] + + with pytest.raises(ValueError, match="requires frame 'world'"): + scene.replace_visualization_layer( + VisualizationLayer( + "grasp/object-cloud", + "camera", + (PointCloudElement("object", np.asarray([[0.0, 0.0, 0.0]])),), + ), + generation=2, + visible=True, + ) + + assert previous.removed is False diff --git a/dimos/manipulation/visualization/viser/visualizer.py b/dimos/manipulation/visualization/viser/visualizer.py index 1b64601128..69201fe343 100644 --- a/dimos/manipulation/visualization/viser/visualizer.py +++ b/dimos/manipulation/visualization/viser/visualizer.py @@ -18,6 +18,7 @@ from contextlib import suppress from typing import TYPE_CHECKING +from dimos.manipulation.visualization.layers import VisualizationLayer from dimos.manipulation.visualization.viser.animation import ( GroupPreviewAnimation, PreviewFrame, @@ -25,6 +26,7 @@ ) from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig from dimos.manipulation.visualization.viser.gui import ViserPanelGui +from dimos.manipulation.visualization.viser.layers import ViserLayerManager from dimos.manipulation.visualization.viser.runtime import ( VISER_URDF_INSTALL_HINT, ViserRuntime, @@ -73,6 +75,7 @@ def __init__( self._server: ViserServer | None = None self._scene: ViserManipulationScene | None = None self._gui: ViserPanelGui | None = None + self._layer_manager: ViserLayerManager | None = None self._session_scene: PlanningSceneInfo | None = None self._operator: object | None = None self._current_states: dict[str, JointState] = {} @@ -120,6 +123,7 @@ def _ensure_started(self) -> None: self._server = None self._scene = None self._gui = None + self._layer_manager = None self._closed = True raise self._runtime = runtime @@ -129,6 +133,15 @@ def _ensure_started(self) -> None: self._closed = False logger.info(f"Viser manipulation visualization: {self.get_visualization_url()}") + def _ensure_layer_manager(self) -> ViserLayerManager | None: + """Create generic layer resources only when a layer is first published.""" + if self._layer_manager is not None: + return self._layer_manager + if self._server is None or self._scene is None or self._closed: + return None + self._layer_manager = ViserLayerManager(self._server, self._scene) + return self._layer_manager + def initialize(self, session: VisualizationSession) -> None: """Initialize Viser robot visuals from a one-shot visualization session.""" self._operator = session.operator @@ -227,6 +240,37 @@ def clear_vis_obstacles(self) -> None: if self._scene is not None: self._scene.clear_vis_obstacles() + def set_layer(self, layer: VisualizationLayer) -> None: + """Queue one complete display-only layer replacement.""" + if self._closed: + return + try: + self._ensure_started() + manager = self._ensure_layer_manager() + if manager is not None: + manager.set_layer(layer) + except Exception: + logger.warning( + "Visualization layer submission failed for '%s'", + layer.id, + exc_info=True, + ) + + def clear_layer(self, layer_id: str) -> None: + """Queue a display-only layer clear while retaining viewer state.""" + if self._closed: + return + try: + self._ensure_started() + if self._layer_manager is not None: + self._layer_manager.clear_layer(layer_id) + except Exception: + logger.warning( + "Visualization layer clear failed for '%s'", + layer_id, + exc_info=True, + ) + def update_state(self, frame: VisualizationStateFrame) -> None: """Update current robot render state from a pushed state frame.""" if self._closed: @@ -322,6 +366,11 @@ def close(self) -> None: self._closed = True errors: list[BaseException] = [] try: + if self._layer_manager is not None: + try: + self._layer_manager.close() + except Exception as e: + errors.append(e) if self._gui is not None: try: self._gui.close() @@ -342,5 +391,6 @@ def close(self) -> None: self._server = None self._scene = None self._gui = None + self._layer_manager = None if errors: raise errors[0] diff --git a/dimos/manipulation/visualization_spec.py b/dimos/manipulation/visualization_spec.py new file mode 100644 index 0000000000..d9a144856f --- /dev/null +++ b/dimos/manipulation/visualization_spec.py @@ -0,0 +1,24 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""RPC contract for display-only manipulation visualization layers.""" + +from typing import Protocol + +from dimos.manipulation.visualization.layers import VisualizationLayer +from dimos.spec.utils import Spec + + +class ManipulationVisualizationSpec(Spec, Protocol): + def set_visualization_layer(self, layer: VisualizationLayer) -> bool: ... diff --git a/dimos/models/segmentation/edge_tam.py b/dimos/models/segmentation/edge_tam.py index ef90f5b144..898685a72a 100644 --- a/dimos/models/segmentation/edge_tam.py +++ b/dimos/models/segmentation/edge_tam.py @@ -28,6 +28,7 @@ from dimos.msgs.sensor_msgs.Image import Image from dimos.perception.detection.detectors.base import Detector +from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.detection.type.detection2d.seg import Detection2DSeg from dimos.utils.data import get_data @@ -46,59 +47,103 @@ class SAM2InferenceState(TypedDict): cached_features: dict[int, Any] -class EdgeTAMProcessor(Detector): - _predictor: "SAM2VideoPredictor" - _inference_state: SAM2InferenceState | None - _frame_count: int - _is_tracking: bool - _buffer_size: int +def _build_model() -> "SAM2VideoPredictor": + """Build the EdgeTAM SAM2 model from the local config and checkpoint.""" + local_config_path = Path(__file__).parent / "configs" / "edgetam.yaml" - def __init__( - self, - ) -> None: - local_config_path = Path(__file__).parent / "configs" / "edgetam.yaml" + if not local_config_path.exists(): + raise FileNotFoundError(f"EdgeTAM config not found at {local_config_path}") + + if not torch.cuda.is_available(): + raise RuntimeError("EdgeTAM requires a CUDA-capable GPU") + + cfg = OmegaConf.load(local_config_path) + + overrides = { + "model.sam_mask_decoder_extra_args.dynamic_multimask_via_stability": True, + "model.sam_mask_decoder_extra_args.dynamic_multimask_stability_delta": 0.05, + "model.sam_mask_decoder_extra_args.dynamic_multimask_stability_thresh": 0.98, + "model.binarize_mask_from_pts_for_mem_enc": True, + "model.fill_hole_area": 8, + } - if not local_config_path.exists(): - raise FileNotFoundError(f"EdgeTAM config not found at {local_config_path}") + for key, value in overrides.items(): + OmegaConf.update(cfg, key, value) - if not torch.cuda.is_available(): - raise RuntimeError("EdgeTAM requires a CUDA-capable GPU") + if cfg.model._target_ != "sam2.sam2_video_predictor.SAM2VideoPredictor": + logger.warning(f"Config target is {cfg.model._target_}, forcing SAM2VideoPredictor") + cfg.model._target_ = "sam2.sam2_video_predictor.SAM2VideoPredictor" - cfg = OmegaConf.load(local_config_path) + predictor: SAM2VideoPredictor = instantiate(cfg.model, _recursive_=True) - overrides = { - "model.sam_mask_decoder_extra_args.dynamic_multimask_via_stability": True, - "model.sam_mask_decoder_extra_args.dynamic_multimask_stability_delta": 0.05, - "model.sam_mask_decoder_extra_args.dynamic_multimask_stability_thresh": 0.98, - "model.binarize_mask_from_pts_for_mem_enc": True, - "model.fill_hole_area": 8, - } + # Suppress the per-frame "propagate in video" tqdm bar from sam2. + import sam2.sam2_video_predictor as _svp - for key, value in overrides.items(): - OmegaConf.update(cfg, key, value) + _svp.tqdm = lambda iterable, *args, **kwargs: iterable - if cfg.model._target_ != "sam2.sam2_video_predictor.SAM2VideoPredictor": - logger.warning(f"Config target is {cfg.model._target_}, forcing SAM2VideoPredictor") - cfg.model._target_ = "sam2.sam2_video_predictor.SAM2VideoPredictor" + checkpoint = get_data("models_edgetam") / "edgetam.pt" + state_dict = torch.load(checkpoint, map_location="cpu", weights_only=True)["model"] + missing_keys, unexpected_keys = predictor.load_state_dict(state_dict) + if missing_keys: + raise RuntimeError("Missing keys in EdgeTAM checkpoint") + if unexpected_keys: + raise RuntimeError("Unexpected keys in EdgeTAM checkpoint") - self._predictor = instantiate(cfg.model, _recursive_=True) + return predictor.to("cuda").eval() - # Suppress the per-frame "propagate in video" tqdm bar from sam2 - import sam2.sam2_video_predictor as _svp - _svp.tqdm = lambda iterable, *a, **kw: iterable +class EdgeTAMImageSegmenter: + """Refine detector boxes into single-image EdgeTAM masks.""" - ckpt_path = str(get_data("models_edgetam") / "edgetam.pt") + def __init__(self) -> None: + from sam2.sam2_image_predictor import SAM2ImagePredictor + + self._predictor = SAM2ImagePredictor(_build_model()) + + def segment( + self, detections: ImageDetections2D[Detection2DBBox] + ) -> ImageDetections2D[Detection2DSeg]: + """Return masks that preserve each input detection's metadata.""" + import cv2 - sd = torch.load(ckpt_path, map_location="cpu", weights_only=True)["model"] - missing_keys, unexpected_keys = self._predictor.load_state_dict(sd) - if missing_keys: - raise RuntimeError("Missing keys in checkpoint") - if unexpected_keys: - raise RuntimeError("Unexpected keys in checkpoint") + if not detections.detections: + return ImageDetections2D(detections.image, []) - self._predictor = self._predictor.to("cuda") - self._predictor.eval() + image = detections.image + rgb = cv2.cvtColor(image.to_opencv(), cv2.COLOR_BGR2RGB) + boxes = np.asarray( + [detection.bbox for detection in detections.detections], dtype=np.float32 + ) + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + self._predictor.set_image(rgb) + masks, _, _ = self._predictor.predict(box=boxes, multimask_output=False) + + masks = masks.reshape(-1, *masks.shape[-2:]) + segmented = [ + Detection2DSeg.from_sam2_result( + mask, + detection.track_id, + image, + class_id=detection.class_id, + name=detection.name, + confidence=detection.confidence, + ) + for detection, mask in zip(detections.detections, masks, strict=True) + ] + return ImageDetections2D(image, segmented) + + +class EdgeTAMProcessor(Detector): + _predictor: "SAM2VideoPredictor" + _inference_state: SAM2InferenceState | None + _frame_count: int + _is_tracking: bool + _buffer_size: int + + def __init__( + self, + ) -> None: + self._predictor = _build_model() self._inference_state = None self._frame_count = 0 diff --git a/dimos/models/vl/moondream.py b/dimos/models/vl/moondream.py index e3cfe744ce..ef5b85b408 100644 --- a/dimos/models/vl/moondream.py +++ b/dimos/models/vl/moondream.py @@ -38,6 +38,7 @@ class MoondreamConfig(HuggingFaceModelConfig, VlModelConfig): model_name: str = "vikhyatk/moondream2" dtype: torch.dtype = torch.bfloat16 auto_resize: tuple[int, int] | None = MOONDREAM_DEFAULT_AUTO_RESIZE + compile_model: bool = False class MoondreamVlModel(HuggingFaceModel, VlModel): @@ -46,13 +47,14 @@ class MoondreamVlModel(HuggingFaceModel, VlModel): @cached_property def _model(self) -> AutoModelForCausalLM: - """Load model with compile() for optimization.""" + """Load the model, optionally enabling its experimental compile path.""" model = AutoModelForCausalLM.from_pretrained( self.config.model_name, trust_remote_code=self.config.trust_remote_code, torch_dtype=self.config.dtype, ).to(self.config.device) - model.compile() + if self.config.compile_model: + model.compile() return model def _to_pil(self, image: Image | np.ndarray[Any, Any]) -> PILImage.Image: diff --git a/dimos/msgs/manipulation_msgs/GraspCandidate.py b/dimos/msgs/manipulation_msgs/GraspCandidate.py new file mode 100644 index 0000000000..f83db6016f --- /dev/null +++ b/dimos/msgs/manipulation_msgs/GraspCandidate.py @@ -0,0 +1,40 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import math +import pickle + +from dimos.msgs.geometry_msgs.Pose import Pose + + +class GraspCandidate: + """A robot TCP pose and its generator-local ranking score.""" + + msg_name = "manipulation_msgs.GraspCandidate" + + def __init__(self, pose: Pose | None = None, score: float = 0.0) -> None: + self.pose = pose if pose is not None else Pose(0.0, 0.0, 0.0) + self.score = float(score) + if not math.isfinite(self.score): + raise ValueError("GraspCandidate.score must be finite") + + def encode(self) -> bytes: + return pickle.dumps({"pose": self.pose, "score": self.score}) + + @classmethod + def decode(cls, data: bytes) -> GraspCandidate: + value = pickle.loads(data) + return cls(value["pose"], value["score"]) diff --git a/dimos/msgs/manipulation_msgs/GraspCandidateArray.py b/dimos/msgs/manipulation_msgs/GraspCandidateArray.py new file mode 100644 index 0000000000..06011017f0 --- /dev/null +++ b/dimos/msgs/manipulation_msgs/GraspCandidateArray.py @@ -0,0 +1,66 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from collections.abc import Iterator +import pickle + +from dimos.msgs.manipulation_msgs.GraspCandidate import GraspCandidate +from dimos.msgs.std_msgs.Header import Header + + +class GraspCandidateArray: + """Ordered grasp proposals sharing one input-cloud header.""" + + msg_name = "manipulation_msgs.GraspCandidateArray" + + def __init__( + self, + header: Header | None = None, + candidates: list[GraspCandidate] | None = None, + selected_index: int = 0, + ) -> None: + self.header = header if header is not None else Header(0.0) + self.candidates = candidates if candidates is not None else [] + self.selected_index = selected_index + + def __len__(self) -> int: + return len(self.candidates) + + def __iter__(self) -> Iterator[GraspCandidate]: + return iter(self.candidates) + + def encode(self) -> bytes: + """Encode using the repository's pickle transport convention.""" + return pickle.dumps( + { + "header": self.header, + "candidates": self.candidates, + "selected_index": self.selected_index, + } + ) + + @classmethod + def decode(cls, data: bytes) -> GraspCandidateArray: + value = pickle.loads(data) + return cls(value["header"], value["candidates"], value.get("selected_index", 0)) + + # Typed LCM transport lets the Rerun bridge subscribe to proposal updates. + def lcm_encode(self) -> bytes: + return self.encode() + + @classmethod + def lcm_decode(cls, data: bytes, **kwargs: object) -> GraspCandidateArray: + return cls.decode(data) diff --git a/dimos/perception/detection/detectors/moondream.py b/dimos/perception/detection/detectors/moondream.py new file mode 100644 index 0000000000..2b005abed9 --- /dev/null +++ b/dimos/perception/detection/detectors/moondream.py @@ -0,0 +1,63 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Text-prompted Moondream object detector adapter.""" + +from dimos.models.vl.moondream import MoondreamVlModel +from dimos.msgs.sensor_msgs.Image import Image +from dimos.perception.detection.detectors.base import Detector +from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D + + +class Moondream2DDetector(Detector): + """Identify prompted objects with Moondream's native box detector.""" + + def __init__(self, model: MoondreamVlModel | None = None, *, max_objects: int = 5) -> None: + if max_objects < 1: + raise ValueError("max_objects must be positive") + self._model = model or MoondreamVlModel() + self._model.start() + self._max_objects = max_objects + self._text_prompts: tuple[str, ...] = () + + def set_prompts( + self, + text: list[str] | None = None, + bboxes: object | None = None, + ) -> None: + """Set one or more text queries; visual prompts are unsupported.""" + if bboxes is not None: + raise ValueError("Moondream detector supports text prompts only") + if text is None: + raise ValueError("Moondream detector requires at least one text prompt") + prompts = tuple(prompt.strip() for prompt in text if prompt.strip()) + if not prompts: + raise ValueError("Moondream detector requires at least one nonempty text prompt") + self._text_prompts = prompts + + def process_image(self, image: Image) -> ImageDetections2D: + """Run prompted object identification on one image.""" + detections = ImageDetections2D(image) + for prompt in self._text_prompts: + result = self._model.query_detections(image, prompt, max_objects=self._max_objects) + detections.detections.extend(result.detections) + return detections + + def describe_image(self, image: Image, question: str) -> str: + """Answer an open-ended question about an image with the loaded VLM.""" + return str(self._model.query(image, question)) + + def stop(self) -> None: + """Release the Moondream model and GPU memory.""" + self._model.stop() diff --git a/dimos/perception/detection/detectors/test_moondream.py b/dimos/perception/detection/detectors/test_moondream.py new file mode 100644 index 0000000000..8c596b03ce --- /dev/null +++ b/dimos/perception/detection/detectors/test_moondream.py @@ -0,0 +1,78 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock + +import numpy as np +import pytest + +from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.perception.detection.detectors.moondream import Moondream2DDetector +from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox +from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D + + +def test_moondream_detector_queries_each_text_prompt() -> None: + image = Image(data=np.zeros((4, 4, 3), dtype=np.uint8), format=ImageFormat.BGR) + model = MagicMock() + model.query_detections.side_effect = [ + ImageDetections2D( + image, + [ + Detection2DBBox( + (0.0, 0.0, 2.0, 2.0), + track_id=0, + class_id=-1, + confidence=1.0, + name="cup", + ts=image.ts, + image=image, + ) + ], + ), + ImageDetections2D( + image, + [ + Detection2DBBox( + (2.0, 2.0, 3.0, 3.0), + track_id=1, + class_id=-1, + confidence=1.0, + name="bottle", + ts=image.ts, + image=image, + ) + ], + ), + ] + detector = Moondream2DDetector(model=model) + + detector.set_prompts(text=["cup", "bottle"]) + detections = detector.process_image(image) + + assert [detection.name for detection in detections] == ["cup", "bottle"] + assert model.query_detections.call_args_list[0].args[1] == "cup" + assert model.query_detections.call_args_list[1].args[1] == "bottle" + detector.stop() + model.start.assert_called_once_with() + model.stop.assert_called_once_with() + + +def test_moondream_detector_rejects_empty_and_visual_prompts() -> None: + detector = Moondream2DDetector(model=MagicMock()) + + with pytest.raises(ValueError, match="nonempty"): + detector.set_prompts(text=[" "]) + with pytest.raises(ValueError, match="text prompts"): + detector.set_prompts(text=["cup"], bboxes=object()) diff --git a/dimos/perception/detection/type/detection2d/seg.py b/dimos/perception/detection/type/detection2d/seg.py index eb6ec734fa..4a510598bb 100644 --- a/dimos/perception/detection/type/detection2d/seg.py +++ b/dimos/perception/detection/type/detection2d/seg.py @@ -34,6 +34,25 @@ class Detection2DSeg(Detection2DBBox): mask: np.ndarray[Any, np.dtype[np.uint8]] # Binary mask [H, W], uint8 0 or 255 + def draw_on(self, img: Any, scale: float = 1.0) -> None: + """Blend the segmentation mask onto a BGR image, then draw its box and label.""" + mask = self.mask + if mask.shape[:2] != img.shape[:2]: + mask = cv2.resize( + mask, + (img.shape[1], img.shape[0]), + interpolation=cv2.INTER_NEAREST, + ) + + selected = mask > 0 + if np.any(selected): + mask_color = np.array([0, 180, 255], dtype=np.float32) + img[selected] = (img[selected].astype(np.float32) * 0.55 + mask_color * 0.45).astype( + np.uint8 + ) + + super().draw_on(img, scale=scale) + @classmethod def from_sam2_result( cls, diff --git a/dimos/perception/detection/type/imageDetections.py b/dimos/perception/detection/type/imageDetections.py index 98fd0e5388..4aa609cc3e 100644 --- a/dimos/perception/detection/type/imageDetections.py +++ b/dimos/perception/detection/type/imageDetections.py @@ -94,4 +94,4 @@ def annotated_image(self, scale: float = 1.0) -> Image: from dimos.msgs.sensor_msgs.Image import Image as ImageMsg - return ImageMsg.from_opencv(img, ts=self.image.ts) + return ImageMsg.from_opencv(img, frame_id=self.image.frame_id, ts=self.image.ts) diff --git a/dimos/perception/experimental/object_scene_registration.py b/dimos/perception/experimental/object_scene_registration.py index 914301afab..8d4cacaff1 100644 --- a/dimos/perception/experimental/object_scene_registration.py +++ b/dimos/perception/experimental/object_scene_registration.py @@ -13,14 +13,15 @@ # limitations under the License. import time -from typing import Any +from typing import Any, Literal import numpy as np from numpy.typing import NDArray +from pydantic import AliasChoices, Field, model_validator from dimos.agents.annotation import skill from dimos.core.core import rpc -from dimos.core.module import Module +from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo @@ -30,6 +31,7 @@ from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.perception.detection.detectors.base import Detector from dimos.perception.detection.detectors.yoloe import Yoloe2DDetector, YoloePromptMode from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.experimental.object import ( @@ -46,8 +48,36 @@ logger = setup_logger() +class ObjectSceneRegistrationConfig(ModuleConfig): + """Configurable detector, segmenter, and RGB-D object reconstruction settings.""" + + target_frame: str = "map" + prompt_mode: YoloePromptMode = YoloePromptMode.LRPC + distance_threshold: float = 0.2 + min_detections_for_permanent: int = 6 + register_objects: bool = True + detect_on_request: bool = False + detector_confidence: float = 0.6 + det: Literal["yoloe", "moondream"] = Field( + default="yoloe", validation_alias=AliasChoices("det", "detector_backend") + ) + seg: Literal["yolo", "edgetam"] = Field( + default="yolo", validation_alias=AliasChoices("seg", "segmentation_backend") + ) + object_voxel_downsample: float = 0.005 + max_distance: float = 0.0 + use_aabb: bool = False + max_obstacle_width: float = 0.0 + + @model_validator(mode="after") + def _require_edgetam_for_moondream(self) -> "ObjectSceneRegistrationConfig": + if self.det == "moondream" and self.seg != "edgetam": + raise ValueError("osr.det=moondream requires osr.seg=edgetam") + return self + + class ObjectSceneRegistrationModule(Module): - """Module for detecting objects in camera images using YOLO-E with 2D and 3D detection.""" + """Module for prompted 2D detection, segmentation, and RGB-D object reconstruction.""" color_image: In[Image] depth_image: In[Image] @@ -56,53 +86,64 @@ class ObjectSceneRegistrationModule(Module): detections_2d: Out[Detection2DArray] detections_3d: Out[Detection3DArray] + annotated_image: Out[Image] objects: Out[list[DetObject]] pointcloud: Out[PointCloud2] - _detector: Yoloe2DDetector | None = None + _detector: Detector | None = None + _segmenter: Any | None = None _camera_info: CameraInfo | None = None _object_db: ObjectDB + _latest_objects: list[Object] + _latest_output_objects: tuple[Object, ...] + _latest_aligned_frames: tuple[Image, Image] | None = None # A tuple assignment/read is atomic, so depth and its transform cannot be # observed from different frames by get_full_scene_pointcloud(). _latest_scene_snapshot: tuple[Image, Transform | None] | None = None + config: ObjectSceneRegistrationConfig - def __init__( - self, - target_frame: str = "map", - prompt_mode: YoloePromptMode = YoloePromptMode.LRPC, - # ObjectDB tuning - distance_threshold: float = 0.2, - min_detections_for_permanent: int = 6, - # Object 3D reconstruction tuning - max_distance: float = 0.0, - use_aabb: bool = False, - max_obstacle_width: float = 0.0, - **kwargs: Any, - ) -> None: + def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) - self._target_frame = target_frame - self._prompt_mode = prompt_mode + self._target_frame = self.config.target_frame + self._prompt_mode = self.config.prompt_mode + self._register_objects = self.config.register_objects + self._detect_on_request = self.config.detect_on_request + self._detector_confidence = self.config.detector_confidence + self._detector_backend = self.config.det + self._segmentation_backend = self.config.seg self._object_db = ObjectDB( - distance_threshold=distance_threshold, - min_detections_for_permanent=min_detections_for_permanent, + distance_threshold=self.config.distance_threshold, + min_detections_for_permanent=self.config.min_detections_for_permanent, ) - self._max_distance = max_distance - self._use_aabb = use_aabb - self._max_obstacle_width = max_obstacle_width + self._latest_objects = [] + self._latest_output_objects = () + self._object_voxel_downsample = self.config.object_voxel_downsample + self._max_distance = self.config.max_distance + self._use_aabb = self.config.use_aabb + self._max_obstacle_width = self.config.max_obstacle_width @rpc def start(self) -> None: super().start() - if self._prompt_mode == YoloePromptMode.LRPC: - model_name = "yoloe-11l-seg-pf.pt" + if self._detector_backend == "moondream": + from dimos.perception.detection.detectors.moondream import Moondream2DDetector + + self._detector = Moondream2DDetector() else: - model_name = "yoloe-11l-seg.pt" + if self._prompt_mode == YoloePromptMode.LRPC: + model_name = "yoloe-11l-seg-pf.pt" + else: + model_name = "yoloe-11l-seg.pt" + self._detector = Yoloe2DDetector( + model_name=model_name, + prompt_mode=self._prompt_mode, + conf=self._detector_confidence, + ) + if self._segmentation_backend == "edgetam": + from dimos.models.segmentation.edge_tam import EdgeTAMImageSegmenter - self._detector = Yoloe2DDetector( - model_name=model_name, - prompt_mode=self._prompt_mode, - ) + self._segmenter = EdgeTAMImageSegmenter() self.camera_info.subscribe(lambda msg: setattr(self, "_camera_info", msg)) @@ -121,8 +162,12 @@ def stop(self) -> None: if self._detector: self._detector.stop() self._detector = None + self._segmenter = None self._object_db.clear() + self._latest_objects = [] + self._latest_output_objects = () + self._latest_aligned_frames = None logger.info("ObjectSceneRegistrationModule stopped") super().stop() @@ -135,37 +180,69 @@ def set_prompts( ) -> None: """Set prompts for detection. Provide either text or bboxes, not both.""" if self._detector is not None: - self._detector.set_prompts(text=text, bboxes=bboxes) + set_prompts = getattr(self._detector, "set_prompts", None) + if not callable(set_prompts): + raise RuntimeError("configured detector does not support prompts") + set_prompts(text=text, bboxes=bboxes) @rpc def select_object(self, track_id: int) -> dict[str, Any] | None: """Get object data by track_id and promote to permanent.""" - for obj in self._object_db.get_all_objects(): + for obj in self._known_objects(): if obj.track_id == track_id: - self._object_db.promote(obj.object_id) + if self._register_objects: + self._object_db.promote(obj.object_id) return obj.to_dict() return None @rpc def get_object_track_ids(self) -> list[int]: """Get track_ids of all permanent objects.""" - return [obj.track_id for obj in self._object_db.get_all_objects()] + return [obj.track_id for obj in self._known_objects()] @rpc def get_detected_objects(self) -> list[dict[str, Any]]: """Get all detected objects with object_id (UUID) and name.""" - return [obj.agent_encode() for obj in self._object_db.get_all_objects()] + return [obj.agent_encode() for obj in self._known_objects()] + + @rpc + def scan_scene(self) -> Detection3DArray: + """Run detection on the latest aligned RGB-D frame and return its 3D detections.""" + frames = self._latest_aligned_frames + if frames is None: + return to_detection3d_array([], frame_id=self._target_frame) + + if not self._register_objects: + self._latest_objects = [] + self._latest_output_objects = () + self._process_images(*frames) + return to_detection3d_array( + list(self._latest_output_objects), + frame_id=self._target_frame, + ts=frames[0].ts, + ) + + @rpc + def describe_scene(self, question: str) -> str: + """Answer an open-ended scene question using the configured Moondream detector.""" + frames = self._latest_aligned_frames + if frames is None: + raise RuntimeError("No aligned RGB-D frame is available") + describe_image = getattr(self._detector, "describe_image", None) + if not callable(describe_image): + raise RuntimeError("Scene description requires osr.det=moondream") + return str(describe_image(frames[0], question)) @rpc def get_object_pointcloud_by_name(self, name: str) -> PointCloud2 | None: """Get pointcloud for an object by class name.""" - objects = self._object_db.find_by_name(name) + objects = [obj for obj in self._known_objects() if obj.name == name] return objects[0].pointcloud if objects else None @rpc def get_object_pointcloud_by_object_id(self, object_id: str) -> PointCloud2 | None: """Get pointcloud for an object by its stable object_id (searches all objects).""" - obj = self._object_db.find_by_object_id(object_id) + obj = next((obj for obj in self._known_objects() if obj.object_id == object_id), None) if obj is None: logger.warning(f"No object found with object_id='{object_id}'") return None @@ -178,7 +255,7 @@ def _get_object_mask(self, object_id: str) -> NDArray[np.uint8] | None: """Get dilated mask for an object by ID.""" import cv2 - for obj in self._object_db.get_all_objects(): + for obj in self._known_objects(): if obj.object_id != object_id: continue if obj.mask is None: @@ -193,6 +270,11 @@ def _get_object_mask(self, object_id: str) -> NDArray[np.uint8] | None: return None + def _known_objects(self) -> list[Object]: + if self._register_objects: + return self._object_db.get_all_objects() + return self._latest_objects + @rpc def get_full_scene_pointcloud( self, @@ -290,6 +372,9 @@ def select(self, track_id: int) -> str: def _on_aligned_frames(self, frames) -> None: # type: ignore[no-untyped-def] color_msg, depth_msg = frames + if self._detect_on_request: + self._latest_aligned_frames = (color_msg, depth_msg) + return self._process_images(color_msg, depth_msg) def _process_images(self, color_msg: Image, depth_msg: Image) -> None: @@ -308,8 +393,26 @@ def _process_images(self, color_msg: Image, depth_msg: Image) -> None: data=depth_cv, format=ImageFormat.DEPTH, frame_id=depth_msg.frame_id, ts=depth_msg.ts ) - # Run 2D detection + # Log each expensive stage separately so a stalled on-demand scan can be localized. + t0 = time.monotonic() + logger.info("Object detection started", detector=self._detector_backend) detections_2d: ImageDetections2D[Any] = self._detector.process_image(color_image) + logger.info( + "Object detection completed", + detector=self._detector_backend, + duration_s=round(time.monotonic() - t0, 3), + detections=len(detections_2d.detections), + ) + if self._segmenter is not None: + t0 = time.monotonic() + logger.info("Object segmentation started", segmenter=self._segmentation_backend) + detections_2d = self._segmenter.segment(detections_2d) + logger.info( + "Object segmentation completed", + segmenter=self._segmentation_backend, + duration_s=round(time.monotonic() - t0, 3), + detections=len(detections_2d.detections), + ) detections_2d_msg = Detection2DArray( detections_length=len(detections_2d.detections), @@ -317,6 +420,7 @@ def _process_images(self, color_msg: Image, depth_msg: Image) -> None: detections=[det.to_ros_detection2d() for det in detections_2d.detections], ) self.detections_2d.publish(detections_2d_msg) + self.annotated_image.publish(detections_2d.annotated_image()) # Process 3D detections self._process_3d_detections(detections_2d, color_image, depth_image) @@ -338,7 +442,9 @@ def _process_3d_detections( self._target_frame, color_image.frame_id, color_image.ts, - 0.1, + # Request-driven scans can use a cached camera frame while + # inference starts; retain temporal alignment within that cache. + 3.0, forward_tolerance=0.2, ) if camera_transform is None: @@ -354,25 +460,33 @@ def _process_3d_detections( depth_image=depth_image, camera_info=self._camera_info, camera_transform=camera_transform, + voxel_downsample=self._object_voxel_downsample, max_distance=self._max_distance, use_aabb=self._use_aabb, max_obstacle_width=self._max_obstacle_width, ) - if not objects: - return - - # Add objects to spatial memory database - self._object_db.add_objects(objects) + if self._register_objects: + if not objects: + return + self._object_db.add_objects(objects) + # Registered mode publishes the complete confirmed scene, not just this frame. + output_objects = self._object_db.get_objects() + else: + self._latest_objects = objects + output_objects = objects - # Publish ALL permanent objects so downstream consumers get the full set, - # not just this frame's batch (which may be a subset of what's on the table). - all_permanent = self._object_db.get_objects() + self._latest_output_objects = tuple(output_objects) - detections_3d = to_detection3d_array(all_permanent) + detections_3d = to_detection3d_array( + output_objects, + frame_id=self._target_frame, + ts=color_image.ts, + ) self.detections_3d.publish(detections_3d) - self.objects.publish(all_permanent) + self.objects.publish(output_objects) - objects_for_pc = all_permanent - aggregated_pc = aggregate_pointclouds(objects_for_pc) + aggregated_pc = aggregate_pointclouds(output_objects) + if not output_objects: + aggregated_pc.frame_id = self._target_frame + aggregated_pc.ts = color_image.ts self.pointcloud.publish(aggregated_pc) - return diff --git a/dimos/perception/experimental/object_scene_registration_spec.py b/dimos/perception/experimental/object_scene_registration_spec.py index 59aae79cab..c48f5548b4 100644 --- a/dimos/perception/experimental/object_scene_registration_spec.py +++ b/dimos/perception/experimental/object_scene_registration_spec.py @@ -15,10 +15,14 @@ from typing import Protocol from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray from dimos.spec.utils import Spec class ObjectSceneRegistrationSpec(Spec, Protocol): + def set_prompts(self, text: list[str] | None = None) -> None: ... + def scan_scene(self) -> Detection3DArray: ... + def describe_scene(self, question: str) -> str: ... def get_object_pointcloud_by_name(self, name: str) -> PointCloud2 | None: ... def get_object_pointcloud_by_object_id(self, object_id: str) -> PointCloud2 | None: ... def get_full_scene_pointcloud( diff --git a/dimos/perception/experimental/test_object_scene_registration_temporal.py b/dimos/perception/experimental/test_object_scene_registration_temporal.py index 99e79eb876..f36f4ce08e 100644 --- a/dimos/perception/experimental/test_object_scene_registration_temporal.py +++ b/dimos/perception/experimental/test_object_scene_registration_temporal.py @@ -17,12 +17,13 @@ from collections.abc import Iterator import sys from typing import Any -from unittest.mock import MagicMock +from unittest.mock import ANY, MagicMock import numpy as np import pytest from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.experimental.object_scene_registration import ObjectSceneRegistrationModule @@ -75,7 +76,56 @@ def test_temporal_tf_lookup_uses_bounded_image_timestamp( _image(12.5), ) - assert tf.calls == [(("map", "camera", 12.5, 0.1), {"forward_tolerance": 0.2})] + assert tf.calls == [(("map", "camera", 12.5, 3.0), {"forward_tolerance": 0.2})] + + +def test_detector_confidence_is_configurable() -> None: + module = ObjectSceneRegistrationModule(detector_confidence=0.4) + try: + assert module._detector_confidence == 0.4 + finally: + module.stop() + + +def test_segmentation_backend_defaults_to_yolo() -> None: + module = ObjectSceneRegistrationModule() + try: + assert module._segmentation_backend == "yolo" + finally: + module.stop() + + with pytest.raises(ValueError, match="seg"): + ObjectSceneRegistrationModule(segmentation_backend="invalid") # type: ignore[arg-type] + with pytest.raises(ValueError, match="det"): + ObjectSceneRegistrationModule(detector_backend="invalid") # type: ignore[arg-type] + with pytest.raises(ValueError, match="requires"): + ObjectSceneRegistrationModule(det="moondream", seg="yolo") + + +def test_edgetam_backend_refines_yolo_detections( + monkeypatch: Any, module: ObjectSceneRegistrationModule +) -> None: + color = Image( + data=np.zeros((2, 2, 3), dtype=np.uint8), + format=ImageFormat.BGR, + frame_id="camera", + ts=4.0, + ) + raw_detections = ImageDetections2D(color, []) + segmented_detections = ImageDetections2D(color, []) + module._detector = MagicMock() + module._detector.process_image.return_value = raw_detections + module._segmenter = MagicMock() + module._segmenter.segment.return_value = segmented_detections + module.detections_2d = MagicMock() + module.annotated_image = MagicMock() + process_3d = MagicMock() + monkeypatch.setattr(module, "_process_3d_detections", process_3d) + + module._process_images(color, _image(4.0)) + + module._segmenter.segment.assert_called_once_with(raw_detections) + process_3d.assert_called_once_with(segmented_detections, color, ANY) def test_failed_lookup_does_not_retry_without_time_or_replace_coherent_cache( @@ -106,7 +156,7 @@ def test_failed_lookup_does_not_retry_without_time_or_replace_coherent_cache( ) assert len(tf.calls) == 2 - assert tf.calls[1] == (("map", "camera", 2.0, 0.1), {"forward_tolerance": 0.2}) + assert tf.calls[1] == (("map", "camera", 2.0, 3.0), {"forward_tolerance": 0.2}) assert module._latest_scene_snapshot == (old_depth, old_transform) @@ -143,3 +193,97 @@ def voxel_down_sample(self, voxel_size: float) -> _PointCloud: module.get_full_scene_pointcloud() result.transform.assert_called_once_with(transform) + + +def test_process_images_publishes_annotated_detection_image( + monkeypatch: Any, module: ObjectSceneRegistrationModule +) -> None: + annotated = MagicMock(spec=Image) + detections = MagicMock(spec=ImageDetections2D) + detections.detections = [] + detections.annotated_image.return_value = annotated + module._detector = MagicMock() + module._detector.process_image.return_value = detections + module.detections_2d = MagicMock() + module.annotated_image = MagicMock() + process_3d = MagicMock() + monkeypatch.setattr(module, "_process_3d_detections", process_3d) + + color = Image( + data=np.zeros((2, 2, 3), dtype=np.uint8), + format=ImageFormat.BGR, + frame_id="camera", + ts=4.0, + ) + module._process_images(color, _image(4.0)) + + module.annotated_image.publish.assert_called_once_with(annotated) + process_3d.assert_called_once() + + +def test_live_mode_publishes_current_objects_without_registration( + monkeypatch: Any, +) -> None: + module = ObjectSceneRegistrationModule(target_frame="camera", register_objects=False) + module._camera_info = MagicMock() + module._object_db.add_objects = MagicMock() + module.detections_3d = MagicMock() + module.objects = MagicMock() + module.pointcloud = MagicMock() + detected_object = MagicMock() + pointcloud = MagicMock() + monkeypatch.setattr( + "dimos.perception.experimental.object_scene_registration.Object.from_2d_to_list", + lambda **_: [detected_object], + ) + monkeypatch.setattr( + "dimos.perception.experimental.object_scene_registration.to_detection3d_array", + lambda *_args, **_kwargs: MagicMock(), + ) + monkeypatch.setattr( + "dimos.perception.experimental.object_scene_registration.aggregate_pointclouds", + lambda _objects: pointcloud, + ) + + ObjectSceneRegistrationModule._process_3d_detections( + module, + MagicMock(spec=ImageDetections2D), + _image(4.0), + _image(4.0), + ) + + module._object_db.add_objects.assert_not_called() + module.objects.publish.assert_called_once_with([detected_object]) + module.pointcloud.publish.assert_called_once_with(pointcloud) + assert module._latest_objects == [detected_object] + module.stop() + + +def test_request_driven_scan_processes_latest_cached_frame(monkeypatch: Any) -> None: + module = ObjectSceneRegistrationModule(target_frame="camera", detect_on_request=True) + color = Image( + data=np.zeros((2, 2, 3), dtype=np.uint8), + format=ImageFormat.BGR, + frame_id="camera", + ts=4.0, + ) + depth = _image(4.0) + module._latest_aligned_frames = (color, depth) + output = MagicMock(frame_id="camera", ts=4.0) + + def process_images(got_color: Image, got_depth: Image) -> None: + assert (got_color, got_depth) == (color, depth) + module._latest_output_objects = (output,) + + detections = MagicMock(spec=Detection3DArray) + monkeypatch.setattr(module, "_process_images", process_images) + monkeypatch.setattr( + "dimos.perception.experimental.object_scene_registration.to_detection3d_array", + lambda *_args, **_kwargs: detections, + ) + + result = module.scan_scene() + + assert result is detections + assert module._latest_output_objects == (output,) + module.stop() diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index c67b10b549..3b786ce140 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -91,6 +91,8 @@ "openarm-planner-coordinator": "dimos.robot.manipulators.openarm.blueprints.planner:openarm_planner_coordinator", "openyam-planner-coordinator": "dimos.robot.manipulators.openyam.blueprints.basic:openyam_planner_coordinator", "path-planner-eval": "dimos.navigation.nav_3d.evaluator.blueprints:path_planner_eval", + "picknplace": "dimos.manipulation.blueprints:picknplace", + "picknplace-agent": "dimos.manipulation.blueprints:picknplace_agent", "teleop-hosted-go2-multicam": "dimos.teleop.hosted.blueprints.cloudflare:teleop_hosted_go2_multicam", "teleop-hosted-go2-transport": "dimos.teleop.hosted.blueprints.cloudflare:teleop_hosted_go2_transport", "teleop-hosted-xarm6": "dimos.teleop.hosted.blueprints.cloudflare:teleop_hosted_xarm6", @@ -151,6 +153,8 @@ "unitree-go2-webrtc-keyboard-teleop": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_webrtc_keyboard_teleop:unitree_go2_webrtc_keyboard_teleop", "unitree-go2-webrtc-rage-keyboard-teleop": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_webrtc_rage_keyboard_teleop:unitree_go2_webrtc_rage_keyboard_teleop", "unity-sim": "dimos.simulation.unity.blueprint:unity_sim", + "xarm-graspgenx": "dimos.robot.manipulators.xarm.blueprints.graspgenx:xarm_graspgenx", + "xarm-graspgenx-agent": "dimos.robot.manipulators.xarm.blueprints.agentic:xarm_graspgenx_agent", "xarm-perception": "dimos.robot.manipulators.xarm.blueprints.perception:xarm_perception", "xarm-perception-agent": "dimos.robot.manipulators.xarm.blueprints.agentic:xarm_perception_agent", "xarm-perception-sim": "dimos.robot.manipulators.xarm.blueprints.simulation:xarm_perception_sim", @@ -214,7 +218,7 @@ "goal-relay": "dimos.navigation.nav_3d.mls_planner.goal_relay.GoalRelay", "google-maps-skill-container": "dimos.agents.skills.google_maps_skill_container.GoogleMapsSkillContainer", "gps-nav-skill-container": "dimos.agents.skills.gps_nav_skill.GpsNavSkillContainer", - "grasping-module": "dimos.manipulation.grasping.grasping.GraspingModule", + "grasp-gen-x-module": "dimos.manipulation.grasping.grasp_gen_x.GraspGenXModule", "gstreamer-camera-module": "dimos.hardware.sensors.camera.gstreamer.gstreamer_camera.GstreamerCameraModule", "hosted-stats-module": "dimos.teleop.hosted.hosted_stats.HostedStatsModule", "joint-trajectory-controller": "dimos.manipulation.control.trajectory_controller.joint_trajectory_controller.JointTrajectoryController", @@ -256,6 +260,7 @@ "pgo": "dimos.navigation.cmu_nav.modules.pgo.pgo.PGO", "phone-teleop-module": "dimos.teleop.phone.phone_teleop_module.PhoneTeleopModule", "pick-and-place-module": "dimos.manipulation.pick_and_place_module.PickAndPlaceModule", + "pick-n-place-module": "dimos.manipulation.picknplace.PickNPlaceModule", "point-lio": "dimos.hardware.sensors.lidar.pointlio.module.PointLio", "pointlio-recorder": "dimos.hardware.sensors.lidar.pointlio.recorder.PointlioRecorder", "quest-teleop-module": "dimos.teleop.quest.quest_teleop_module.QuestTeleopModule", diff --git a/dimos/robot/manipulators/common/agent_prompts.py b/dimos/robot/manipulators/common/agent_prompts.py index a010bf8018..31e02c1507 100644 --- a/dimos/robot/manipulators/common/agent_prompts.py +++ b/dimos/robot/manipulators/common/agent_prompts.py @@ -59,14 +59,15 @@ ## Pick & Place - **pick **: Pick up a detected object by name. Use the EXACT name from \ look/scan_objects output. When duplicates exist, pass the object_id shown in brackets \ -(e.g. [id=abc12345]). Example: "pick the cup", "grab the spray can" +(e.g. [id=abc12345]). On GraspGenX-enabled stacks, pick ranks learned grasp proposals, \ +checks motion feasibility before moving, and can verify calibrated closure feedback. Example: \ +"pick the cup", "grab the spray can" - **place **: Place a held object at explicit world-frame coordinates. \ Example: "place it at 0.4, 0.3, 0.1" - **drop_on **: Drop a held object onto another detected object. \ Automatically compensates for camera occlusion. Example: "drop it in the bowl", \ "put it on the box" - **place_back**: Return a held object to its original pick position. -- **pick_and_place **: Pick then place in one command. ## Motion - **move_to_pose [roll pitch yaw]**: Move end-effector to an absolute \ @@ -99,13 +100,18 @@ - NEVER open the gripper while holding an object unless the user asks or you are \ executing place/drop_on. The gripper stays closed during movement. - After pick or place, return to init with **go_init** unless another action follows. +- If pick reports that the object may be held, do not open the gripper automatically. \ +Report the failure phase and ask the user before releasing or recovering. +- If pick fails before closure, call **reset** if the robot entered FAULT, then \ +**scan_objects** before retrying. Do not clear all perception obstacles merely to force \ +a plan through a changed scene. # Coordinate System World frame (meters): X = forward, Y = left, Z = up. Z = 0 is robot base. Typical working area: X 0.3-0.7, Y -0.5 to 0.5, Z 0.05-0.5. # Error Recovery -If planning fails with COLLISION_AT_START: call **clear_perception_obstacles**, then \ -**reset**, then retry. -After any planning failure, call **reset** before more planning or motion. +If planning fails with COLLISION_AT_START, inspect or rescan the scene. Clear perception \ +obstacles only when they are known to be stale. After any robot motion fault, call \ +**reset** before more planning or motion. """ diff --git a/dimos/robot/manipulators/xarm/blueprints/agentic.py b/dimos/robot/manipulators/xarm/blueprints/agentic.py index 073ede3607..465370a77b 100644 --- a/dimos/robot/manipulators/xarm/blueprints/agentic.py +++ b/dimos/robot/manipulators/xarm/blueprints/agentic.py @@ -24,6 +24,7 @@ MANIPULATION_AGENT_SYSTEM_PROMPT, ) from dimos.robot.manipulators.xarm.blueprints.basic import xarm7_planner_coordinator +from dimos.robot.manipulators.xarm.blueprints.graspgenx import xarm_graspgenx from dimos.robot.manipulators.xarm.blueprints.perception import xarm_perception from dimos.robot.manipulators.xarm.blueprints.simulation import xarm_perception_sim @@ -39,6 +40,12 @@ McpClient.blueprint(system_prompt=MANIPULATION_AGENT_SYSTEM_PROMPT), ) +xarm_graspgenx_agent = autoconnect( + xarm_graspgenx, + McpServer.blueprint(), + McpClient.blueprint(system_prompt=MANIPULATION_AGENT_SYSTEM_PROMPT), +) + xarm_perception_sim_agent = autoconnect( xarm_perception_sim, McpServer.blueprint(), diff --git a/dimos/robot/manipulators/xarm/blueprints/graspgenx.py b/dimos/robot/manipulators/xarm/blueprints/graspgenx.py new file mode 100644 index 0000000000..c753194c97 --- /dev/null +++ b/dimos/robot/manipulators/xarm/blueprints/graspgenx.py @@ -0,0 +1,61 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GraspGenX-enabled real-hardware xArm perception blueprint.""" + +from __future__ import annotations + +import math + +from dimos.core.coordination.blueprints import autoconnect +from dimos.manipulation.grasping.grasp_gen_x import GraspGenXModule +from dimos.manipulation.pick_and_place_module import PickAndPlaceModule +from dimos.robot.manipulators.xarm.blueprints.perception import xarm_perception +from dimos.robot.manipulators.xarm.config import make_xarm7_model_config +from dimos.robot.manipulators.xarm.grasp_config import make_xarm_graspgenx_config + +_graspgenx_config = make_xarm_graspgenx_config() + +xarm_graspgenx = autoconnect( + xarm_perception, + PickAndPlaceModule.blueprint( + robots=[ + make_xarm7_model_config( + name="arm", + add_gripper=True, + pitch=math.radians(45), + tf_extra_links=["link7"], + ) + ], + planning_timeout=10.0, + visualization={"backend": "meshcat"}, + floor_z=-0.02, + heuristic_grasp_fallback=False, + planning_frame="world", + grasp_approach_vector=(0.0, 0.0, -1.0), + grasp_verification={ + # Enable only after completing the hardware calibration recorded + # in the grasp-pipeline OpenSpec change. + "enabled": False, + "open_position": 0.85, + "closed_position": 0.0, + "held_threshold": 0.02, + "timeout": 2.0, + "poll_interval": 0.05, + }, + ), + GraspGenXModule.blueprint( + **_graspgenx_config.model_dump(exclude={"rpc_transport", "tf_transport", "g"}) + ), +).global_config(n_workers=5) diff --git a/dimos/robot/manipulators/xarm/blueprints/perception.py b/dimos/robot/manipulators/xarm/blueprints/perception.py index f187e7ab50..ead451ca37 100644 --- a/dimos/robot/manipulators/xarm/blueprints/perception.py +++ b/dimos/robot/manipulators/xarm/blueprints/perception.py @@ -45,6 +45,7 @@ planning_timeout=10.0, visualization={"backend": "meshcat"}, floor_z=-0.02, + heuristic_grasp_fallback=True, ), RealSenseCamera.blueprint( base_frame_id="link7", diff --git a/dimos/robot/manipulators/xarm/blueprints/simulation.py b/dimos/robot/manipulators/xarm/blueprints/simulation.py index fb6e21f09a..7e7d2fa1de 100644 --- a/dimos/robot/manipulators/xarm/blueprints/simulation.py +++ b/dimos/robot/manipulators/xarm/blueprints/simulation.py @@ -36,6 +36,7 @@ robots=[make_xarm7_sim_robot_config()], planning_timeout=10.0, visualization={"backend": "meshcat"}, + heuristic_grasp_fallback=True, ), MujocoSimModule.blueprint(**make_xarm7_sim_module_kwargs(XARM7_SIM_PATH)), ObjectSceneRegistrationModule.blueprint(target_frame="world"), diff --git a/dimos/robot/manipulators/xarm/grasp_config.py b/dimos/robot/manipulators/xarm/grasp_config.py new file mode 100644 index 0000000000..b70ced75a6 --- /dev/null +++ b/dimos/robot/manipulators/xarm/grasp_config.py @@ -0,0 +1,63 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GraspGenX geometry for the UFACTORY xArm gripper.""" + +from __future__ import annotations + +from dimos.manipulation.grasping.grasp_gen_x import ( + GraspGenXConfig, + SweepVolumeGripperConfig, +) + +# Geometry was derived from UFACTORY's xarm_ros gripper URDF and collision +# meshes at commit 0b5118eb6bf664fc3891c14b203e6ecbd5095dca: +# - link_tcp is 0.172 m along +Z from xarm_gripper_base_link +# - link_tcp's closing axis is 90 degrees counter-clockwise around local +Z +# from GraspGenX's local +X closing axis +# - the inner finger volume is approximately 0.085 x 0.032 x 0.067 m +# The model's grasp frame is the gripper base; DimOS plans for link_tcp. +XARM_GRASP_FRAME_TO_TCP = ( + (0.0, -1.0, 0.0, 0.0), + (1.0, 0.0, 0.0, 0.0), + (0.0, 0.0, 1.0, 0.172), + (0.0, 0.0, 0.0, 1.0), +) + +# Inverse of ``XARM_GRASP_FRAME_TO_TCP``. Rerun receives TCP poses, while the +# sweep geometry below is expressed in the GraspGenX gripper-base frame. +XARM_TCP_TO_GRASP_FRAME = ( + (0.0, 1.0, 0.0, 0.0), + (-1.0, 0.0, 0.0, 0.0), + (0.0, 0.0, 1.0, -0.172), + (0.0, 0.0, 0.0, 1.0), +) + +XARM_GRIPPER_SWEEP = SweepVolumeGripperConfig( + extents_open=(0.085, 0.032, 0.067), + offset_open=(0.0, 0.0, 0.1285), + extents_half_open=(0.0425, 0.032, 0.067), + offset_half_open=(0.0, 0.0, 0.1285), + fingertip_depth=0.162, + family="revolute_2f", +) + + +def make_xarm_graspgenx_config() -> GraspGenXConfig: + """Return the import-safe learned-grasp deployment configuration.""" + return GraspGenXConfig( + gripper=XARM_GRIPPER_SWEEP, + grasp_frame_to_tcp=XARM_GRASP_FRAME_TO_TCP, + max_candidates=100, + ) diff --git a/docs/capabilities/manipulation/agentic.md b/docs/capabilities/manipulation/agentic.md index 3b3b03f9ef..6d57f2a1b5 100644 --- a/docs/capabilities/manipulation/agentic.md +++ b/docs/capabilities/manipulation/agentic.md @@ -43,6 +43,56 @@ uv run dimos stop Use `dimos log -f` to follow the log while the run is active. +## Learned grasp-to-pick pipeline + +The real-hardware `xarm-graspgenx-agent` blueprint adds GraspGenX proposals to +the xArm perception stack. Install the optional runtime and start it with: + +```bash +uv sync --extra graspgenx --extra manipulation --inexact +uv run dimos run xarm-graspgenx-agent +``` + +`pick` remains the only high-level picking tool. It resolves one current +object, obtains that object's planning-frame point cloud, requests ranked +GraspGenX candidates, and rejects candidates that fail pre-grasp, grasp, or +retreat inverse kinematics. During planning, the selected target is +temporarily removed from the collision scene while all other obstacles remain +active. The selected candidate then runs through prepare, approach, grasp, +close, verify, and retreat phases. + +Use the stable object ID returned by `scan_objects` whenever names are +ambiguous. A name is accepted only when it identifies exactly one current +detection; an object-ID prefix must also be unique. Existing +`xarm-perception` and `xarm-perception-sim` blueprints retain their explicit +heuristic grasp fallback and do not load the optional GraspGenX runtime. + +The learned pipeline configuration lives in +`dimos/robot/manipulators/xarm/grasp_config.py`. It records the xArm gripper +sweep volume and the transform from GraspGenX's gripper frame to the planned +TCP. `PickAndPlaceModuleConfig` controls the planning frame, maximum point +cloud age, candidate-check limit, TCP approach direction, approach/retreat +offsets, heuristic fallback, and closure-feedback verification thresholds. +Changing the frame transform, approach direction, or closure threshold +requires robot-specific calibration. + +Failures are phase-specific and stop motion immediately. Before closure, a +failed transaction leaves the gripper in its current safe state. After a +successful close command, failures never automatically reopen the gripper; +the result includes `object_may_be_held=true`, and an operator or agent should +inspect state before issuing another motion. Target collision geometry is +restored on every exit path, and restoration errors are reported without +hiding the primary failure. + +The current verification is a closure-position proxy: an xArm gripper that +stops above the calibrated empty-close threshold is treated as holding +something. It does not measure grasp force, detect slip, or prove that the +intended object was acquired. Force/torque or tactile feedback is required for +those stronger guarantees. The shipped learned-grasp blueprint keeps this +proxy disabled until the open, empty-close, and representative held-object +positions have been measured on the target xArm; enable +`grasp_verification.enabled` only after recording that calibration. + ## Daily interaction For normal interactive use, start the human-friendly terminal client: diff --git a/docs/usage/visualization.md b/docs/usage/visualization.md index d3a147df21..5204f85706 100644 --- a/docs/usage/visualization.md +++ b/docs/usage/visualization.md @@ -89,6 +89,23 @@ if __name__ == "__main__": Every LCM stream, such as `color_image` (output by CameraModule), that uses a data type (like `Image`) that has a `.to_rerun` method will get rendered (`rr.log`) using the LCM topic as the rerun entity path. In other words: to render something, simply log it to a stream and it will automatically be available in rerun. +## Grasp Proposal Visualization + +The manipulation package includes an interactive Viser demo for inspecting the +banana object point cloud and the score-ranked GraspGenX proposals together: + +```bash +uv sync --extra manipulation --extra graspgenx +uv run python -m dimos.manipulation.demo_grasp_visualization \ + --max-candidates 20 +``` + +Open the URL printed by the command. The **Grasp / Object Cloud** and +**Grasp / Proposals** controls can be toggled independently or together from +their parent **All** control. Proposals are ordered from highest to lowest score +and colored from green to orange. Press Ctrl-C to stop the demo and release the +Viser and GraspGenX resources. + ## Performance Tuning ### Symptom: Slow Map Updates diff --git a/openspec/changes/add-grasp-pipeline-skill/.openspec.yaml b/openspec/changes/add-grasp-pipeline-skill/.openspec.yaml new file mode 100644 index 0000000000..f205fc727f --- /dev/null +++ b/openspec/changes/add-grasp-pipeline-skill/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-07-29 diff --git a/openspec/changes/add-grasp-pipeline-skill/README.md b/openspec/changes/add-grasp-pipeline-skill/README.md new file mode 100644 index 0000000000..29d33b2729 --- /dev/null +++ b/openspec/changes/add-grasp-pipeline-skill/README.md @@ -0,0 +1,3 @@ +# add-grasp-pipeline-skill + +Add an end-to-end manipulation skill that resolves an object point cloud, requests GraspGenX proposals, validates and ranks feasible candidates, and executes a safe pick. diff --git a/openspec/changes/add-grasp-pipeline-skill/design.md b/openspec/changes/add-grasp-pipeline-skill/design.md new file mode 100644 index 0000000000..5682ce8d87 --- /dev/null +++ b/openspec/changes/add-grasp-pipeline-skill/design.md @@ -0,0 +1,153 @@ +## Context + +GraspGenX is an import-safe, dedicated-worker module implementing `GraspGenSpec.propose_grasps(PointCloud2) -> GraspCandidateArray`. Object perception already implements `ObjectSceneRegistrationSpec`, including stable-ID/name point-cloud lookup, and emits world-frame `DetObject` instances consumed by `PickAndPlaceModule`. + +The current `pick` skill already owns planning, execution, gripper control, perception obstacle integration, and `place_back` state, but `_generate_grasps_for_pick` produces a single hand-authored pose. Its sequence begins moving after only the pre-grasp plan succeeds, uses fixed waits for gripper commands, and reports success without checking whether an object prevented full closure. + +The natural missing seam is orchestration inside `PickAndPlaceModule`: + +```text +object name / id + | + v +ObjectSceneRegistrationSpec -----> object PointCloud2 (world) + | + v +GraspGenSpec --------------------> ranked TCP candidates + | + v +PickAndPlaceModule + resolve -> validate -> feasibility gate -> execute -> verify -> retreat + | | + +------ planning world / coordinator -------+ +``` + +The model worker remains independent and optional. The high-level skill owns the transaction because it already owns manipulation state and can guarantee cleanup across planning, gripper, and obstacle mutations. + +## Goals / Non-Goals + +**Goals:** + +- Make `pick` a complete learned-grasp pipeline with precise, testable phase and failure semantics. +- Reuse the existing perception, proposal, planning, coordinator, and `SkillResult` interfaces. +- Reject bad candidates before physical motion where possible. +- Preserve non-target collision checking and guarantee planning-scene cleanup. +- Verify physical closure using gripper feedback on the initial xArm integration. +- Keep the high-level agent interface stable. + +**Non-Goals:** + +- Retraining, fine-tuning, or changing GraspGenX inference. +- Adding grasp-quality calibration across proposal backends. +- Visual-servoing or closed-loop pose correction during the final approach. +- Force/torque-based grasp verification, slip detection, or automatic regrasp after contact. +- General attached-object collision geometry during subsequent place motions; that should be a follow-up capability. +- Enabling GraspGenX on every manipulator blueprint in the first change. + +## Decisions + +### 1. Deepen `PickAndPlaceModule` instead of adding a peer skill container + +`PickAndPlaceModule` will declare injected perception and grasp proposal Spec attributes and will orchestrate the transaction behind its existing `pick` skill. This keeps planning state, obstacle state, execution, gripper control, and `_last_pick_pose` under one lifecycle owner. + +Alternative: create a separate `GraspPipelineSkillContainer` that calls manipulation RPCs. Rejected because the current manipulation API has no transaction-level Spec, would expose partially coordinated state across RPC threads, and would duplicate cleanup and error translation. + +### 2. Preserve `pick` as the public skill + +The existing signature remains `pick(object_name, object_id=None, robot_name=None)`. Internally, candidate generation becomes a provider strategy: injected `GraspGenSpec` first, with the existing heuristic generator only when explicitly enabled in `PickAndPlaceModuleConfig`. + +Alternative: add `grasp_pick` or `pick_with_graspgenx`. Rejected because agents would have to choose between overlapping high-level skills and the orchestration is backend-independent even though GraspGenX is the first provider. + +### 3. Resolve through the perception RPC, not the local detection snapshot + +The snapshot remains useful for agent display and obstacle synchronization, but the pipeline obtains the proposal input using `ObjectSceneRegistrationSpec` by stable ID or unique name. The returned point cloud and proposal header must match the configured planning frame (initially `world`). Frame mismatch is an error; this change does not add a hidden TF dependency. + +Name-only lookup must first establish uniqueness from the current detection snapshot. This avoids the existing perception RPC's “first matching name” behavior silently selecting the wrong duplicate. + +### 4. Separate candidate feasibility from physical execution + +Candidates remain in generator score order. For each candidate up to `max_grasp_candidates_to_check`, the pipeline: + +1. validates finite rigid-pose data and frame agreement; +2. derives pre-grasp and retreat poses using the configured approach-axis offset; +3. checks IK/collision feasibility for all three targets without dispatching motion; +4. chooses the first candidate passing the gate. + +The actual plans are regenerated from live state immediately before each phase because stored plans become stale after execution. A planning or execution failure after motion begins terminates the transaction rather than jumping to another candidate from a changed robot state. + +Alternative: attempt candidates sequentially and retry after any failure. Rejected because after the first approach the robot is no longer at the common evaluated start state, making retry safety and cleanup ambiguous. + +### 5. Treat target obstacle exclusion as transaction state + +Before feasibility checks, the pipeline calls the existing targeted `WorldMonitor.remove_object_obstacle(object_id)` path. Other perception and static obstacles remain. A `try/finally` transaction boundary refreshes perception obstacles on every return path. + +The obstacle monitor can receive live updates concurrently, so the implementation must ensure the target is not re-added during the exclusion window. The preferred extension is a scoped suppression API owned by `WorldObstacleMonitor` (for example, an object-ID suppression context managed under its existing lock), rather than repeatedly deleting the obstacle from orchestration code. + +Alternative: clear all perception obstacles. Rejected because it removes collision protection for the rest of the scene. Permanently delete the target obstacle. Rejected because failure paths would leave the planning world inconsistent. + +### 6. Model the pipeline as an explicit transaction + +A private transaction object records the selected object, proposal source, current phase, candidate rank/score, target-suppression handle, closure state, and cleanup status. A module lock rejects concurrent `pick` calls. The phases are: + +```text +RESOLVE -> PROPOSE -> SELECT -> PREPARE -> APPROACH -> GRASP + | + v + CLOSE -> VERIFY -> RETREAT -> DONE +``` + +No automatic rollback motion is promised. Before gripper closure, failures leave the gripper state explicit in the result. After closure, failures never auto-open the gripper because an object may be held. + +### 7. Use feedback-based closure verification with robot-specific configuration + +The initial xArm blueprint configures: + +- open and closed command endpoints in the units already expected by the coordinator path; +- a held-object closure threshold and comparison direction; +- command and verification timeouts plus polling interval. + +The pipeline first checks the close command result, then polls `get_gripper`. Reaching the empty-closed region is a verification failure; remaining beyond the held threshold is success. Configuration validation ensures the threshold lies between the open and closed endpoints. + +This is a contact proxy, not proof against slip. The result should say “grasp verified by gripper closure feedback,” not claim force or object identity verification. + +Alternative: fixed sleep followed by unconditional success. Rejected because it cannot distinguish an accepted command from a successful physical pick. + +### 8. Return structured phase-specific failures + +Extend manipulation errors with at least: + +- `GRASP_PROVIDER_UNAVAILABLE` +- `GRASP_INPUT_INVALID` +- `GRASP_FRAME_MISMATCH` +- `GRASP_VERIFICATION_FAILED` +- `PICK_BUSY` + +Existing `OBJECT_NOT_DETECTED`, `GRASP_GENERATION_FAILED`, `GRASP_ATTEMPTS_EXHAUSTED`, `PLANNING_FAILED`, `GRIPPER_FAILED`, execution errors, and timeouts remain applicable. Human-readable details include phase, candidate rank/score when selected, and whether the gripper may hold an object. + +## Risks / Trade-offs + +- [Single-view point clouds can produce geometrically plausible but poor grasps] → retain score ordering, validate scene feasibility, expose candidate rank/score, and leave visual servoing/regrasp for follow-up. +- [Gripper aperture is an imperfect grasp signal, especially for thin objects] → make thresholds robot-specific, test boundary behavior, and describe verification as a closure proxy. +- [Planning feasibility checks may be expensive across many GPU proposals] → cap candidates checked, stop at the first feasible candidate, and record rejection metrics for tuning. +- [The target can be re-added by asynchronous perception during a pick] → add scoped suppression inside the obstacle monitor under its lock and test live-update behavior. +- [A target-free collision world permits intended finger/object contact but cannot model post-grasp payload collisions] → keep all non-target obstacles and explicitly defer attached-object geometry. +- [GraspGenX increases GPU memory and startup time] → retain a dedicated worker, lazy optional runtime imports, and blueprint-level opt-in. +- [Planning can still fail after an earlier feasibility gate because the robot/world changed] → regenerate plans from live state and stop safely rather than retrying from an unanalysed state. + +## Migration Plan + +1. Add configuration, error codes, and private transaction/candidate helpers behind the unchanged `pick` signature. +2. Add scoped target-obstacle suppression and unit tests without enabling it in shipped blueprints. +3. Wire the perception and proposal Specs into `PickAndPlaceModule`; keep heuristic fallback explicitly enabled in legacy blueprints during transition. +4. Add a distinct GraspGenX-enabled xArm perception blueprint with xArm-specific gripper sweep/TCP and verification configuration, then compose its agentic variant. Keep the existing blueprint dependency footprint unchanged. +5. Validate in deterministic unit tests, recorded/replay perception, MuJoCo where sensor support permits, and finally real xArm hardware with a guarded test matrix. +6. Update the agent prompt, blueprint registry if a new runnable blueprint is introduced, and manipulation documentation. + +Rollback is blueprint-level: remove the GraspGenX module and restore explicit heuristic fallback. The public `pick` signature does not require caller migration. + +## Open Questions + +- What xArm closure threshold has been validated for the physical gripper, and does it need object-width-aware tolerance? +- Does the current perception obstacle monitor need to freeze only the target ID, or should it snapshot all obstacles for the short execution window? +- Is candidate feasibility via existing IK/collision APIs sufficiently predictive, or should the first version generate full approach/grasp/retreat paths in a cloned planning context? +- Which approach axis encoded by the GraspGenX TCP transform should define pre-grasp and retreat offsets for the configured xArm gripper? diff --git a/openspec/changes/add-grasp-pipeline-skill/proposal.md b/openspec/changes/add-grasp-pipeline-skill/proposal.md new file mode 100644 index 0000000000..c4b8d6e2ae --- /dev/null +++ b/openspec/changes/add-grasp-pipeline-skill/proposal.md @@ -0,0 +1,30 @@ +## Why + +GraspGenX can now produce ranked grasp poses, but the agent-facing `pick` skill still uses one heuristic pose and has no end-to-end path from a detected object to a verified physical pick. A pipeline is needed now to connect object point-cloud lookup, learned proposals, motion feasibility, collision-world handling, gripper actuation, and clear recovery semantics. + +## What Changes + +- Upgrade the existing `pick` skill to resolve a unique detected object and obtain its world-frame point cloud through the perception RPC interface. +- Request ranked TCP grasp candidates through `GraspGenSpec`, preserving the generator's score order while rejecting invalid or motion-infeasible candidates. +- Execute a safe pick sequence: pre-grasp approach, gripper open, grasp, gripper close, and retreat. +- Temporarily exclude only the target object from planning collisions while preserving all other scene obstacles, and restore a consistent planning scene on every exit path. +- Verify grasp closure using configured gripper feedback when available and return structured failure codes that distinguish perception, proposal, feasibility, execution, and verification failures. +- Keep the heuristic grasp path available only as an explicit configuration fallback for blueprints that do not include a grasp proposal module. +- Add a GraspGenX-enabled xArm perception manipulation blueprint with gripper-specific configuration, leaving the existing non-GPU blueprint available, and update the manipulation agent prompt to describe the learned-pick behavior. + +## Capabilities + +### New Capabilities + +- `grasp-pipeline-skill`: End-to-end behavior and failure semantics for resolving an object, proposing and selecting feasible grasps, executing a pick, and verifying the result. + +### Modified Capabilities + +None. + +## Impact + +- Affected modules: `PickAndPlaceModule`, `GraspGenSpec`, `ObjectSceneRegistrationSpec`, manipulation error types, and xArm perception blueprints/prompts. +- The new xArm learned-pick blueprint requires the `graspgenx` optional dependency, a dedicated GraspGenX worker, and robot-specific sweep-volume/TCP configuration; existing xArm blueprints remain runnable without that extra. +- Existing `pick(object_name, object_id, robot_name)` callers remain source-compatible; observed candidate selection and failure results become more precise. +- No change is proposed to GraspGenX inference itself or to generic motion-planner algorithms. diff --git a/openspec/changes/add-grasp-pipeline-skill/specs/grasp-pipeline-skill/spec.md b/openspec/changes/add-grasp-pipeline-skill/specs/grasp-pipeline-skill/spec.md new file mode 100644 index 0000000000..7ecf052a06 --- /dev/null +++ b/openspec/changes/add-grasp-pipeline-skill/specs/grasp-pipeline-skill/spec.md @@ -0,0 +1,92 @@ +## ADDED Requirements + +### Requirement: Unique object resolution +The `pick` skill SHALL resolve exactly one detected object before requesting grasp proposals. It SHALL prefer an explicitly supplied stable object ID, SHALL reject ambiguous ID prefixes or names, and SHALL return `OBJECT_NOT_DETECTED` when no object matches. + +#### Scenario: Object ID selects one of several same-named objects +- **WHEN** the caller supplies an object ID that uniquely identifies one detected object +- **THEN** the pipeline uses that object's point cloud regardless of other objects with the same name + +#### Scenario: Object name is ambiguous +- **WHEN** the caller supplies only a name and multiple current detections match it +- **THEN** the pipeline performs no robot motion and returns a failure that asks the caller to provide an object ID + +### Requirement: Proposal input and frame contract +The pipeline SHALL retrieve the selected object's `PointCloud2` through `ObjectSceneRegistrationSpec`, SHALL reject empty or stale input according to configured limits, and SHALL require the proposal frame to match the manipulation planning frame. It MUST NOT silently interpret a candidate in a different frame. + +#### Scenario: Valid world-frame point cloud +- **WHEN** perception returns a non-empty, sufficiently recent object point cloud in the manipulation planning frame +- **THEN** the pipeline passes that cloud unchanged to `GraspGenSpec.propose_grasps` + +#### Scenario: Proposal frame differs from planning frame +- **WHEN** the returned candidate array identifies a frame other than the configured manipulation planning frame +- **THEN** the pipeline performs no robot motion and returns a frame-mismatch failure + +### Requirement: Ranked feasibility selection +The pipeline SHALL examine candidates in descending generator-score order, up to a configurable attempt limit. It SHALL reject non-finite or malformed poses and candidates whose pre-grasp, grasp, or retreat targets fail kinematic or collision feasibility checks. Generator scores SHALL be treated only as relative ranking values, not calibrated probabilities. + +#### Scenario: Highest-scored candidate is infeasible +- **WHEN** the first candidate cannot satisfy approach or grasp feasibility and a lower-scored candidate can +- **THEN** the pipeline selects the first feasible lower-scored candidate without moving for the rejected candidate + +#### Scenario: No candidate is feasible +- **WHEN** every candidate within the configured attempt limit fails validation or feasibility +- **THEN** the pipeline performs no gripper closure, leaves the robot in a safe pre-pick state, and returns `GRASP_ATTEMPTS_EXHAUSTED` with rejection counts by reason + +### Requirement: Target-aware collision scene +The pipeline SHALL keep non-target scene obstacles active while checking and executing a pick. It SHALL exclude only the selected target object from collision checking for the grasp approach and SHALL restore a consistent perception-derived planning scene on success, failure, cancellation, or exception. + +#### Scenario: Target is registered as an obstacle +- **WHEN** the selected object is present in the planning world as a perception obstacle +- **THEN** the pipeline removes that object's obstacle before grasp feasibility checks while retaining all other object and static obstacles + +#### Scenario: Execution fails after target exclusion +- **WHEN** any later planning, execution, gripper, verification, or retreat step fails +- **THEN** cleanup refreshes or restores the perception obstacle state before the skill returns + +### Requirement: Safe pick execution +For a selected feasible candidate, the pipeline SHALL execute the ordered phases `PREPARE`, `APPROACH`, `GRASP`, `CLOSE`, `VERIFY`, and `RETREAT`. It SHALL stop at the first failed phase, SHALL report that phase in the result, and MUST NOT open the gripper automatically after closure because the robot may be holding the object. + +#### Scenario: Successful pick sequence +- **WHEN** all motion plans execute, gripper commands are accepted, verification succeeds, and retreat completes +- **THEN** the skill returns success including the selected candidate rank and score and stores the grasp pose for `place_back` + +#### Scenario: Retreat fails after closure +- **WHEN** gripper closure and verification succeed but retreat planning or execution fails +- **THEN** the skill returns a retreat failure, leaves the gripper closed, and reports that the object may still be held + +### Requirement: Grasp verification +The learned-pick blueprint SHALL configure gripper-feedback verification. Verification SHALL poll feedback until a configurable timeout and SHALL distinguish an object-blocked closure from a fully closed empty gripper using robot-specific command units and thresholds. + +#### Scenario: Feedback indicates an object is held +- **WHEN** the final gripper position remains on the configured held-object side of the closure threshold before timeout +- **THEN** verification succeeds and the pipeline proceeds to retreat + +#### Scenario: Feedback indicates an empty close +- **WHEN** the gripper reaches the configured empty-closed region +- **THEN** the pipeline returns `GRASP_VERIFICATION_FAILED`, leaves the gripper closed, and does not report a successful pick + +### Requirement: Explicit heuristic fallback +Blueprints without a grasp proposal provider SHALL fail learned-pick requests by default. A blueprint MAY explicitly enable the existing heuristic pose generator as a fallback, and the skill result SHALL identify when that fallback was used. + +#### Scenario: Grasp provider is unavailable and fallback is disabled +- **WHEN** `pick` is called without an injected `GraspGenSpec` +- **THEN** the skill performs no motion and returns `GRASP_PROVIDER_UNAVAILABLE` + +#### Scenario: Grasp provider is unavailable and fallback is enabled +- **WHEN** `pick` is called without an injected `GraspGenSpec` on a blueprint that explicitly enables heuristic fallback +- **THEN** the pipeline uses the heuristic candidate path and identifies the proposal source in its result + +### Requirement: Single active pick transaction +The module SHALL allow at most one pick pipeline transaction at a time. A concurrent request SHALL be rejected without changing motion, gripper, proposal, or planning-scene state. + +#### Scenario: Concurrent pick request +- **WHEN** a second `pick` call arrives while another pick transaction is active +- **THEN** the second call returns a busy failure and does not enter any pipeline phase + +### Requirement: Blueprint and agent exposure +The xArm perception manipulation stack SHALL compose `PickAndPlaceModule`, `ObjectSceneRegistrationModule`, and `GraspGenXModule` with compatible world-frame and gripper/TCP configuration. The agent prompt SHALL continue to expose `pick` as the single high-level picking skill and SHALL describe its object-ID disambiguation and failure recovery behavior. + +#### Scenario: Learned-pick blueprint is built +- **WHEN** the xArm learned-pick blueprint is constructed with the `graspgenx` extra installed +- **THEN** blueprint Spec injection resolves one perception provider and one grasp proposal provider for the pick module diff --git a/openspec/changes/add-grasp-pipeline-skill/tasks.md b/openspec/changes/add-grasp-pipeline-skill/tasks.md new file mode 100644 index 0000000000..a726ba6e48 --- /dev/null +++ b/openspec/changes/add-grasp-pipeline-skill/tasks.md @@ -0,0 +1,51 @@ +## 1. Contracts and Configuration + +- [x] 1.1 Add phase-specific grasp pipeline error codes to `ManipulationSkillError` and unit-test their `SkillResult` serialization/logging behavior. +- [x] 1.2 Add validated `PickAndPlaceModuleConfig` fields for provider fallback, planning frame, input age, candidate limit, pre-grasp/retreat offsets, and gripper-feedback verification. +- [x] 1.3 Declare optional injected `ObjectSceneRegistrationSpec` and `GraspGenSpec` dependencies on `PickAndPlaceModule`, and add blueprint build tests for present, absent, and ambiguous providers. +- [x] 1.4 Define private typed transaction, phase, candidate, rejection, and verification-result models so state that changes together is not stored in parallel fields. + +## 2. Target-Aware Planning Scene + +- [x] 2.1 Add a scoped target-object suppression API to `WorldObstacleMonitor` and its `WorldMonitor` facade using the existing monitor lock. +- [x] 2.2 Ensure live perception updates cannot re-add a suppressed target while other object obstacles continue to add/update normally. +- [x] 2.3 Ensure suppression exit refreshes or restores the target on normal return, exception, cancellation, and partial setup failure. +- [x] 2.4 Add deterministic unit tests covering nested/duplicate suppression requests, concurrent perception updates, failed obstacle mutations, and cleanup. + +## 3. Object Resolution and Candidate Selection + +- [x] 3.1 Implement unique object resolution by stable ID or unambiguous current name, returning actionable failures without motion. +- [x] 3.2 Retrieve and validate the selected object's point cloud through `ObjectSceneRegistrationSpec`, including non-empty data, timestamp age, and planning-frame checks. +- [x] 3.3 Call `GraspGenSpec.propose_grasps`, validate the candidate-array header and poses, and preserve stable descending generator-score order. +- [x] 3.4 Derive pre-grasp and retreat targets from the configured TCP approach axis and candidate pose. +- [x] 3.5 Implement no-motion feasibility gating for pre-grasp, grasp, and retreat targets, capped by configuration and reporting rejection counts by reason. +- [x] 3.6 Retain the heuristic generator only behind explicit fallback configuration and identify the selected proposal source in results. +- [x] 3.7 Add unit tests for duplicate names, ID prefixes, missing/stale/wrong-frame clouds, provider failures, malformed candidates, stable score ties, candidate limits, and lower-ranked feasible selection. + +## 4. Pick Transaction and Verification + +- [x] 4.1 Add a single-active-pick guard that rejects concurrent transactions without mutating robot, gripper, proposal, or obstacle state. +- [x] 4.2 Implement the `PREPARE`, `APPROACH`, `GRASP`, `CLOSE`, `VERIFY`, and `RETREAT` phase runner behind the existing `pick` signature. +- [x] 4.3 Check every planning, execution, wait, and gripper-command result; regenerate motion plans from live state at each phase and terminate after the first post-motion failure. +- [x] 4.4 Replace fixed grasp sleeps with timeout-bounded gripper feedback polling and robot-specific held/empty threshold evaluation. +- [x] 4.5 Preserve a closed gripper on every post-closure failure, include “object may be held” context, and store `_last_pick_pose` only after verified closure and successful retreat. +- [x] 4.6 Guarantee transaction and target-suppression cleanup through one exit path while retaining the primary failure if cleanup also fails. +- [x] 4.7 Add phase-by-phase unit tests for success, command rejection, planning/execution failure, timeout, empty close, retreat failure, cleanup failure, and concurrent calls. + +## 5. Blueprint and Agent Integration + +- [x] 5.1 Define reviewed xArm sweep-volume, grasp-frame-to-TCP, approach-axis, and closure-verification configuration without performing model or hardware work at import time. +- [x] 5.2 Add a distinct GraspGenX-enabled xArm perception blueprint and agentic blueprint that compose exactly one perception provider, proposal provider, manipulation module, MCP server, and MCP client. +- [x] 5.3 Keep existing xArm perception blueprints free of the `graspgenx` runtime requirement and explicitly configure their intended heuristic fallback behavior. +- [x] 5.4 Update the manipulation agent prompt to keep `pick` as the sole high-level pick tool and document exact-name/object-ID disambiguation plus safe recovery. +- [x] 5.5 Regenerate `dimos/robot/all_blueprints.py` through `test_all_blueprints_generation.py` and verify the new names appear in `dimos list`. + +## 6. End-to-End Validation and Documentation + +- [x] 6.1 Add integration tests with fake perception, proposal, planner/coordinator, and gripper feedback providers that exercise the full RPC/Spec-wired pipeline. +- [x] 6.2 Add a replay or fixture-based test proving object/proposal/planning frame consistency with real `PointCloud2` and `GraspCandidateArray` messages. +- [ ] 6.3 Validate the GraspGenX-enabled blueprint startup and one successful/one infeasible candidate flow with the `graspgenx` extra on a GPU-capable environment. +- [ ] 6.4 Calibrate and record the xArm empty-close versus held-object threshold across representative object widths before enabling verification on hardware. +- [ ] 6.5 Run guarded real-xArm tests for success, empty grasp, unreachable proposals, execution interruption, and retreat failure; verify the gripper never auto-opens after closure. +- [x] 6.6 Update manipulation capability documentation with architecture, configuration, failure semantics, and the distinction between closure-proxy verification and force/slip verification. +- [x] 6.7 Run focused pytest suites, blueprint-generation validation, formatting/lint checks, and mypy on touched modules. diff --git a/openspec/changes/validate-connected-grasp-sequences/.openspec.yaml b/openspec/changes/validate-connected-grasp-sequences/.openspec.yaml new file mode 100644 index 0000000000..f205fc727f --- /dev/null +++ b/openspec/changes/validate-connected-grasp-sequences/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-07-29 diff --git a/openspec/changes/validate-connected-grasp-sequences/design.md b/openspec/changes/validate-connected-grasp-sequences/design.md new file mode 100644 index 0000000000..a9f7c79bf9 --- /dev/null +++ b/openspec/changes/validate-connected-grasp-sequences/design.md @@ -0,0 +1,131 @@ +## Context + +`PickAndPlaceModule` currently gates ranked grasp candidates by solving collision-aware IK independently for pre-grasp, grasp, and retreat poses. This proves that each pose is individually reachable, but not that collision-free paths connect the robot's current state to those poses in order. + +The generic planner contract already accepts explicit `start` and `goal` joint states. The RoboPlan adapter is the exception: it rejects a valid explicit start when it differs from the authoritative live state, even though RoboPlan's native RRT accepts arbitrary start and goal configurations. The pick pipeline also performs an optional low-height safety lift immediately before approach, so candidate validation must account for that shared segment. + +Upstream perception remains responsible for supplying a target-only segmented cloud in the manipulation planning frame. This change starts at grasp proposals and does not add detection or segmentation behavior. + +## Goals / Non-Goals + +**Goals:** + +- Require a connected, collision-free sequence before accepting a grasp candidate. +- Align RoboPlan with the explicit-start planner contract without mutating authoritative live state. +- Include a required safety lift as one shared prerequisite rather than repeating it for every candidate. +- Keep validation motion-free and retain execution-time replanning from measured state. +- Validate the behavior with compact, deterministic fixtures and guarded GraspGenX integration coverage. +- Provide a contributor-facing offline demo of the complete proposal-and-planning pipeline. + +**Non-Goals:** + +- Snapshotting or freezing the planning scene across validation calls. +- Attaching generated object geometry to the gripper for retreat planning. +- Coordinated arm-and-gripper planning or validation of finger sweeps. +- Reusing dry-run paths for physical execution. +- Adding perception, point-cloud segmentation, a new timeout, or backend-specific public error codes. + +## Decisions + +### 1. RoboPlan will honor an explicit planning start + +Both robot-scoped and selected-group RoboPlan planning entry points will pass the supplied normalized start and goal to native RRT. The live planning context remains authoritative for unselected joints, other robots, and the rest of the scene. Planning from a hypothetical selected-joint start MUST NOT update the live context. + +Alternative: clone and overwrite the whole live scene for every request. Rejected because RoboPlan already represents the selected planning group's start explicitly, while the live scene is still needed for unselected state. + +### 2. Connected validation will be a side-effect-free planning helper + +A manipulation planning helper will accept an ordered pose sequence and an optional explicit selected-joint start. For each pose it will: + +1. solve collision-aware IK seeded by the current sequence endpoint; +2. plan from that endpoint to the IK solution; +3. use the returned path endpoint as the next segment's IK seed and planning start. + +It will return the first failed pose index and the final endpoint on success. It will not store a preview/execution plan, change manipulation state, or dispatch motion. + +Alternative: call the existing stateful `plan_to_pose` API three times. Rejected because failed screening would fault manipulation state, overwrite the execution plan, and prevent clean evaluation of lower-ranked candidates. + +### 3. The safety lift is planned once as shared preparation + +The current low-height check will be factored so the required lift target can be computed without moving. If no lift is required, candidate validation begins at authoritative current state. If a lift is required, it is dry-run planned once: + +- failure aborts the pick in `PREPARE` with `PLANNING_FAILED`; +- success supplies a common hypothetical endpoint for every candidate. + +The lift failure is not a candidate rejection because changing candidates cannot make the shared preparation segment feasible. + +### 4. Candidate validation is connected and ordered + +For each candidate in generator-score order, the pipeline derives pre-grasp and retreat poses and validates: + +```text +shared start -> pre-grasp -> grasp -> retreat +``` + +The first failed segment increments the existing stage-level rejection counter. A candidate is selected only after all three paths succeed. No motion or gripper command occurs while candidates are screened. + +Alternative: retain independent collision-aware IK checks. Rejected because disconnected feasible configurations do not establish an executable grasp sequence. + +### 5. Execution replans rather than reusing validation paths + +The selected candidate retains its poses, not the dry-run paths. Physical execution continues to plan each phase from fresh measured state immediately before dispatch. If an execution-time plan fails after motion begins, the transaction stops and does not try a different candidate. + +Alternative: execute the exact dry-run paths. Rejected because controller error, settling, preparation motion, and live scene changes can make stored paths stale. + +### 6. MVP validation uses the latest live scene and bounded collision fidelity + +Every segment uses the latest available non-target scene. The target remains suppressed during validation and execution, but non-target updates are not frozen. Retreat collision checking covers the robot and gripper state represented in the planning scene; it does not include attached-object geometry. The dry run does not model separate open and closed finger configurations. + +These limits are explicit so a successful result is not described as payload-safe or as a validated finger sweep. + +### 7. Reporting remains stable and stage-oriented + +Public rejection counts remain `pre_grasp_infeasible`, `grasp_infeasible`, and `retreat_infeasible`. IK status, planner status, timeout, and no-path details are logged for diagnosis rather than added to the skill-result contract. + +### 8. Tests use a small reusable fixture set + +Tests will be layered: + +- orchestration unit tests with deterministic IK and planner results; +- RoboPlan contract tests proving arbitrary starts and live-state preservation; +- an open scene plus one parameterized blocker for segment failures; +- a recorded target-only cloud for GraspGenX structural smoke coverage; +- one guarded GraspGenX-to-real-planner integration requiring at least one complete sequence, without asserting exact stochastic poses. + +This avoids maintaining a separate fixture file for every rejection case. + +### 9. The MVP includes an offline pipeline demo + +A dedicated `dimos.manipulation.demo_grasp_pipeline` CLI will load the recorded target-only cloud, relocate it into the synthetic xArm workspace, run real GraspGenX, evaluate ranked candidates with the same connected planning primitive used by `pick`, and stop after selecting a complete plan. It will never create a coordinator client or dispatch execution. + +The demo will write a human-readable terminal summary plus machine-readable JSON containing candidate scores, rejection stages, the selected candidate, and every safety-lift/pre-grasp/grasp/retreat joint path. It will also render the relocated cloud with the selected grasp. +The selected candidate and saved plan remain in the robot TCP frame. Because the +GraspGenX sweep geometry is defined in its grasp/base frame, visualization converts +the TCP pose back through the configured `grasp_frame_to_tcp` transform before +drawing the gripper. + +Alternative: direct contributors to the guarded pytest integration. Rejected because a test runner does not expose candidate-level outcomes or reusable plan artifacts and is not an adequate manual validation interface. + +## Risks / Trade-offs + +- [A scene update can occur between connected segment checks] → Use the latest scene for every segment and replan again before execution. +- [Execution-time IK can choose a different configuration than validation] → Seed from fresh measured state and require a complete collision-free execution plan before each dispatch. +- [Retreat can be safe for the robot but unsafe for the held object's volume] → State the MVP limitation explicitly and defer attached-object planning. +- [Current gripper geometry may differ from approach or retreat geometry] → Do not claim finger-sweep validation; add coordinated gripper planning only as a later capability. +- [Ranked screening multiplies planner calls] → Retain the configured candidate cap and existing planner timeout; planning is expected to be cheap for the MVP. +- [GraspGenX output may be nondeterministic] → Assert structural and feasibility invariants rather than exact candidate poses. + +## Migration Plan + +1. Land explicit-start RoboPlan behavior with backend contract tests. +2. Add the side-effect-free connected-sequence planning helper. +3. Integrate shared safety-lift and candidate-sequence validation. +4. Add deterministic and guarded integration coverage. +5. Add the offline pipeline demo and artifact output. +6. Enable the behavior through the existing `pick` pipeline with no caller migration. + +Rollback restores independent IK screening and the RoboPlan live-start guard. The public `pick` signature and configuration remain unchanged. + +## Open Questions + +No material MVP design questions remain. Payload attachment, coordinated gripper geometry, and physical gripper-threshold tuning are explicitly deferred follow-up work. diff --git a/openspec/changes/validate-connected-grasp-sequences/proposal.md b/openspec/changes/validate-connected-grasp-sequences/proposal.md new file mode 100644 index 0000000000..2d1a28ee89 --- /dev/null +++ b/openspec/changes/validate-connected-grasp-sequences/proposal.md @@ -0,0 +1,30 @@ +## Why + +The grasp pipeline currently treats independently collision-free IK targets as sufficient evidence that a candidate is feasible. A candidate can therefore pass selection even when no connected motion exists from the current robot state through pre-grasp, grasp, and retreat. + +## What Changes + +- Make RoboPlan honor the planner contract's explicit start state instead of requiring it to match the live scene state. +- Validate an optional shared safety lift before evaluating candidate-specific motion. +- Accept a grasp candidate only when connected plans succeed from the shared start through pre-grasp, grasp, and retreat, with each segment starting at the preceding segment's endpoint. +- Keep validation motion-free and replan every physical execution segment from fresh measured state. +- Preserve stage-level candidate rejection reasons while logging detailed IK and planner outcomes. +- Add deterministic unit and backend-contract coverage plus a compact GraspGenX-to-planner integration test strategy. +- Add a no-hardware CLI demo that runs the complete proposal-and-planning pipeline and saves inspectable candidate and path artifacts. + +## Capabilities + +### New Capabilities + +- `connected-grasp-sequence-validation`: Defines explicit-start planning, shared preparation, connected candidate validation, execution replanning, failure reporting, and MVP collision-model boundaries. + +### Modified Capabilities + +None. + +## Impact + +- Affects RoboPlan planner adaptation, manipulation planning helpers, and `PickAndPlaceModule` candidate selection. +- Extends the pending grasp-pipeline behavior without changing the public `pick` signature or adding configuration. +- Uses the existing planning timeout and planning-scene update behavior. +- Adds no new runtime dependency; GPU-backed GraspGenX validation remains guarded by the existing optional dependency and test environment. diff --git a/openspec/changes/validate-connected-grasp-sequences/specs/connected-grasp-sequence-validation/spec.md b/openspec/changes/validate-connected-grasp-sequences/specs/connected-grasp-sequence-validation/spec.md new file mode 100644 index 0000000000..8b63adf402 --- /dev/null +++ b/openspec/changes/validate-connected-grasp-sequences/specs/connected-grasp-sequence-validation/spec.md @@ -0,0 +1,78 @@ +## ADDED Requirements + +### Requirement: Explicit-start RoboPlan planning +RoboPlan planning SHALL accept a valid explicit selected-joint start that differs from the authoritative live robot state. The requested start SHALL seed native planning, while unselected joints and other robots SHALL retain their latest live scene state. Planning MUST NOT mutate authoritative live state. + +#### Scenario: Hypothetical start differs from live state +- **WHEN** a caller requests a path from a valid selected-joint start that differs from the live selected-joint state +- **THEN** RoboPlan attempts the path from the requested start and leaves the live state unchanged + +#### Scenario: Unselected scene state is retained +- **WHEN** a selected planning group is planned from a hypothetical start +- **THEN** unselected joints, other robots, and non-target scene geometry retain their latest live values during collision checking + +### Requirement: Shared safety-lift validation +The pick pipeline SHALL determine whether its existing safety lift is required before candidate evaluation. It SHALL dry-run a required lift exactly once and SHALL use the successful lift endpoint as the common start for candidate validation. + +#### Scenario: Required safety lift is feasible +- **WHEN** the end effector requires a safety lift and a connected lift plan succeeds +- **THEN** every grasp candidate is evaluated from the planned lift endpoint without executing the lift during validation + +#### Scenario: Required safety lift is infeasible +- **WHEN** the end effector requires a safety lift and no lift plan succeeds +- **THEN** the pick aborts in `PREPARE` with `PLANNING_FAILED` and does not attribute the failure to any candidate + +### Requirement: Connected candidate sequence +The pipeline SHALL accept a grasp candidate only when connected plans succeed in order from the shared start to pre-grasp, from the pre-grasp path endpoint to grasp, and from the grasp path endpoint to retreat. Each IK solve SHALL be seeded from the preceding endpoint, and each plan SHALL begin at that same endpoint. + +#### Scenario: All candidate segments connect +- **WHEN** pre-grasp, grasp, and retreat each have a collision-free path beginning at the preceding path endpoint +- **THEN** the candidate passes motion-feasibility validation + +#### Scenario: Individually reachable poses are disconnected +- **WHEN** all three target poses have collision-free IK solutions but any required connecting path fails +- **THEN** the candidate is rejected without robot motion + +#### Scenario: Higher-ranked candidate has a blocked segment +- **WHEN** a higher-ranked candidate has a failed connected segment and a lower-ranked candidate has a complete connected sequence +- **THEN** the pipeline selects the lower-ranked candidate without moving for the rejected candidate + +### Requirement: Motion-free validation and fresh execution planning +Candidate validation SHALL NOT store an execution plan, dispatch robot motion, or issue gripper commands. After selection, every physical motion segment SHALL be replanned from freshly measured state before execution. + +#### Scenario: Candidate is selected +- **WHEN** a candidate passes connected dry-run validation +- **THEN** no validation path is executed and approach planning begins from fresh measured state + +#### Scenario: Execution replan fails after motion begins +- **WHEN** any execution-time replan fails after the robot has moved +- **THEN** the transaction stops at that phase and does not attempt a different grasp candidate + +### Requirement: Live-scene MVP collision contract +Each validation segment SHALL use the latest available planning scene with the selected target suppressed and non-target obstacles active. The MVP SHALL NOT claim atomic scene validation, attached-object clearance, or separate open-versus-closed finger-sweep validation. + +#### Scenario: Non-target scene changes during validation +- **WHEN** a non-target obstacle update arrives between two segment planning calls +- **THEN** the later segment uses the updated scene without restarting or freezing the complete validation sequence + +#### Scenario: Retreat validation succeeds +- **WHEN** the grasp-to-retreat path is collision-free for the robot and gripper state represented in the planning scene +- **THEN** retreat passes the MVP gate without asserting clearance for attached target geometry + +### Requirement: Stage-level rejection reporting +The pipeline SHALL report candidate rejection by the first failed sequence stage using `pre_grasp_infeasible`, `grasp_infeasible`, or `retreat_infeasible`. Backend-specific IK and planner outcomes SHALL remain diagnostic details and SHALL NOT become new public rejection categories in this change. + +#### Scenario: Grasp segment planning fails +- **WHEN** pre-grasp succeeds and the pre-grasp-to-grasp segment fails through IK, timeout, collision, or no path +- **THEN** the pipeline increments `grasp_infeasible` and records the detailed cause in diagnostic logs + +### Requirement: Offline full-pipeline demo +The repository SHALL provide a no-hardware CLI demo that runs a recorded target-only cloud through GraspGenX proposal generation and connected RoboPlan validation. The demo SHALL stop before execution and SHALL save candidate outcomes, the selected candidate, and all connected joint paths as inspectable artifacts. + +#### Scenario: Contributor runs the pipeline demo +- **WHEN** the recorded fixture, GraspGenX runtime, model, CUDA, and xArm planning data are available +- **THEN** the command selects only a candidate with a complete connected sequence and writes a summary, joint paths, and selected-grasp visualization without dispatching robot motion + +#### Scenario: No candidate has a complete sequence +- **WHEN** every candidate within the configured demo limit fails connected validation +- **THEN** the command writes the rejection summary, exits unsuccessfully, and does not issue motion or gripper commands diff --git a/openspec/changes/validate-connected-grasp-sequences/tasks.md b/openspec/changes/validate-connected-grasp-sequences/tasks.md new file mode 100644 index 0000000000..6fc3b3b2d4 --- /dev/null +++ b/openspec/changes/validate-connected-grasp-sequences/tasks.md @@ -0,0 +1,37 @@ +## 1. Align RoboPlan Explicit-Start Planning + +- [x] 1.1 Update robot-scoped and selected-group RoboPlan planning to pass any valid normalized explicit start to native RRT while retaining live unselected scene state. +- [x] 1.2 Add backend contract tests proving hypothetical starts are honored, paths begin at the requested start, and authoritative live robot state is not mutated. +- [x] 1.3 Extend multi-robot/auxiliary-joint coverage to prove hypothetical selected-group planning preserves other live scene state. + +## 2. Add Connected Dry-Run Planning + +- [x] 2.1 Add a side-effect-free manipulation helper that checks an ordered pose sequence from current state or an explicit selected-joint start and returns the first failed index plus the final endpoint. +- [x] 2.2 Chain each successful path endpoint into both the next IK seed and the next planner start, using collision-aware IK and the existing planning timeout. +- [x] 2.3 Log stage-specific IK and planner failure details without storing a generated execution plan or changing manipulation state. +- [x] 2.4 Add unit tests for complete success, IK failure, path failure, malformed/empty planner paths, explicit sequence starts, endpoint chaining, and no state or plan mutation. + +## 3. Integrate Shared Preparation and Candidate Validation + +- [x] 3.1 Factor the existing low-height safety-lift target calculation from its physical execution so it can be dry-run planned. +- [x] 3.2 Plan a required safety lift once in `PREPARE`, abort with `PLANNING_FAILED` on failure, and pass its endpoint as the shared start for all candidates. +- [x] 3.3 Replace independent pre-grasp/grasp/retreat IK screening with connected sequence checks in ranked candidate order. +- [x] 3.4 Preserve `pre_grasp_infeasible`, `grasp_infeasible`, and `retreat_infeasible` counters based on the first failed segment while keeping backend details diagnostic-only. +- [x] 3.5 Verify candidate screening issues no motion or gripper commands and that physical execution still replans every segment from fresh measured state. +- [x] 3.6 Add pick-pipeline tests for no-lift and shared-lift starts, lift failure, each failed candidate segment, lower-ranked fallback, exhausted candidates, and execution-time replan failure. + +## 4. Validate MVP Scene and Integration Boundaries + +- [x] 4.1 Add compact open-scene and parameterized-blocker coverage for successful sequences and blocked pre-grasp, grasp, and retreat paths without creating separate scene fixtures per case. +- [x] 4.2 Verify connected validation keeps the target suppressed, retains non-target obstacles, and allows later segments to observe live scene updates. +- [x] 4.3 Document and test the MVP boundary that retreat excludes attached-object geometry and uses the gripper state currently represented in the planning scene. +- [x] 4.4 Add a guarded GraspGenX smoke test using one recorded target-only planning-frame cloud and assert structural proposal invariants rather than exact stochastic poses. +- [x] 4.5 Add a guarded GraspGenX-to-real-planner integration test requiring at least one complete connected sequence in the reusable open scene. +- [x] 4.6 Run focused manipulation tests, RoboPlan tests, static typing, formatting, and lint checks; record environment-only skips for optional GPU integration. + +## 5. Add Offline Full-Pipeline Demo + +- [x] 5.1 Expose connected dry-run segment paths without storing an execution plan or changing manipulation state. +- [x] 5.2 Add an import-safe CLI that loads and relocates the recorded target-only cloud, runs real GraspGenX and xArm RoboPlan, and never dispatches hardware motion. +- [x] 5.3 Save candidate rejection outcomes, selected-candidate metadata, safety-lift/pre-grasp/grasp/retreat joint paths, and a selected-grasp visualization under a caller-selected output directory. +- [x] 5.4 Add hermetic demo tests for successful selection, exhausted candidates, artifact contents, CLI wiring, and no execution/gripper calls; run focused and static quality gates. diff --git a/openspec/specs/manipulation-visualization-layers/spec.md b/openspec/specs/manipulation-visualization-layers/spec.md new file mode 100644 index 0000000000..876d479d44 --- /dev/null +++ b/openspec/specs/manipulation-visualization-layers/spec.md @@ -0,0 +1,148 @@ +# Manipulation Visualization Layers + +## Purpose + +Define display-only manipulation visualization layers, their backend-neutral +geometry models, Viser lifecycle and rendering behavior, and the standalone +banana grasp-proposal visualization workflow. + +## Requirements + +### Requirement: Display-only visualization layer interface +The manipulation visualization interface SHALL accept complete `VisualizationLayer` replacements through `set_layer` and SHALL clear a layer's rendered contents through `clear_layer`. These operations MUST NOT mutate the planning world, create collision geometry, or affect collision and motion-planning results. + +#### Scenario: Visual layer is published +- **WHEN** a caller submits a valid visualization layer +- **THEN** the visualization backend accepts the layer for display without invoking planning-world mutation + +#### Scenario: Layer is cleared +- **WHEN** a caller clears a registered visualization layer +- **THEN** its visual elements disappear while its registration and visibility preference remain + +#### Scenario: Visualization and collision representations share a source object +- **WHEN** a display-only element and a planning obstacle describe the same physical object +- **THEN** adding, replacing, clearing, or failing to render the visual element does not add, replace, remove, or otherwise modify the planning obstacle + +### Requirement: Backend-neutral layer and element models +A `VisualizationLayer` SHALL have a nonempty stable hierarchical ID, one nonempty coordinate frame, a first-registration default visibility value, and elements with IDs unique within that layer. The initial supported element types SHALL be point clouds and line sets represented by owned NumPy array snapshots rather than backend handles or robotics-domain messages. + +#### Scenario: Valid point cloud +- **WHEN** a point-cloud element contains finite `N x 3` points and either no colors or matching RGB colors +- **THEN** the model snapshots and accepts the element + +#### Scenario: Valid line set +- **WHEN** a line-set element contains finite `N x 3` vertices and `M x 2` in-range vertex indices +- **THEN** the model snapshots and accepts the element + +#### Scenario: Caller mutates source arrays +- **WHEN** a caller changes a source NumPy array after constructing or submitting an element +- **THEN** the accepted visual snapshot remains unchanged + +#### Scenario: Invalid geometry +- **WHEN** an element has non-finite coordinates, inconsistent color counts, invalid edge indices, or a non-positive explicit size +- **THEN** the layer is rejected before any current rendered generation is modified + +#### Scenario: Duplicate element identity +- **WHEN** one layer contains two elements with the same ID +- **THEN** the layer is rejected + +### Requirement: Complete and atomic layer replacement +Publishing a layer with an existing ID SHALL replace its complete contents. Viser MUST expose either the previous valid generation or the complete new generation and MUST NOT expose a partial mixture when validation or rendering fails. + +#### Scenario: Successful replacement +- **WHEN** every element in a replacement layer renders successfully +- **THEN** Viser displays the complete new generation and removes the previous generation + +#### Scenario: Failed replacement +- **WHEN** any element in a replacement layer fails validation or rendering +- **THEN** Viser removes partial replacement handles, retains the complete previous valid generation, and reports a visualization warning + +#### Scenario: Empty replacement +- **WHEN** a registered layer is cleared +- **THEN** Viser removes its element handles but retains the layer entry and viewer state + +### Requirement: Best-effort latest-wins updates +Layer submission SHALL NOT wait for Viser scene rendering. Viser SHALL retain at most the newest pending operation for each layer, SHALL order replacement and clear operations for the same layer, and SHALL contain visualization failures without failing a caller's robotics operation. + +#### Scenario: Replacements arrive faster than rendering +- **WHEN** multiple replacements for one layer arrive before its pending update renders +- **THEN** Viser may skip intermediate replacements but eventually renders the newest layer value + +#### Scenario: Clear supersedes replacement +- **WHEN** a clear operation supersedes an older pending replacement for the same layer +- **THEN** the older replacement does not reappear after the clear + +#### Scenario: Renderer is unavailable +- **WHEN** Viser is disconnected, closed, or raises while processing an update +- **THEN** the caller is not failed or blocked on scene rendering and the failure is logged as a visualization warning + +### Requirement: Hierarchical layer selector +Viser SHALL derive a hierarchical selector from slash-separated layer IDs. Leaf layers SHALL be independently toggleable, group controls SHALL toggle all descendants, and replacing or clearing contents SHALL preserve viewer-owned visibility. + +#### Scenario: Hierarchical registration +- **WHEN** layers `grasp/object-cloud` and `grasp/proposals` are first submitted +- **THEN** Viser displays them as independently controlled children of a `Grasp` group + +#### Scenario: Hidden layer is replaced +- **WHEN** a user hides a layer and its producer publishes replacement contents +- **THEN** Viser updates the contents while leaving the layer hidden + +#### Scenario: Parent group is toggled +- **WHEN** a user changes a group checkbox +- **THEN** Viser applies that visibility to every descendant leaf layer + +#### Scenario: Multiple clients view the selector +- **WHEN** more than one Viser client is connected +- **THEN** they observe the same server-global layer visibility state + +### Requirement: Viser point-cloud rendering +Viser SHALL render point-cloud elements with source RGB when present and cyan otherwise. It SHALL use a default point size of 5 mm unless overridden and SHALL cap display at 20,000 points per element without modifying the submitted element. + +#### Scenario: Colored cloud under the cap +- **WHEN** a point-cloud element provides RGB colors and at most 20,000 points +- **THEN** Viser renders every point with its matching source color + +#### Scenario: Uncolored cloud +- **WHEN** a point-cloud element provides no colors +- **THEN** Viser renders it using the cyan fallback + +#### Scenario: Cloud exceeds the render cap +- **WHEN** a point-cloud element contains more than 20,000 points +- **THEN** Viser deterministically samples at most 20,000 points, applies the same sample to colors, and leaves the source element unchanged + +### Requirement: Viser line-set rendering +Viser SHALL render generic line-set elements from vertices and indexed edges, preserving supplied uniform or per-line colors and explicit positive line width. + +#### Scenario: Gripper wireframe line set +- **WHEN** a line set describes a parallel-jaw gripper proposal +- **THEN** Viser renders every indexed segment at the supplied pose and appearance + +#### Scenario: Multiple proposal elements +- **WHEN** a layer contains line sets with unique IDs for multiple grasp ranks +- **THEN** Viser renders each proposal independently within the same layer lifecycle + +### Requirement: Interactive banana grasp visualization demo +The repository SHALL provide an interactive visualization-only command that loads the existing banana segmented object cloud, runs the real GraspGenX proposal provider, publishes the cloud and up to the configured number of ranked proposal wireframes to Viser, prints the visualization URL, and remains active until interrupted. + +#### Scenario: Default demo run +- **WHEN** a contributor runs `python -m dimos.manipulation.demo_grasp_visualization` +- **THEN** the demo publishes the banana to `grasp/object-cloud`, publishes at most 20 green-to-orange ranked gripper wireframes to `grasp/proposals`, prints the Viser URL, and performs no planning or execution + +#### Scenario: Candidate limit override +- **WHEN** a contributor supplies a positive `--max-candidates` value +- **THEN** the demo publishes no more than that number of score-ranked proposals + +#### Scenario: Demo shutdown +- **WHEN** the contributor interrupts the demo +- **THEN** it closes the visualization and GraspGenX resources cleanly + +#### Scenario: Optional dependency is unavailable +- **WHEN** GraspGenX or the Viser manipulation dependency cannot be initialized +- **THEN** the demo exits with actionable installation guidance + +### Requirement: Visualization-focused verification +Automated verification SHALL cover layer-model validation and ownership, Viser rendering and lifecycle, hierarchical visibility, atomic failure behavior, latest-wins ordering, and demo wiring without requiring a browser, robot, planner, or live perception system. + +#### Scenario: Automated test environment +- **WHEN** the visualization test suite runs with fake Viser scene handles and injected demo dependencies +- **THEN** it validates the specified behavior without launching a browser or executing robot motion